| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import functools |
| 4 | import inspect |
| 5 | |
| 6 | import torch |
| 7 | import torch.nn as nn |
| 8 | |
| 9 | from vllm.config import get_cached_compilation_config |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.model_executor.utils import maybe_disable_graph_partition |
| 12 | from vllm.platforms import current_platform |
| 13 | |
| 14 | logger = init_logger(__name__) |
| 15 | |
| 16 | # Dictionary of all custom ops (classes, indexed by registered name). |
| 17 | # To check if an op with a name is enabled, call .enabled() on the class. |
| 18 | # Examples: |
| 19 | # - MyOp.enabled() |
| 20 | # - op_registry["my_op"].enabled() |
| 21 | op_registry: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {} |
| 22 | op_registry_oot: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {} |
| 23 | |
| 24 | |
| 25 | def maybe_get_oot_by_class(class_type: type) -> type: |
| 26 | class_name = class_type.__name__ |
| 27 | if class_name in op_registry_oot: |
| 28 | return op_registry_oot[class_name] |
| 29 | return class_type |
| 30 | |
| 31 | |
| 32 | class PluggableLayer(nn.Module): |
| 33 | """Base class for pluggable layers. |
| 34 | |
| 35 | A PluggableLayer is a *module-composing* abstraction: it may instantiate other |
| 36 | ``torch.nn.Module`` objects as sub-layers, and its functionality depends on |
| 37 | these sub-layers following a generalized invocation sequence. Also, it is stateful |
| 38 | and may hold parameters or buffers. |
| 39 | |
| 40 | Unlike :class:`CustomOp`, PluggableLayer does NOT provide per-platform |
| 41 | ``forward_*`` dispatch. Instead, it supports out-of-tree (OOT) replacement |
| 42 | of the entire layer class at instantiation time, allowing customized |
| 43 | initialization and submodule composition. |
| 44 | """ |
| 45 | |
| 46 | def __new__(cls, *args, **kwargs): |
| 47 | try: |
| 48 | layer_class_name = cls.__name__ |
| 49 | except AttributeError: |
| 50 | raise TypeError( |
| 51 | f"Cannot instantiate '{cls.__name__}': its 'name' attribute " |
| 52 | f"was not set, possibly because it was not decorated with " |
| 53 | f"@PluggableLayer.register, or it's the PluggableLayer itself." |
| 54 | ) from None |
| 55 | |
| 56 | if layer_class_name not in op_registry_oot: |
| 57 | layer_cls_to_instantiate = cls |
| 58 | else: |
| 59 | layer_cls_to_instantiate = op_registry_oot[layer_class_name] |
| 60 | logger.debug( |
| 61 | "Instantiating pluggable layer: %s using %s", |
| 62 | layer_class_name, |
| 63 | str(layer_cls_to_instantiate), |
| 64 | ) |
| 65 | return super().__new__(layer_cls_to_instantiate) |
| 66 | |
| 67 | # Decorator to register pluggable layers. |
| 68 | @classmethod |
| 69 | def register(cls, name: str): |
| 70 | def decorator(op_cls): |
| 71 | assert name not in op_registry, f"Duplicate op name: {name}" |
| 72 | op_cls.name = name |
| 73 | op_registry[name] = op_cls |
| 74 | return op_cls |
| 75 | |
| 76 | return decorator |
| 77 | |
| 78 | # Decorator to register out-of-tree(oot) pluggable layers. |
| 79 | # For OOT pluggable layers: |
| 80 | # if in-tree layer class is registered with an oot_custom_layer, |
| 81 | # the oot_custom_layer will be used instead. |
| 82 | @classmethod |
| 83 | def register_oot(cls, _decorated_layer_cls=None, name: str | None = None): |
| 84 | def decorator(layer_cls): |
| 85 | reg_name = name if name is not None else cls.__name__ |
| 86 | assert reg_name not in op_registry_oot, f"Duplicate layer name: {reg_name}" |
| 87 | layer_cls.name = reg_name |
| 88 | op_registry_oot[reg_name] = layer_cls |
| 89 | return layer_cls |
| 90 | |
| 91 | if _decorated_layer_cls is None: |
| 92 | # Called with parentheses: @PluggableLayer.register_oot() |
| 93 | # or @PluggableLayer.register_oot(name="...") |
| 94 | return decorator |
| 95 | elif isinstance(_decorated_layer_cls, type): # Check if it's a class |
| 96 | # Called without parentheses: @PluggableLayer.register_oot |
| 97 | return decorator(_decorated_layer_cls) |
| 98 | else: |
| 99 | raise TypeError("Decorator can only be applied to classes.") |
| 100 | |
| 101 | |
| 102 | class CustomOp(nn.Module): |
| 103 | """Base class for custom ops. |
| 104 | Dispatches the forward method to the appropriate backend. |
| 105 | """ |
| 106 | |
| 107 | def __new__(cls, *args, **kwargs): |
| 108 | try: |
| 109 | op_name = cls.__name__ |
| 110 | except AttributeError: |
| 111 | raise TypeError( |
| 112 | f"Cannot instantiate '{cls.__name__}': its 'name' attribute " |
| 113 | f"was not set, possibly because it was not decorated with " |
| 114 | f"@CustomOp.register, or it's the CustomOp base class itself." |
| 115 | ) from None |
| 116 | |
| 117 | if op_name not in op_registry_oot: |
| 118 | op_cls_to_instantiate = cls |
| 119 | else: |
| 120 | op_cls_to_instantiate = op_registry_oot[op_name] |
| 121 | logger.debug( |
| 122 | "Instantiating custom op: %s using %s", |
| 123 | op_name, |
| 124 | str(op_cls_to_instantiate), |
| 125 | ) |
| 126 | return super().__new__(op_cls_to_instantiate) |
| 127 | |
| 128 | def __init__(self, *, enforce_enable: bool = False, compile_native: bool = False): |
| 129 | super().__init__() |
| 130 | self._enforce_enable = enforce_enable |
| 131 | self._forward_method = self.dispatch_forward(compile_native=compile_native) |
| 132 | |
| 133 | def forward(self, *args, **kwargs): |
| 134 | return self._forward_method(*args, **kwargs) |
| 135 | |
| 136 | def forward_native(self, *args, **kwargs): |
| 137 | """PyTorch-native implementation; OOT plugins override via |
| 138 | :meth:`forward_oot`.""" |
| 139 | raise NotImplementedError |
| 140 | |
| 141 | def forward_oot(self, *args, **kwargs): |
| 142 | # By default, OOT ops fall back to the PyTorch-native implementation. |
| 143 | return self.forward_native(*args, **kwargs) |
| 144 | |
| 145 | def dispatch_forward(self, compile_native: bool): |
| 146 | # The hw-agnostic CustomOp dispatches only between forward_native |
| 147 | # (in-tree, used on every platform) and forward_oot (overridden by |
| 148 | # an OOT plugin via register_oot). |
| 149 | compilation_config = get_cached_compilation_config() |
| 150 | |
| 151 | # NOTE(shen-shanshan): CustomOp object can be enforce enabled, e.g., |
| 152 | # enable device-specific kernels in ViT models when enabling graph |
| 153 | # mode. By default, it will follow the compilation_config to determine |
| 154 | # whether enable itself. |
| 155 | # This enforce_enable mechanism will be removed after we adding a |
| 156 | # separate compilation_config for multi-modal part. |
| 157 | enabled = self._enforce_enable or self.enabled() |
| 158 | if enabled: |
| 159 | compilation_config.enabled_custom_ops.update([self.__class__.name]) |
| 160 | else: |
| 161 | compilation_config.disabled_custom_ops.update([self.__class__.name]) |
| 162 | |
| 163 | if not enabled: |
| 164 | # Compile forward_native to avoid eager torch ops if inside |
| 165 | # opaque torch custom op (e.g. fused_moe, unified_attention, etc.) |
| 166 | return self.maybe_compile(self.forward_native, enable=compile_native) |
| 167 | |
| 168 | if current_platform.is_out_of_tree(): |
| 169 | return self.forward_oot |
| 170 | return self.forward_native |
| 171 | |
| 172 | def maybe_compile(self, fn, *, enable: bool = True): |
| 173 | """Compile fn if compilation enabled. |
| 174 | Useful for CustomOp instances called from within a torch custom op, |
| 175 | meaning the forward call is hidden from the model-level torch.compile. |
| 176 | |
| 177 | NOTE: this does not enable fusion across ops, so opaque custom ops |
| 178 | should still be unwrapped wherever possible. |
| 179 | """ |
| 180 | from vllm.config.compilation import CompilationMode |
| 181 | |
| 182 | # Do not compile if compilation disabled |
| 183 | if not enable: |
| 184 | return fn |
| 185 | |
| 186 | # Do not compile if global compilation disabled |
| 187 | compilation_config = get_cached_compilation_config() |
| 188 | if compilation_config.mode == CompilationMode.NONE: |
| 189 | return fn |
| 190 | |
| 191 | # If eager backend is used, do not compile either |
| 192 | if compilation_config.backend == "eager": |
| 193 | return fn |
| 194 | |
| 195 | compile_options = maybe_disable_graph_partition( |
| 196 | current_platform.simple_compile_backend |
| 197 | ) |
| 198 | backend = current_platform.simple_compile_backend |
| 199 | |
| 200 | dynamic_arg_dims = getattr(self.__class__, "_dynamic_arg_dims", None) |
| 201 | if dynamic_arg_dims is not None: |
| 202 | compiled_fn = torch.compile( |
| 203 | fn, |
| 204 | dynamic=False, |
| 205 | backend=backend, |
| 206 | options=compile_options, |
| 207 | ) |
| 208 | sig = inspect.signature(fn) |
| 209 | |
| 210 | @functools.wraps(fn) |
| 211 | def wrapper(*args, **kwargs): |
| 212 | bound = sig.bind(*args, **kwargs) |
| 213 | bound.apply_defaults() |
| 214 | for name, dims in dynamic_arg_dims.items(): |
| 215 | arg = bound.arguments.get(name) |
| 216 | if arg is not None and isinstance(arg, torch.Tensor): |
| 217 | dims_list = [dims] if isinstance(dims, int) else dims |
| 218 | for d in dims_list: |
| 219 | real_d = arg.ndim + d if d < 0 else d |
| 220 | torch._dynamo.mark_dynamic(arg, real_d) |
| 221 | return compiled_fn(*args, **kwargs) |
| 222 | |
| 223 | return wrapper |
| 224 | |
| 225 | # dynamic=True to avoid recompilations |
| 226 | return torch.compile( |
| 227 | fn, |
| 228 | dynamic=True, |
| 229 | backend=backend, |
| 230 | options=compile_options, |
| 231 | ) |
| 232 | |
| 233 | @classmethod |
| 234 | def enabled(cls) -> bool: |
| 235 | # if no name, then it was not registered |
| 236 | compilation_config = get_cached_compilation_config() |
| 237 | custom_ops = compilation_config.custom_ops |
| 238 | if not hasattr(cls, "name"): |
| 239 | logger.warning_once( |
| 240 | "Custom op %s was not registered, which means it won't appear " |
| 241 | "in the op registry. It will be enabled/disabled based on the " |
| 242 | "global settings.", |
| 243 | cls.__name__, |
| 244 | ) |
| 245 | return CustomOp.default_on() |
| 246 | |
| 247 | enabled = f"+{cls.name}" in custom_ops |
| 248 | disabled = f"-{cls.name}" in custom_ops |
| 249 | assert not (enabled and disabled), f"Cannot enable and disable {cls.name}" |
| 250 | |
| 251 | return (CustomOp.default_on() or enabled) and not disabled |
| 252 | |
| 253 | @staticmethod |
| 254 | def default_on() -> bool: |
| 255 | """Behavior controlled by `CompilationConfig.custom_ops`: On by default if |
| 256 | 'all', off by default if 'none'. |
| 257 | When PyTorch Inductor is used, 'none' is the default value, |
| 258 | otherwise 'all'. |
| 259 | """ |
| 260 | compilation_config = get_cached_compilation_config() |
| 261 | count_none = compilation_config.custom_ops.count("none") |
| 262 | count_all = compilation_config.custom_ops.count("all") |
| 263 | assert count_none + count_all == 1 |
| 264 | |
| 265 | return not count_none > 0 or count_all > 0 |
| 266 | |
| 267 | # Decorator to register custom ops. |
| 268 | @classmethod |
| 269 | def register( |
| 270 | cls, |
| 271 | name: str, |
| 272 | dynamic_arg_dims: dict[str, int | list[int]] | None = None, |
| 273 | ): |
| 274 | def decorator(op_cls): |
| 275 | assert name not in op_registry, f"Duplicate op name: {name}" |
| 276 | op_cls.name = name |
| 277 | op_cls._dynamic_arg_dims = dynamic_arg_dims |
| 278 | op_registry[name] = op_cls |
| 279 | return op_cls |
| 280 | |
| 281 | return decorator |
| 282 | |
| 283 | # Decorator to register out-of-tree(oot) custom ops. |
| 284 | # For OOT custom ops: |
| 285 | # if in-tree layer class is registered with an oot_custom_op layer, |
| 286 | # the oot_custom_op layer will be used instead. |
| 287 | # Example: |
| 288 | # - @UnquantizedFusedMoEMethod.register_oot |
| 289 | # class HPUUnquantizedFusedMoEMethod(UnquantizedFusedMoEMethod) |
| 290 | # or |
| 291 | # - @CustomOP.register_oot(name="UnquantizedFusedMoEMethod") |
| 292 | @classmethod |
| 293 | def register_oot(cls, _decorated_op_cls=None, name: str | None = None): |
| 294 | def decorator(op_cls): |
| 295 | reg_name = name if name is not None else cls.__name__ |
| 296 | assert reg_name not in op_registry_oot, f"Duplicate op name: {reg_name}" |
| 297 | op_cls.name = reg_name |
| 298 | op_registry_oot[reg_name] = op_cls |
| 299 | return op_cls |
| 300 | |
| 301 | if _decorated_op_cls is None: |
| 302 | # Called with parentheses: @CustomOP.register_oot() |
| 303 | # or @CustomOP.register_oot(name="...") |
| 304 | # So, _decorated_op_cls is None. |
| 305 | # We return the actual decorator function. |
| 306 | return decorator |
| 307 | elif isinstance(_decorated_op_cls, type): # Check if it's a class |
| 308 | # Called without parentheses: @CustomOP.register_oot |
| 309 | # The first argument is the class itself. |
| 310 | # We call the 'decorator' function immediately with the class. |
| 311 | return decorator(_decorated_op_cls) |
| 312 | else: |
| 313 | # Handle other unexpected cases if necessary |
| 314 | raise TypeError("Decorator can only be applied to classes.") |