| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | import inspect |
| 5 | import traceback |
| 6 | from collections.abc import Callable |
| 7 | from pathlib import Path |
| 8 | from typing import Any, ClassVar, Literal, overload |
| 9 | |
| 10 | import regex as re |
| 11 | import torch |
| 12 | from torch.library import Library, infer_schema |
| 13 | |
| 14 | from vllm.ir.tolerances import DEFAULT_TOLERANCES, ToleranceSpec |
| 15 | from vllm.ir.util import hash_source, weak_cache |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.logging_utils import lazy, tensors_str_no_data |
| 18 | |
| 19 | InputGenerator = Callable[..., tuple[Any, ...]] |
| 20 | |
| 21 | vllm_ir_torch_lib = Library("vllm_ir", "FRAGMENT") # IR op lib; monkeypatch in tests. |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | |
| 26 | def _torch_ops_subtree(lib: Any) -> Any: |
| 27 | """``torch.ops`` subtree for ``lib.ns``; fall back if doc mocks replace ``ns``.""" |
| 28 | ns = getattr(lib, "ns", None) |
| 29 | if isinstance(ns, str): |
| 30 | return getattr(torch.ops, ns) |
| 31 | return torch.ops.vllm_ir |
| 32 | |
| 33 | |
| 34 | _NAME_PATTERN = re.compile(r"^[a-z_][a-z_0-9]*$") |
| 35 | |
| 36 | RESERVED_PROVIDERS = ["native", "unfused"] |
| 37 | """Providers that are reserved and cannot be used for custom implementations.""" |
| 38 | |
| 39 | |
| 40 | def _validate_name(name: str, entity_type: str) -> None: |
| 41 | """Validate that a name matches the required pattern `[a-z_][a-z_0-9]*`.""" |
| 42 | if not _NAME_PATTERN.match(name): |
| 43 | raise ValueError( |
| 44 | f"{entity_type} name '{name}' is invalid. " |
| 45 | f"Names must start with a letter or underscore, " |
| 46 | f"followed by lowercase letters, underscores, or digits only." |
| 47 | ) |
| 48 | |
| 49 | |
| 50 | _ENABLE_TORCH_WRAP: bool = True |
| 51 | """Global override flag to control torch op layer wrapping.""" |
| 52 | |
| 53 | |
| 54 | def set_default_torch_wrap(enable: bool = True) -> None: |
| 55 | """Permanently set the torch wrap flag.""" |
| 56 | global _ENABLE_TORCH_WRAP |
| 57 | _ENABLE_TORCH_WRAP = enable |
| 58 | |
| 59 | |
| 60 | @contextlib.contextmanager |
| 61 | def enable_torch_wrap(enable: bool = True): |
| 62 | """Context manager to enable/disable torch custom op wrapping for vLLM IR ops. |
| 63 | When torch wrapping is disabled, the torch custom op layer is skipped |
| 64 | and IR ops dispatch directly to the implementation. |
| 65 | Helpful for avoiding torch dispatch overhead in eager mode |
| 66 | and avoiding the need for lowering for platforms not using Inductor. |
| 67 | """ |
| 68 | global _ENABLE_TORCH_WRAP |
| 69 | old = _ENABLE_TORCH_WRAP |
| 70 | try: |
| 71 | _ENABLE_TORCH_WRAP = enable |
| 72 | yield |
| 73 | finally: |
| 74 | _ENABLE_TORCH_WRAP = old |
| 75 | |
| 76 | |
| 77 | # 0-param decorator overload (no inplace) |
| 78 | @overload |
| 79 | def register_op(f: Callable[..., Any]) -> "IrOp": ... |
| 80 | |
| 81 | |
| 82 | # parametrized decorator with allow_inplace=False (default) |
| 83 | @overload |
| 84 | def register_op( |
| 85 | *, |
| 86 | name: str | None = None, |
| 87 | activations: list[str] | None = None, |
| 88 | allow_inplace: Literal[False] = False, |
| 89 | ) -> Callable[[Callable[..., Any]], "IrOp"]: ... |
| 90 | |
| 91 | |
| 92 | # parametrized decorator with allow_inplace=True |
| 93 | @overload |
| 94 | def register_op( |
| 95 | *, |
| 96 | name: str | None = None, |
| 97 | activations: list[str] | None = None, |
| 98 | allow_inplace: Literal[True], |
| 99 | ) -> Callable[[Callable[..., Any]], "IrOpInplace"]: ... |
| 100 | |
| 101 | |
| 102 | def register_op( |
| 103 | f: Callable | None = None, |
| 104 | *, |
| 105 | name: str | None = None, |
| 106 | activations: list[str] | None = None, |
| 107 | allow_inplace: bool = False, |
| 108 | ) -> "IrOp | Callable[[Callable], IrOp]": |
| 109 | """Register a new vLLM IR op. |
| 110 | |
| 111 | Args: |
| 112 | f: the native implementation of the op |
| 113 | name: the name of the op, defaults to the function name |
| 114 | activations: list of activation params, defaults to params starting with 'x' |
| 115 | allow_inplace: add a maybe_inplace overload that allows inplace impls |
| 116 | |
| 117 | Returns: |
| 118 | the IrOp object if f is provided, otherwise a decorator |
| 119 | |
| 120 | Example usage: |
| 121 | ```python |
| 122 | @vllm.ir.register_op |
| 123 | def my_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 124 | return x + y |
| 125 | |
| 126 | |
| 127 | @vllm.ir.register_op(name="custom_mul") |
| 128 | def multiply(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 129 | return x * y |
| 130 | ``` |
| 131 | |
| 132 | """ |
| 133 | |
| 134 | def decorator(_f: Callable): |
| 135 | op_name: str = _f.__name__ if name is None else name |
| 136 | _validate_name(op_name, "Op") |
| 137 | assert op_name not in IrOp.registry, f"Op '{op_name}' is already registered." |
| 138 | # Slice out the decorator function frames from the stack |
| 139 | stack = traceback.format_stack()[:-2] |
| 140 | if allow_inplace: |
| 141 | op: IrOp = IrOpInplace(op_name, _f, activations, stack) |
| 142 | else: |
| 143 | op = IrOp(op_name, _f, activations, stack) |
| 144 | IrOp.registry[op_name] = op |
| 145 | return op |
| 146 | |
| 147 | if f is not None: |
| 148 | return decorator(f) |
| 149 | |
| 150 | return decorator |
| 151 | |
| 152 | |
| 153 | class IrOp: |
| 154 | registry: ClassVar[dict[str, "IrOp"]] = {} |
| 155 | |
| 156 | name: str |
| 157 | impls: dict[str, "IrOpImpl"] |
| 158 | allow_inplace: bool = False |
| 159 | |
| 160 | def __init__( |
| 161 | self, |
| 162 | name: str, |
| 163 | native_impl: Callable, |
| 164 | activations: list[str] | None = None, |
| 165 | registration_stack: list[str] | None = None, |
| 166 | ): |
| 167 | self._py_signature = inspect.signature(native_impl) |
| 168 | if any( |
| 169 | p.kind == inspect.Parameter.KEYWORD_ONLY |
| 170 | for p in self._py_signature.parameters.values() |
| 171 | ): |
| 172 | raise ValueError( |
| 173 | f"Op {name} has keyword-only arguments which are not currently " |
| 174 | f"supported. That's because kwargs are not allowed during lowering." |
| 175 | ) |
| 176 | |
| 177 | # By convention, we consider parameters starting with 'x' as activations. |
| 178 | if activations is None: |
| 179 | activations = [ |
| 180 | p.name |
| 181 | for p in self._py_signature.parameters.values() |
| 182 | if p.name.startswith("x") |
| 183 | ] |
| 184 | |
| 185 | self.name = name |
| 186 | self._docstring = inspect.getdoc(native_impl) or "" |
| 187 | self._registration_stack = registration_stack or [] |
| 188 | self.impls: dict[str, IrOpImpl] = {} |
| 189 | self.activations = activations |
| 190 | self.activation_indices = [ |
| 191 | i |
| 192 | for i, p in enumerate(self._py_signature.parameters.values()) |
| 193 | if p.name in activations |
| 194 | ] |
| 195 | self._priority_impls: list[IrOpImpl] = [] |
| 196 | self._schema_str = infer_schema(native_impl, mutates_args=[]) |
| 197 | self._input_generator: InputGenerator | None = None |
| 198 | self._tolerance_overrides: ToleranceSpec = {} |
| 199 | |
| 200 | # native implementation |
| 201 | self.impls["native"] = IrOpImpl( |
| 202 | self, |
| 203 | "native", |
| 204 | native_impl, |
| 205 | # always supported |
| 206 | supported=True, |
| 207 | supports_args=None, |
| 208 | registration_stack=self._registration_stack, |
| 209 | ) |
| 210 | |
| 211 | # By default, fake routes directly to native, |
| 212 | # can be overridden by register_fake |
| 213 | self._fake_fn = native_impl |
| 214 | |
| 215 | # torch registration (resolve ``torch.ops`` subtree from ``lib.ns``) |
| 216 | lib = vllm_ir_torch_lib |
| 217 | lib.define(self.name + self._schema_str) |
| 218 | # CompositeExplicitAutograd is not decomposed |
| 219 | # by ATen IR normalization in AOTAutograd |
| 220 | lib.impl(self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd") |
| 221 | lib._register_fake(self.name, self._fake_call) |
| 222 | torch_ops = _torch_ops_subtree(lib) |
| 223 | assert hasattr(torch_ops, name) |
| 224 | self.torch_op: torch._ops.OpOverload = getattr(torch_ops, name).default |
| 225 | |
| 226 | def register_fake(self, fn: Callable) -> Callable: |
| 227 | """Register a fake impl for the torch custom op. If this method is not called, |
| 228 | the native implementation is used directly for the fake implementation. |
| 229 | """ |
| 230 | self._fake_fn = fn |
| 231 | return fn |
| 232 | |
| 233 | def _fake_call(self, *args, **kwargs) -> Any: |
| 234 | """Call to the fake implementation of the op. We use indirection because we want |
| 235 | users to be able to register fake later but also want it to fall back to native |
| 236 | directly by default, instead of going through the dispatching mechanism. |
| 237 | """ |
| 238 | return self._fake_fn(*args, **kwargs) |
| 239 | |
| 240 | def register_impl( |
| 241 | self, |
| 242 | provider: str, |
| 243 | *, |
| 244 | supported: bool = True, |
| 245 | supports_args: Callable[..., bool] | None = None, |
| 246 | inplace: bool = False, |
| 247 | ) -> Callable[[Callable[..., Any]], "IrOpImpl"]: |
| 248 | """Register an implementation for this custom op. |
| 249 | |
| 250 | Args: |
| 251 | provider: The name of the provider, must be unique. |
| 252 | supported: Static support check, use this to check platform support. |
| 253 | supports_args: Dynamic arg support check, used for types and shapes. |
| 254 | inplace: Does this op reuse activation input memory for outputs |
| 255 | |
| 256 | Returns: |
| 257 | A decorator that registers the implementation. |
| 258 | |
| 259 | The decorated function must have the same semantics and signature as |
| 260 | the native implementation. |
| 261 | |
| 262 | The provider name must be unique and not one of the RESERVED_PROVIDERS. |
| 263 | The supported and supports_args parameters should not be used to implement |
| 264 | custom enablement logic based on global state (e.g. environment variables). |
| 265 | Instead, supported param should only be used to check for platform support |
| 266 | (e.g. whether a specific hardware or library is available). |
| 267 | supports_args should be used to check whether the provided arguments are |
| 268 | compatible with the implementation. |
| 269 | For custom enablement logic, set op impl priority. |
| 270 | |
| 271 | Example: |
| 272 | ```python |
| 273 | @my_op.register_impl("my_provider", supported=torch.cuda.is_available()) |
| 274 | def my_provider_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: ... |
| 275 | ``` |
| 276 | |
| 277 | """ |
| 278 | assert provider not in RESERVED_PROVIDERS, ( |
| 279 | f"Provider name {provider} is reserved." |
| 280 | ) |
| 281 | _validate_name(provider, "Provider") |
| 282 | |
| 283 | def _register_impl(f: Callable): |
| 284 | # Slice out the decorator function from the stack |
| 285 | stack = traceback.format_stack()[:-1] |
| 286 | impl = IrOpImpl(self, provider, f, supported, supports_args, inplace, stack) |
| 287 | self.impls[provider] = impl |
| 288 | |
| 289 | if self.get_priority(): |
| 290 | logger.warning( |
| 291 | "Warning: registering new impl %s for op %s while priority is set.", |
| 292 | provider, |
| 293 | self.name, |
| 294 | ) |
| 295 | |
| 296 | return impl |
| 297 | |
| 298 | return _register_impl |
| 299 | |
| 300 | def _inner_call(self, *args, **kwargs) -> Any: |
| 301 | """Eager call to torch op lands here. When torch wrapping is disabled, |
| 302 | __call__ routes straight here instead of going through torch op dispatching. |
| 303 | """ |
| 304 | impl = self.dispatch(*args, **kwargs) |
| 305 | |
| 306 | # Default overload must be functional, |
| 307 | # use func_impl_fn to correctly handle inplace impls. |
| 308 | return impl.func_impl_fn(*args, **kwargs) |
| 309 | |
| 310 | def apply_arg_defaults(self, args) -> tuple: |
| 311 | """Return args with default values applied. |
| 312 | Defaults are taken from the native implementation signature. |
| 313 | |
| 314 | SHOULD NOT BE USED IN THE DISPATCH PATH (SLOW). |
| 315 | Only for Inductor lowering. |
| 316 | """ |
| 317 | bound_args = self._py_signature.bind(*args) |
| 318 | bound_args.apply_defaults() |
| 319 | return bound_args.args |
| 320 | |
| 321 | def dispatch(self, *args, **kwargs) -> "IrOpImpl": |
| 322 | """Dispatch to the appropriate implementation based on current priority |
| 323 | and argument support checks. Returns the selected IrOpImpl. |
| 324 | |
| 325 | THIS FUNCTION IS ON THE HOT PATH (OP DISPATCH), MUST BE FAST. |
| 326 | """ |
| 327 | if not self._priority_impls: |
| 328 | if not torch.compiler.is_compiling(): |
| 329 | # Logging not compatible with Dynamo tracing |
| 330 | # (this code is exposed when torch wrapping is disabled) |
| 331 | logger.warning_once( |
| 332 | "Priority not set for op %s, using native implementation.", |
| 333 | self.name, |
| 334 | ) |
| 335 | return self.impls["native"] |
| 336 | |
| 337 | for impl in self._priority_impls: |
| 338 | if not impl.supported: |
| 339 | raise ValueError( |
| 340 | f"Implementation {impl.provider} for op {self.name} not supported. " |
| 341 | f"All implementations in priority list must be supported." |
| 342 | ) |
| 343 | if impl.supports_args(*args, **kwargs): |
| 344 | return impl |
| 345 | |
| 346 | if not torch.compiler.is_compiling(): |
| 347 | logger.debug( |
| 348 | "Skipping provider %s because it does not support " |
| 349 | "%s with args=%s kwargs=%s", |
| 350 | impl.provider, |
| 351 | self.name, |
| 352 | lazy(lambda: tensors_str_no_data(args)), |
| 353 | lazy(lambda: tensors_str_no_data(kwargs)), |
| 354 | ) |
| 355 | |
| 356 | raise RuntimeError( |
| 357 | "Priority set incorrectly: the last implementation must " |
| 358 | "support all args (can be native). This is likely an internal bug" |
| 359 | ) |
| 360 | |
| 361 | def __call__(self, *args, **kwargs) -> Any: |
| 362 | if not _ENABLE_TORCH_WRAP: |
| 363 | return self._inner_call(*args, **kwargs) |
| 364 | |
| 365 | return self.torch_op(*args, **kwargs) |
| 366 | |
| 367 | def __repr__(self) -> str: |
| 368 | """Return unambiguous string representation.""" |
| 369 | return f"IrOp('{self.name}')" |
| 370 | |
| 371 | def __str__(self) -> str: |
| 372 | """Return human-readable string representation using docstring.""" |
| 373 | if not self._docstring: |
| 374 | return f"IrOp('{self.name}')" |
| 375 | first_line = self._docstring.split("\n")[0].strip() |
| 376 | return f"IrOp('{self.name}') - {first_line}" |
| 377 | |
| 378 | def get_priority(self) -> list[str]: |
| 379 | """Get the current dispatch priority for implementations for this op.""" |
| 380 | return [p.provider for p in self._priority_impls] |
| 381 | |
| 382 | def _filter_priority_impls(self, priority: list[str]) -> list["IrOpImpl"]: |
| 383 | assert all(p in self.impls for p in priority), ( |
| 384 | "All providers in priority must be registered implementations." |
| 385 | ) |
| 386 | filtered_impls: list[IrOpImpl] = [] |
| 387 | for p in priority: |
| 388 | impl = self.impls[p] |
| 389 | if not impl.supported: |
| 390 | # Skip unsupported implementations |
| 391 | continue |
| 392 | |
| 393 | filtered_impls.append(impl) |
| 394 | |
| 395 | # If all args are supported, skip other implementations |
| 396 | if impl.supports_all_args: |
| 397 | return filtered_impls |
| 398 | |
| 399 | logger.warning_once( |
| 400 | "Op %s: No implementation in priority list supports all args, " |
| 401 | "execution fallback to native is possible. To silence this warning, " |
| 402 | "explicitly add 'native' to the end of the priority list", |
| 403 | self.name, |
| 404 | ) |
| 405 | filtered_impls.append(self.impls["native"]) |
| 406 | return filtered_impls |
| 407 | |
| 408 | def set_default(self, priority: list[str]) -> None: |
| 409 | """Permanently set the dispatch priority for this op. Use this for |
| 410 | process-lifetime setup (e.g., worker startup). For scoped overrides, |
| 411 | use ``set_priority`` instead. |
| 412 | """ |
| 413 | self._priority_impls = self._filter_priority_impls(priority) |
| 414 | logger.debug( |
| 415 | "Priority for vllm.ir.%s set to %s", |
| 416 | self.name, |
| 417 | lazy(lambda: [p.provider for p in self._priority_impls]), |
| 418 | ) |
| 419 | |
| 420 | @contextlib.contextmanager |
| 421 | def set_priority(self, priority: list[str]): |
| 422 | """Context manager setting the dispatch priority for this op's impls.""" |
| 423 | old_priority_impls = self._priority_impls |
| 424 | try: |
| 425 | self._priority_impls = self._filter_priority_impls(priority) |
| 426 | logger.debug( |
| 427 | "Priority for vllm.ir.%s set to %s", |
| 428 | self.name, |
| 429 | lazy(lambda: [p.provider for p in self._priority_impls]), |
| 430 | ) |
| 431 | yield |
| 432 | finally: |
| 433 | self._priority_impls = old_priority_impls |
| 434 | |
| 435 | def supported_providers(self) -> list[str]: |
| 436 | return [p.provider for p in self.impls.values() if p.supported] |
| 437 | |
| 438 | @property |
| 439 | def has_input_generator(self) -> bool: |
| 440 | return self._input_generator is not None |
| 441 | |
| 442 | def register_input_generator(self, fn: InputGenerator) -> InputGenerator: |
| 443 | self._input_generator = fn |
| 444 | return fn |
| 445 | |
| 446 | def generate_inputs(self, **kwargs: Any) -> tuple[Any, ...]: |
| 447 | if self._input_generator is None: |
| 448 | raise RuntimeError( |
| 449 | f"No input generator registered for op '{self.name}'. " |
| 450 | f"Use @ir.ops.{self.name}.register_input_generator" |
| 451 | ) |
| 452 | return self._input_generator(**kwargs) |
| 453 | |
| 454 | def override_tolerance( |
| 455 | self, dtype: torch.dtype, *, atol: float, rtol: float |
| 456 | ) -> None: |
| 457 | self._tolerance_overrides[dtype] = {"atol": atol, "rtol": rtol} |
| 458 | |
| 459 | def get_tolerance(self, dtype: torch.dtype) -> dict[str, float]: |
| 460 | if dtype in self._tolerance_overrides: |
| 461 | return self._tolerance_overrides[dtype] |
| 462 | if dtype in DEFAULT_TOLERANCES: |
| 463 | return DEFAULT_TOLERANCES[dtype] |
| 464 | raise ValueError( |
| 465 | f"No tolerance defined for dtype {dtype} in op '{self.name}'. " |
| 466 | f"Use op.override_tolerance({dtype}, atol=..., rtol=...) " |
| 467 | f"or add {dtype} to DEFAULT_TOLERANCES." |
| 468 | ) |
| 469 | |
| 470 | |
| 471 | class IrOpInplace(IrOp): |
| 472 | """IR op with inplace support via maybe_inplace.""" |
| 473 | |
| 474 | maybe_inplace: "IrOpInplaceOverload" |
| 475 | allow_inplace: bool = True |
| 476 | |
| 477 | def __init__( |
| 478 | self, |
| 479 | name: str, |
| 480 | native_impl: Callable, |
| 481 | activations: list[str] | None = None, |
| 482 | registration_stack: list[str] | None = None, |
| 483 | ): |
| 484 | super().__init__(name, native_impl, activations, registration_stack) |
| 485 | |
| 486 | # Create the inplace overload |
| 487 | self.maybe_inplace = IrOpInplaceOverload(self) |
| 488 | |
| 489 | |
| 490 | class IrOpInplaceOverload: |
| 491 | def __init__(self, op: IrOp): |
| 492 | params, returns = op._schema_str.split(" -> ") |
| 493 | n_outputs = returns.count("Tensor") |
| 494 | |
| 495 | assert returns.count("Tensor") == len(op.activations), ( |
| 496 | "Inplace overload requires the same number of outputs as activations." |
| 497 | ) |
| 498 | |
| 499 | assert returns.count(",") == n_outputs - 1, ( |
| 500 | "Inplace overload only supports Tensor outputs for now." |
| 501 | ) |
| 502 | |
| 503 | self.op = op |
| 504 | self.name = f"{op.name}.maybe_inplace" |
| 505 | self._schema_str = infer_schema( |
| 506 | op.impls["native"].impl_fn, mutates_args=op.activations |
| 507 | ) |
| 508 | |
| 509 | # torch registration (resolve ``torch.ops`` subtree from ``lib.ns``) |
| 510 | lib = vllm_ir_torch_lib |
| 511 | lib.define(self.name + self._schema_str) |
| 512 | lib.impl(self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd") |
| 513 | # fake goes to default overload for now |
| 514 | lib._register_fake(self.name, self.op._fake_call) |
| 515 | |
| 516 | torch_ops = _torch_ops_subtree(lib) |
| 517 | assert hasattr(getattr(torch_ops, self.op.name), "maybe_inplace") |
| 518 | self.torch_op = getattr(torch_ops, self.op.name).maybe_inplace |
| 519 | |
| 520 | def __call__(self, *args, **kwargs) -> Any: |
| 521 | if not _ENABLE_TORCH_WRAP: |
| 522 | return self._inner_call(*args, **kwargs) |
| 523 | |
| 524 | return self.torch_op(*args, **kwargs) |
| 525 | |
| 526 | def _inner_call(self, *args, **kwargs) -> Any: |
| 527 | # Calling the maybe_inplace overload means we can use inplace impls directly. |
| 528 | impl = self.op.dispatch(*args, **kwargs) |
| 529 | return impl.impl_fn(*args, **kwargs) |
| 530 | |
| 531 | |
| 532 | class IrOpImpl: |
| 533 | def __init__( |
| 534 | self, |
| 535 | op: IrOp, |
| 536 | provider: str, |
| 537 | impl_fn: Callable, |
| 538 | supported: bool, |
| 539 | supports_args: Callable[..., bool] | None, |
| 540 | inplace: bool = False, |
| 541 | registration_stack: list[str] | None = None, |
| 542 | ): |
| 543 | assert provider not in op.impls, ( |
| 544 | f"Implementation for provider {provider} already registered." |
| 545 | ) |
| 546 | # Native also uses this path, so we allow it here. |
| 547 | assert provider == "native" or provider not in RESERVED_PROVIDERS, ( |
| 548 | f"Provider name {provider} is reserved." |
| 549 | ) |
| 550 | |
| 551 | # Enforce the exact same schema as the native implementation. |
| 552 | # This takes care of names, types, and defaults. |
| 553 | schema = infer_schema(impl_fn, mutates_args=[]) |
| 554 | if schema != op._schema_str: |
| 555 | raise ValueError( |
| 556 | f"Implementation for provider {provider} has schema '{schema}' which " |
| 557 | f"does not match native schema '{op._schema_str}' for op {op.name}." |
| 558 | ) |
| 559 | |
| 560 | if supports_args is not None: |
| 561 | if not callable(supports_args): |
| 562 | raise ValueError( |
| 563 | f"supports_args for provider {provider} must be a callable" |
| 564 | ) |
| 565 | |
| 566 | # We also manually validate the supports_args signature. |
| 567 | # Matching signatures allow faster dispatch on the hotpath. |
| 568 | |
| 569 | # Check that supports_args does not have keyword-only parameters |
| 570 | supports_args_signature = inspect.signature(supports_args) |
| 571 | params = supports_args_signature.parameters |
| 572 | if any(p.kind == inspect.Parameter.KEYWORD_ONLY for p in params.values()): |
| 573 | raise ValueError( |
| 574 | f"supports_args for provider {provider} " |
| 575 | f"cannot have keyword-only parameters" |
| 576 | ) |
| 577 | |
| 578 | # Check that supports_args has the same total number of parameters |
| 579 | op_params = op._py_signature.parameters |
| 580 | if len(params) != len(op_params): |
| 581 | raise ValueError( |
| 582 | f"supports_args for provider {provider} must have the same number " |
| 583 | f"of parameters ({len(params)}) as the native implementation " |
| 584 | f"({len(op_params)})" |
| 585 | ) |
| 586 | |
| 587 | # Check that names and defaults match for supports_args |
| 588 | for p, op_p in zip(params.values(), op_params.values()): |
| 589 | if p.name != op_p.name: |
| 590 | raise ValueError( |
| 591 | f"supports_args for provider {provider} has parameter " |
| 592 | f"'{p.name}' which does not match native parameter " |
| 593 | f"'{op_p.name}'" |
| 594 | ) |
| 595 | if p.default != op_p.default: |
| 596 | raise ValueError( |
| 597 | f"supports_args for provider {provider} has parameter " |
| 598 | f"'{p.name}' with default {p.default} which does not match " |
| 599 | f"native default {op_p.default}'" |
| 600 | ) |
| 601 | |
| 602 | if inplace: |
| 603 | assert op.allow_inplace, ( |
| 604 | f"Inplace implementation cannot be registered for op {op.name}" |
| 605 | f" that does not allow inplace." |
| 606 | ) |
| 607 | |
| 608 | self.op = op |
| 609 | self.provider = provider |
| 610 | self.impl_fn = impl_fn |
| 611 | self.supported = supported |
| 612 | self._supports_args = supports_args |
| 613 | self.inplace = inplace |
| 614 | self._registration_stack = registration_stack or [] |
| 615 | |
| 616 | @property |
| 617 | def supports_all_args(self) -> bool: |
| 618 | """Check if this implementation supports all args unconditionally.""" |
| 619 | return self._supports_args is None |
| 620 | |
| 621 | def supports_args(self, *args, **kwargs) -> bool: |
| 622 | if self._supports_args is None: |
| 623 | return True |
| 624 | |
| 625 | return self._supports_args(*args, **kwargs) |
| 626 | |
| 627 | @weak_cache |
| 628 | def uuid(self): |
| 629 | """Compile-time hash to uniquely determine whether the implementation has |
| 630 | changed. Used by vllm-compile hash mechanism and torch.compile lowering pass |
| 631 | uuid to control the vLLM compile cache and AOTAutograd/Inductor caches |
| 632 | respectively. |
| 633 | |
| 634 | Source file contents do not change so we cache uuid. |
| 635 | TODO(luka): Cache the file hash as multiple impls are likely in the same file. |
| 636 | """ |
| 637 | sources = [Path(inspect.getfile(self.impl_fn))] |
| 638 | return hash_source(*sources) |
| 639 | |
| 640 | def func_impl_fn(self, *args, **kwargs) -> Any: |
| 641 | """Copy any inputs in activations if this is an inplace impl, |
| 642 | to ensure functional semantics. |
| 643 | """ |
| 644 | if not self.inplace: |
| 645 | return self.impl_fn(*args, **kwargs) |
| 646 | |
| 647 | # copy activations to ensure functional semantics |
| 648 | new_args = list(args) |
| 649 | for i in self.op.activation_indices: |
| 650 | assert isinstance(args[i], torch.Tensor) |
| 651 | new_args[i] = args[i].clone() |
| 652 | |
| 653 | return self.impl_fn(*new_args, **kwargs) |