vllm.model_executor.hw_agnostic.custom_op

Classes:

• CustomOp –

  Base class for custom ops.

• PluggableLayer –

  Base class for pluggable layers.

CustomOp

Bases: Module

Base class for custom ops. Dispatches the forward method to the appropriate backend.

Methods:

• default_on –

  Behavior controlled by CompilationConfig.custom_ops: On by default if

• forward_native –

  PyTorch-native implementation; OOT plugins override via

• maybe_compile –

  Compile fn if compilation enabled.

Source code in vllm/model_executor/hw_agnostic/custom_op.py

class CustomOp(nn.Module):
    """Base class for custom ops.
    Dispatches the forward method to the appropriate backend.
    """

    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

        return super().__new__(_class_to_instantiate(cls, op_name))

    def __init__(self, *, enforce_enable: bool = False, compile_native: bool = False):
        super().__init__()
        self._enforce_enable = enforce_enable
        self._forward_method = self.dispatch_forward(compile_native=compile_native)

    def forward(self, *args, **kwargs):
        return self._forward_method(*args, **kwargs)

    def forward_native(self, *args, **kwargs):
        """PyTorch-native implementation; OOT plugins override via
        :meth:`forward_oot`."""
        raise NotImplementedError

    def forward_oot(self, *args, **kwargs):
        # By default, OOT ops fall back to the PyTorch-native implementation.
        return self.forward_native(*args, **kwargs)

    def dispatch_forward(self, compile_native: bool):
        # The hw-agnostic CustomOp dispatches only between forward_native
        # (in-tree, used on every platform) and forward_oot (overridden by
        # an OOT plugin via register_oot).
        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_out_of_tree():
            return self.forward_oot
        return self.forward_native

    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,
        )

    @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
        assert not (enabled and disabled), f"Cannot enable and disable {cls.name}"

        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")
        assert count_none + count_all == 1

        return not count_none > 0 or count_all > 0

    # Decorator to register custom ops.
    @classmethod
    def register(
        cls,
        name: str,
        dynamic_arg_dims: dict[str, int | list[int]] | None = None,
    ):
        def decorator(op_cls):
            assert name not in op_registry, f"Duplicate op name: {name}"
            op_cls.name = name
            op_cls._dynamic_arg_dims = dynamic_arg_dims
            op_registry[name] = op_cls
            return op_cls

        return decorator

    # Decorator to register out-of-tree(oot) custom ops.
    # For OOT custom ops:
    #   if in-tree layer class is registered with an oot_custom_op layer,
    #   the oot_custom_op layer will be used instead.
    # Example:
    # - @UnquantizedFusedMoEMethod.register_oot
    #   class HPUUnquantizedFusedMoEMethod(UnquantizedFusedMoEMethod)
    # or
    # - @CustomOP.register_oot(name="UnquantizedFusedMoEMethod")
    @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.")

default_on() staticmethod

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'.

Source code in vllm/model_executor/hw_agnostic/custom_op.py

@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")
    assert count_none + count_all == 1

    return not count_none > 0 or count_all > 0

forward_native(*args, **kwargs)

PyTorch-native implementation; OOT plugins override via :meth:forward_oot.

Source code in vllm/model_executor/hw_agnostic/custom_op.py

def forward_native(self, *args, **kwargs):
    """PyTorch-native implementation; OOT plugins override via
    :meth:`forward_oot`."""
    raise NotImplementedError

maybe_compile(fn, *, enable=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.

Source code in vllm/model_executor/hw_agnostic/custom_op.py

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,
    )

PluggableLayer

Bases: 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.

Source code in vllm/model_executor/hw_agnostic/custom_op.py

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

        return super().__new__(_class_to_instantiate(cls, layer_class_name))

    # 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_to_instantiate(cls, name)

The out-of-tree override registered for cls as name, else cls.

The in-tree and the hw-agnostic path keep separate registries, both keyed on the bare class name, so a class can be registered in the hw-agnostic path but not in the in-tree one, and vice versa. This catches the two cases where an override written against the in-tree class would go unnoticed:

• Registered in the hw-agnostic path, but not derived from cls: Python would skip its __init__, so this raises.
• Registered in the in-tree path, but not in the hw-agnostic one: it would never be used, so this warns.

Source code in vllm/model_executor/hw_agnostic/custom_op.py

def _class_to_instantiate(cls: type, name: str) -> type:
    """The out-of-tree override registered for `cls` as `name`, else `cls`.

    The in-tree and the hw-agnostic path keep separate registries, both keyed on
    the bare class name, so a class can be registered in the hw-agnostic path
    but not in the in-tree one, and vice versa. This catches the two cases where
    an override written against the in-tree class would go unnoticed:

    * Registered in the hw-agnostic path, but not derived from `cls`: Python
      would skip its `__init__`, so this raises.
    * Registered in the in-tree path, but not in the hw-agnostic one: it would
      never be used, so this warns.
    """
    override = op_registry_oot.get(name)
    if override is None:
        from vllm.model_executor.custom_op import op_registry_oot as in_tree_oot

        if name in in_tree_oot:
            logger.warning_once(
                "Out-of-tree override %s is registered for the in-tree %s, so the "
                "hw-agnostic path does not use it. To opt in, subclass "
                "hw_agnostic.resolve(<module>, %r) and register it with that "
                "class's register_oot.",
                in_tree_oot[name].__name__,
                name,
                name,
            )
        return cls
    if not issubclass(override, cls):
        raise TypeError(
            f"Out-of-tree override {override.__name__!r} for {name!r} does not "
            f"derive from {cls.__module__}.{name}, the class the hw-agnostic path "
            f"instantiates. Subclass hw_agnostic.resolve(<module>, {name!r}) "
            f"instead."
        )
    logger.debug("Instantiating %s using %s", name, override)
    return override