vLLM Explained

model_executor/custom_op.py: CustomOp, PluggableLayer and their registries

AdvancedCode mapabout 7 min read + 3 min deep dives

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.

Filevllm/model_executor/custom_op.py, 356 lines of Python. Two module dicts, one helper function, two classes.
Main classesPluggableLayer (vllm/model_executor/custom_op.py:32) and CustomOp (vllm/model_executor/custom_op.py:102), both subclasses of nn.Module.
Runs inEvery 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 byThe 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).
Callsget_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 fileThe 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#

Figure: the life of one CustomOp object. All line numbers refer to vllm/model_executor/custom_op.py unless the box names another file.
Open full size ↗

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#

What it calls#

Key classes and functions#

NameLineWhat it does
op_registryvllm/model_executor/custom_op.py:21Registered op name → class, for CustomOp and PluggableLayer.
op_registry_ootvllm/model_executor/custom_op.py:22In-tree class name (or the given name) → out-of-tree class.
maybe_get_oot_by_class(cls)vllm/model_executor/custom_op.py:25-29The OOT replacement of a class, else the class.
PluggableLayer.__new__vllm/model_executor/custom_op.py:46-65Builds the OOT class when the class name is in op_registry_oot.
PluggableLayer.register(name)vllm/model_executor/custom_op.py:68-76Sets name, adds to op_registry.
PluggableLayer.register_ootvllm/model_executor/custom_op.py:82-99Adds an OOT layer class.
CustomOp.__new__vllm/model_executor/custom_op.py:107-126Same swap as above.
CustomOp.__init__, forwardvllm/model_executor/custom_op.py:128-134Binds the method once; forward calls it.
forward_native ... forward_ootvllm/model_executor/custom_op.py:136-170Seven methods and their default bodies.
dispatch_forwardvllm/model_executor/custom_op.py:172-205Enabled check + platform choice.
maybe_compilevllm/model_executor/custom_op.py:207-266Optional 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-307Base mode: "all" → true, "none" → false.
CustomOp.register(name, dynamic_arg_dims)vllm/model_executor/custom_op.py:310-323Sets name and _dynamic_arg_dims, adds to op_registry.
CustomOp.register_ootvllm/model_executor/custom_op.py:334-356Adds 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)
  1. Lines 108-115. The try block reads cls.__name__. Every Python class has __name__, so the AttributeError branch, with its text about the name attribute, cannot run for a class. (Derived from Python semantics.)
  2. 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).
  3. Lines 120-125. The OOT class is used, with a DEBUG log.
  4. 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
  1. 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.
  2. Line 183. enforce_enable from __init__ overrides the list. The comment says it exists for the vision encoder of multimodal models and will be removed (lines 177-182).
  3. Lines 184-187. The counters enabled_custom_ops and disabled_custom_ops record self.__class__.name.
  4. Lines 189-192. Disabled: forward_native, compiled only with compile_native=True.
  5. Lines 194-205. Enabled: the first matching platform test wins; CUDA is the final else. An UnspecifiedPlatform therefore gets forward_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
  1. Line 273. A class with no name attribute was never registered: WARNING once, and the base mode decides.
  2. Lines 282-289. Both +name and -name: ValueError. CompilationConfig already rejects this earlier (vllm/config/compilation.py:1027-1034).
  3. Line 290. (default_on() or enabled) and not disabled.
  4. Lines 300-305. Exactly one base mode, else ValueError.
  5. 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.")
  1. Line 337. The key is name, or the name of the class on which the decorator is called. On CustomOp.register_oot() without name, the key would be "CustomOp" (derived).
  2. Line 338. A duplicate key fails the assert.
  3. Line 339. The OOT class's name becomes the key. The +name switch for it is then +RotaryEmbedding, not the in-tree op name (derived; see CustomOp, out-of-tree ops and IR op providers).
  4. 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 dive

Data: the registries and the attributes this file sets

ValueTypeSet at
op_registrydict[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_ootdict[str, type], key = class name such as "RotaryEmbedding"vllm/model_executor/custom_op.py:88, vllm/model_executor/custom_op.py:340
cls.namestr class attribute; inherited by subclassesvllm/model_executor/custom_op.py:72, vllm/model_executor/custom_op.py:318, vllm/model_executor/custom_op.py:339
cls._dynamic_arg_dimsdict[str, int | list[int]] | Nonevllm/model_executor/custom_op.py:319
self._enforce_enable, self._forward_methodbool, bound method or compiled functionvllm/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 dive

maybe_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,
        )
  1. Lines 218-228. Return fn unchanged when compilation is off for this call, for the whole config (mode NONE), or with the "eager" backend.
  2. Lines 230-233. The backend is the platform's simple_compile_backend ("inductor" by default, vllm/platforms/interface.py:172), with graph partition off.
  3. 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).
  4. 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 dive

Processes and timing

  • Import time. register runs when Python imports a layer module. register_oot runs when the plugin module is imported.
  • Model build. __new__, __init__ and dispatch_forward run once per layer object, inside set_current_vllm_config. A CustomOp built when no config is set raises AssertionError (vllm/config/vllm.py:3576-3586).
  • Forward. forward only 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 dive

Error paths: what raises, what only logs

SituationResultWhere
Duplicate register nameAssertionError "Duplicate op name"vllm/model_executor/custom_op.py:317, vllm/model_executor/custom_op.py:71
Duplicate OOT keyAssertionErrorvllm/model_executor/custom_op.py:338, vllm/model_executor/custom_op.py:86
register_oot("X")TypeErrorvllm/model_executor/custom_op.py:354-356
Unregistered CustomOp subclassWARNING once; base mode decidesvllm/model_executor/custom_op.py:273-280
+x and -x bothValueErrorvllm/model_executor/custom_op.py:284-288
No or two base modesValueErrorvllm/model_executor/custom_op.py:302-305
A forward_cuda or forward_native that the subclass does not define, when chosenNotImplementedError at the first forwardvllm/model_executor/custom_op.py:136-145
Deep dive

Configuration that this file reads

SettingDefaultSet atEffect 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-1826enabled(), default_on()
compilation_config.modeVLLM_COMPILE above O0; NONE with --enforce-eager or breakable CUDA graphsvllm/config/vllm.py:1806-1810, vllm/config/vllm.py:1745-1746, vllm/config/vllm.py:849-851maybe_compile line 223
compilation_config.backendplatform's compile backendvllm/config/compilation.py:1080-1081maybe_compile line 227

No environment variable is read in this file.

Check yourself#

  1. 1
    Which dict does __new__ read, and with which key?
    easy
    Show answer
  2. 2
    An enabled op runs on ROCm. The subclass defines only forward_native and forward_cuda. Which method runs?
    medium
    Show answer

    dispatch_forward returns forward_hip (vllm/model_executor/custom_op.py:194-195). Its default body calls forward_cuda (vllm/model_executor/custom_op.py:147-149).

  3. 3
    What is the difference between PluggableLayer and CustomOp?
    medium
    Show answer

    Both swap the class in __new__ through op_registry_oot. Only CustomOp has the forward_* methods and dispatch_forward (vllm/model_executor/custom_op.py:33-44).

  4. 4
    An op registered with dynamic_arg_dims={"x": -1} is disabled and built with compile_native=True under VLLM_COMPILE. What does a call with a 2-D tensor x do?
    hard
    Show answer

    maybe_compile compiled forward_native with dynamic=False. The wrapper binds the arguments, turns -1 into x.ndim - 1 = 1, calls mark_dynamic(x, 1), then calls the compiled function (vllm/model_executor/custom_op.py:235-258).