wehub-resource-sync
665 lines
24 KiB
Python
665 lines
24 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import contextlib
|
|
import inspect
|
|
import traceback
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
from typing import Any, ClassVar, Literal, overload
|
|
|
|
import regex as re
|
|
import torch
|
|
from torch.library import Library, infer_schema
|
|
|
|
from vllm.ir.tolerances import DEFAULT_TOLERANCES, ToleranceSpec
|
|
from vllm.ir.util import hash_source, weak_cache
|
|
from vllm.logger import init_logger
|
|
from vllm.logging_utils import lazy, tensors_str_no_data
|
|
|
|
InputGenerator = Callable[..., tuple[Any, ...]]
|
|
|
|
vllm_ir_torch_lib = Library("vllm_ir", "FRAGMENT") # IR op lib; monkeypatch in tests.
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _torch_ops_subtree(lib: Any) -> Any:
|
|
"""``torch.ops`` subtree for ``lib.ns``; fall back if doc mocks replace ``ns``."""
|
|
ns = getattr(lib, "ns", None)
|
|
if isinstance(ns, str):
|
|
return getattr(torch.ops, ns)
|
|
return torch.ops.vllm_ir
|
|
|
|
|
|
_NAME_PATTERN = re.compile(r"^[a-z_][a-z_0-9]*$")
|
|
|
|
RESERVED_PROVIDERS = ["native", "unfused"]
|
|
"""Providers that are reserved and cannot be used for custom implementations."""
|
|
|
|
|
|
def _validate_name(name: str, entity_type: str) -> None:
|
|
"""Validate that a name matches the required pattern `[a-z_][a-z_0-9]*`."""
|
|
if not _NAME_PATTERN.match(name):
|
|
raise ValueError(
|
|
f"{entity_type} name '{name}' is invalid. "
|
|
f"Names must start with a letter or underscore, "
|
|
f"followed by lowercase letters, underscores, or digits only."
|
|
)
|
|
|
|
|
|
_ENABLE_TORCH_WRAP: bool = True
|
|
"""Global override flag to control torch op layer wrapping."""
|
|
|
|
|
|
def set_default_torch_wrap(enable: bool = True) -> None:
|
|
"""
|
|
Permanently set the torch wrap flag.
|
|
"""
|
|
global _ENABLE_TORCH_WRAP
|
|
_ENABLE_TORCH_WRAP = enable
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def enable_torch_wrap(enable: bool = True):
|
|
"""
|
|
Context manager to enable/disable torch custom op wrapping for vLLM IR ops.
|
|
When torch wrapping is disabled, the torch custom op layer is skipped
|
|
and IR ops dispatch directly to the implementation.
|
|
Helpful for avoiding torch dispatch overhead in eager mode
|
|
and avoiding the need for lowering for platforms not using Inductor.
|
|
"""
|
|
|
|
global _ENABLE_TORCH_WRAP
|
|
old = _ENABLE_TORCH_WRAP
|
|
try:
|
|
_ENABLE_TORCH_WRAP = enable
|
|
yield
|
|
finally:
|
|
_ENABLE_TORCH_WRAP = old
|
|
|
|
|
|
# 0-param decorator overload (no inplace)
|
|
@overload
|
|
def register_op(f: Callable[..., Any]) -> "IrOp": ...
|
|
|
|
|
|
# parametrized decorator with allow_inplace=False (default)
|
|
@overload
|
|
def register_op(
|
|
*,
|
|
name: str | None = None,
|
|
activations: list[str] | None = None,
|
|
allow_inplace: Literal[False] = False,
|
|
) -> Callable[[Callable[..., Any]], "IrOp"]: ...
|
|
|
|
|
|
# parametrized decorator with allow_inplace=True
|
|
@overload
|
|
def register_op(
|
|
*,
|
|
name: str | None = None,
|
|
activations: list[str] | None = None,
|
|
allow_inplace: Literal[True],
|
|
) -> Callable[[Callable[..., Any]], "IrOpInplace"]: ...
|
|
|
|
|
|
def register_op(
|
|
f: Callable | None = None,
|
|
*,
|
|
name: str | None = None,
|
|
activations: list[str] | None = None,
|
|
allow_inplace: bool = False,
|
|
) -> "IrOp | Callable[[Callable], IrOp]":
|
|
"""
|
|
Register a new vLLM IR op.
|
|
|
|
Args:
|
|
f: the native implementation of the op
|
|
name: the name of the op, defaults to the function name
|
|
activations: list of activation params, defaults to params starting with 'x'
|
|
allow_inplace: add a maybe_inplace overload that allows inplace impls
|
|
|
|
Returns:
|
|
the IrOp object if f is provided, otherwise a decorator
|
|
|
|
Example usage:
|
|
```python
|
|
@vllm.ir.register_op
|
|
def my_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x + y
|
|
|
|
|
|
@vllm.ir.register_op(name="custom_mul")
|
|
def multiply(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x * y"""
|
|
|
|
def decorator(_f: Callable):
|
|
op_name: str = _f.__name__ if name is None else name
|
|
_validate_name(op_name, "Op")
|
|
assert op_name not in IrOp.registry, f"Op '{op_name}' is already registered."
|
|
# Slice out the decorator function frames from the stack
|
|
stack = traceback.format_stack()[:-2]
|
|
if allow_inplace:
|
|
op: IrOp = IrOpInplace(op_name, _f, activations, stack)
|
|
else:
|
|
op = IrOp(op_name, _f, activations, stack)
|
|
IrOp.registry[op_name] = op
|
|
return op
|
|
|
|
if f is not None:
|
|
return decorator(f)
|
|
|
|
return decorator
|
|
|
|
|
|
class IrOp:
|
|
registry: ClassVar[dict[str, "IrOp"]] = {}
|
|
|
|
name: str
|
|
impls: dict[str, "IrOpImpl"]
|
|
allow_inplace: bool = False
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
native_impl: Callable,
|
|
activations: list[str] | None = None,
|
|
registration_stack: list[str] | None = None,
|
|
):
|
|
self._py_signature = inspect.signature(native_impl)
|
|
if any(
|
|
p.kind == inspect.Parameter.KEYWORD_ONLY
|
|
for p in self._py_signature.parameters.values()
|
|
):
|
|
raise ValueError(
|
|
f"Op {name} has keyword-only arguments which are not currently "
|
|
f"supported. That's because kwargs are not allowed during lowering."
|
|
)
|
|
|
|
# By convention, we consider parameters starting with 'x' as activations.
|
|
if activations is None:
|
|
activations = [
|
|
p.name
|
|
for p in self._py_signature.parameters.values()
|
|
if p.name.startswith("x")
|
|
]
|
|
|
|
self.name = name
|
|
self._docstring = inspect.getdoc(native_impl) or ""
|
|
self._registration_stack = registration_stack or []
|
|
self.impls: dict[str, IrOpImpl] = {}
|
|
self.activations = activations
|
|
self.activation_indices = [
|
|
i
|
|
for i, p in enumerate(self._py_signature.parameters.values())
|
|
if p.name in activations
|
|
]
|
|
self._priority_impls: list[IrOpImpl] = []
|
|
self._schema_str = infer_schema(native_impl, mutates_args=[])
|
|
self._input_generator: InputGenerator | None = None
|
|
self._tolerance_overrides: ToleranceSpec = {}
|
|
|
|
# native implementation
|
|
self.impls["native"] = IrOpImpl(
|
|
self,
|
|
"native",
|
|
native_impl,
|
|
# always supported
|
|
supported=True,
|
|
supports_args=None,
|
|
registration_stack=self._registration_stack,
|
|
)
|
|
|
|
# By default, fake routes directly to native,
|
|
# can be overridden by register_fake
|
|
self._fake_fn = native_impl
|
|
|
|
# torch registration (resolve ``torch.ops`` subtree from ``lib.ns``)
|
|
lib = vllm_ir_torch_lib
|
|
lib.define(self.name + self._schema_str)
|
|
# CompositeExplicitAutograd is not decomposed
|
|
# by ATen IR normalization in AOTAutograd
|
|
lib.impl(self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd")
|
|
lib._register_fake(self.name, self._fake_call)
|
|
torch_ops = _torch_ops_subtree(lib)
|
|
assert hasattr(torch_ops, name)
|
|
self.torch_op: torch._ops.OpOverload = getattr(torch_ops, name).default
|
|
|
|
def register_fake(self, fn: Callable) -> Callable:
|
|
"""
|
|
Register a fake impl for the torch custom op. If this method is not called,
|
|
the native implementation is used directly for the fake implementation.
|
|
"""
|
|
self._fake_fn = fn
|
|
return fn
|
|
|
|
def _fake_call(self, *args, **kwargs) -> Any:
|
|
"""
|
|
Call to the fake implementation of the op. We use indirection because we want
|
|
users to be able to register fake later but also want it to fall back to native
|
|
directly by default, instead of going through the dispatching mechanism.
|
|
"""
|
|
return self._fake_fn(*args, **kwargs)
|
|
|
|
def register_impl(
|
|
self,
|
|
provider: str,
|
|
*,
|
|
supported: bool = True,
|
|
supports_args: Callable[..., bool] | None = None,
|
|
inplace: bool = False,
|
|
) -> Callable[[Callable[..., Any]], "IrOpImpl"]:
|
|
"""
|
|
Register an implementation for this custom op.
|
|
Args:
|
|
provider: The name of the provider, must be unique.
|
|
supported: Static support check, use this to check platform support.
|
|
supports_args: Dynamic arg support check, used for types and shapes.
|
|
inplace: Does this op reuse activation input memory for outputs
|
|
|
|
Returns:
|
|
A decorator that registers the implementation.
|
|
|
|
The decorated function must have the same semantics and signature as
|
|
the native implementation.
|
|
|
|
The provider name must be unique and not one of the RESERVED_PROVIDERS.
|
|
The supported and supports_args parameters should not be used to implement
|
|
custom enablement logic based on global state (e.g. environment variables).
|
|
Instead, supported param should only be used to check for platform support
|
|
(e.g. whether a specific hardware or library is available).
|
|
supports_args should be used to check whether the provided arguments are
|
|
compatible with the implementation.
|
|
For custom enablement logic, set op impl priority.
|
|
|
|
Example:
|
|
```python
|
|
@my_op.register_impl("my_provider", supported=torch.cuda.is_available())
|
|
def my_provider_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: ...
|
|
```
|
|
|
|
"""
|
|
assert provider not in RESERVED_PROVIDERS, (
|
|
f"Provider name {provider} is reserved."
|
|
)
|
|
_validate_name(provider, "Provider")
|
|
|
|
def _register_impl(f: Callable):
|
|
# Slice out the decorator function from the stack
|
|
stack = traceback.format_stack()[:-1]
|
|
impl = IrOpImpl(self, provider, f, supported, supports_args, inplace, stack)
|
|
self.impls[provider] = impl
|
|
|
|
if self.get_priority():
|
|
logger.warning(
|
|
"Warning: registering new impl %s for op %s while priority is set.",
|
|
provider,
|
|
self.name,
|
|
)
|
|
|
|
return impl
|
|
|
|
return _register_impl
|
|
|
|
def _inner_call(self, *args, **kwargs) -> Any:
|
|
"""
|
|
Eager call to torch op lands here. When torch wrapping is disabled,
|
|
__call__ routes straight here instead of going through torch op dispatching.
|
|
"""
|
|
impl = self.dispatch(*args, **kwargs)
|
|
|
|
# Default overload must be functional,
|
|
# use func_impl_fn to correctly handle inplace impls.
|
|
return impl.func_impl_fn(*args, **kwargs)
|
|
|
|
def apply_arg_defaults(self, args) -> tuple:
|
|
"""
|
|
Return args with default values applied.
|
|
Defaults are taken from the native implementation signature.
|
|
|
|
SHOULD NOT BE USED IN THE DISPATCH PATH (SLOW).
|
|
Only for Inductor lowering.
|
|
"""
|
|
bound_args = self._py_signature.bind(*args)
|
|
bound_args.apply_defaults()
|
|
return bound_args.args
|
|
|
|
def dispatch(self, *args, **kwargs) -> "IrOpImpl":
|
|
"""
|
|
Dispatch to the appropriate implementation based on current priority
|
|
and argument support checks. Returns the selected IrOpImpl.
|
|
|
|
THIS FUNCTION IS ON THE HOT PATH (OP DISPATCH), MUST BE FAST.
|
|
"""
|
|
if not self._priority_impls:
|
|
if not torch.compiler.is_compiling():
|
|
# Logging not compatible with Dynamo tracing
|
|
# (this code is exposed when torch wrapping is disabled)
|
|
logger.warning_once(
|
|
"Priority not set for op %s, using native implementation.",
|
|
self.name,
|
|
)
|
|
return self.impls["native"]
|
|
|
|
for impl in self._priority_impls:
|
|
if not impl.supported:
|
|
raise ValueError(
|
|
f"Implementation {impl.provider} for op {self.name} not supported. "
|
|
f"All implementations in priority list must be supported."
|
|
)
|
|
if impl.supports_args(*args, **kwargs):
|
|
return impl
|
|
|
|
if not torch.compiler.is_compiling():
|
|
logger.debug(
|
|
"Skipping provider %s because it does not support "
|
|
"%s with args=%s kwargs=%s",
|
|
impl.provider,
|
|
self.name,
|
|
lazy(lambda: tensors_str_no_data(args)),
|
|
lazy(lambda: tensors_str_no_data(kwargs)),
|
|
)
|
|
|
|
raise RuntimeError(
|
|
"Priority set incorrectly: the last implementation must "
|
|
"support all args (can be native). This is likely an internal bug"
|
|
)
|
|
|
|
def __call__(self, *args, **kwargs) -> Any:
|
|
if not _ENABLE_TORCH_WRAP:
|
|
return self._inner_call(*args, **kwargs)
|
|
|
|
return self.torch_op(*args, **kwargs)
|
|
|
|
def __repr__(self) -> str:
|
|
"""Return unambiguous string representation."""
|
|
return f"IrOp('{self.name}')"
|
|
|
|
def __str__(self) -> str:
|
|
"""Return human-readable string representation using docstring."""
|
|
if not self._docstring:
|
|
return f"IrOp('{self.name}')"
|
|
first_line = self._docstring.split("\n")[0].strip()
|
|
return f"IrOp('{self.name}') - {first_line}"
|
|
|
|
def get_priority(self) -> list[str]:
|
|
"""Get the current dispatch priority for implementations for this op."""
|
|
return [p.provider for p in self._priority_impls]
|
|
|
|
def _filter_priority_impls(self, priority: list[str]) -> list["IrOpImpl"]:
|
|
assert all(p in self.impls for p in priority), (
|
|
"All providers in priority must be registered implementations."
|
|
)
|
|
filtered_impls: list[IrOpImpl] = []
|
|
for p in priority:
|
|
impl = self.impls[p]
|
|
if not impl.supported:
|
|
# Skip unsupported implementations
|
|
continue
|
|
|
|
filtered_impls.append(impl)
|
|
|
|
# If all args are supported, skip other implementations
|
|
if impl.supports_all_args:
|
|
return filtered_impls
|
|
|
|
logger.warning_once(
|
|
"Op %s: No implementation in priority list supports all args, "
|
|
"execution fallback to native is possible. To silence this warning, "
|
|
"explicitly add 'native' to the end of the priority list",
|
|
self.name,
|
|
)
|
|
filtered_impls.append(self.impls["native"])
|
|
return filtered_impls
|
|
|
|
def set_default(self, priority: list[str]) -> None:
|
|
"""
|
|
Permanently set the dispatch priority for this op. Use this for
|
|
process-lifetime setup (e.g., worker startup). For scoped overrides,
|
|
use ``set_priority`` instead.
|
|
"""
|
|
self._priority_impls = self._filter_priority_impls(priority)
|
|
logger.debug(
|
|
"Priority for vllm.ir.%s set to %s",
|
|
self.name,
|
|
lazy(lambda: [p.provider for p in self._priority_impls]),
|
|
)
|
|
|
|
@contextlib.contextmanager
|
|
def set_priority(self, priority: list[str]):
|
|
"""
|
|
Context manager to set the dispatch priority for implementations for this op.
|
|
"""
|
|
old_priority_impls = self._priority_impls
|
|
try:
|
|
self._priority_impls = self._filter_priority_impls(priority)
|
|
logger.debug(
|
|
"Priority for vllm.ir.%s set to %s",
|
|
self.name,
|
|
lazy(lambda: [p.provider for p in self._priority_impls]),
|
|
)
|
|
yield
|
|
finally:
|
|
self._priority_impls = old_priority_impls
|
|
|
|
def supported_providers(self) -> list[str]:
|
|
return [p.provider for p in self.impls.values() if p.supported]
|
|
|
|
@property
|
|
def has_input_generator(self) -> bool:
|
|
return self._input_generator is not None
|
|
|
|
def register_input_generator(self, fn: InputGenerator) -> InputGenerator:
|
|
self._input_generator = fn
|
|
return fn
|
|
|
|
def generate_inputs(self, **kwargs: Any) -> tuple[Any, ...]:
|
|
if self._input_generator is None:
|
|
raise RuntimeError(
|
|
f"No input generator registered for op '{self.name}'. "
|
|
f"Use @ir.ops.{self.name}.register_input_generator"
|
|
)
|
|
return self._input_generator(**kwargs)
|
|
|
|
def override_tolerance(
|
|
self, dtype: torch.dtype, *, atol: float, rtol: float
|
|
) -> None:
|
|
self._tolerance_overrides[dtype] = {"atol": atol, "rtol": rtol}
|
|
|
|
def get_tolerance(self, dtype: torch.dtype) -> dict[str, float]:
|
|
if dtype in self._tolerance_overrides:
|
|
return self._tolerance_overrides[dtype]
|
|
if dtype in DEFAULT_TOLERANCES:
|
|
return DEFAULT_TOLERANCES[dtype]
|
|
raise ValueError(
|
|
f"No tolerance defined for dtype {dtype} in op '{self.name}'. "
|
|
f"Use op.override_tolerance({dtype}, atol=..., rtol=...) "
|
|
f"or add {dtype} to DEFAULT_TOLERANCES."
|
|
)
|
|
|
|
|
|
class IrOpInplace(IrOp):
|
|
"""IR op with inplace support via maybe_inplace."""
|
|
|
|
maybe_inplace: "IrOpInplaceOverload"
|
|
allow_inplace: bool = True
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
native_impl: Callable,
|
|
activations: list[str] | None = None,
|
|
registration_stack: list[str] | None = None,
|
|
):
|
|
super().__init__(name, native_impl, activations, registration_stack)
|
|
|
|
# Create the inplace overload
|
|
self.maybe_inplace = IrOpInplaceOverload(self)
|
|
|
|
|
|
class IrOpInplaceOverload:
|
|
def __init__(self, op: IrOp):
|
|
params, returns = op._schema_str.split(" -> ")
|
|
n_outputs = returns.count("Tensor")
|
|
|
|
assert returns.count("Tensor") == len(op.activations), (
|
|
"Inplace overload requires the same number of outputs as activations."
|
|
)
|
|
|
|
assert returns.count(",") == n_outputs - 1, (
|
|
"Inplace overload only supports Tensor outputs for now."
|
|
)
|
|
|
|
self.op = op
|
|
self.name = f"{op.name}.maybe_inplace"
|
|
self._schema_str = infer_schema(
|
|
op.impls["native"].impl_fn, mutates_args=op.activations
|
|
)
|
|
|
|
# torch registration (resolve ``torch.ops`` subtree from ``lib.ns``)
|
|
lib = vllm_ir_torch_lib
|
|
lib.define(self.name + self._schema_str)
|
|
lib.impl(self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd")
|
|
# fake goes to default overload for now
|
|
lib._register_fake(self.name, self.op._fake_call)
|
|
|
|
torch_ops = _torch_ops_subtree(lib)
|
|
assert hasattr(getattr(torch_ops, self.op.name), "maybe_inplace")
|
|
self.torch_op = getattr(torch_ops, self.op.name).maybe_inplace
|
|
|
|
def __call__(self, *args, **kwargs) -> Any:
|
|
if not _ENABLE_TORCH_WRAP:
|
|
return self._inner_call(*args, **kwargs)
|
|
|
|
return self.torch_op(*args, **kwargs)
|
|
|
|
def _inner_call(self, *args, **kwargs) -> Any:
|
|
# Calling the maybe_inplace overload means we can use inplace impls directly.
|
|
impl = self.op.dispatch(*args, **kwargs)
|
|
return impl.impl_fn(*args, **kwargs)
|
|
|
|
|
|
class IrOpImpl:
|
|
def __init__(
|
|
self,
|
|
op: IrOp,
|
|
provider: str,
|
|
impl_fn: Callable,
|
|
supported: bool,
|
|
supports_args: Callable[..., bool] | None,
|
|
inplace: bool = False,
|
|
registration_stack: list[str] | None = None,
|
|
):
|
|
assert provider not in op.impls, (
|
|
f"Implementation for provider {provider} already registered."
|
|
)
|
|
# Native also uses this path, so we allow it here.
|
|
assert provider == "native" or provider not in RESERVED_PROVIDERS, (
|
|
f"Provider name {provider} is reserved."
|
|
)
|
|
|
|
# Enforce the exact same schema as the native implementation.
|
|
# This takes care of names, types, and defaults.
|
|
schema = infer_schema(impl_fn, mutates_args=[])
|
|
if schema != op._schema_str:
|
|
raise ValueError(
|
|
f"Implementation for provider {provider} has schema '{schema}' which "
|
|
f"does not match native schema '{op._schema_str}' for op {op.name}."
|
|
)
|
|
|
|
if supports_args is not None:
|
|
if not callable(supports_args):
|
|
raise ValueError(
|
|
f"supports_args for provider {provider} must be a callable"
|
|
)
|
|
|
|
# We also manually validate the supports_args signature.
|
|
# Matching signatures allow faster dispatch on the hotpath.
|
|
|
|
# Check that supports_args does not have keyword-only parameters
|
|
supports_args_signature = inspect.signature(supports_args)
|
|
params = supports_args_signature.parameters
|
|
if any(p.kind == inspect.Parameter.KEYWORD_ONLY for p in params.values()):
|
|
raise ValueError(
|
|
f"supports_args for provider {provider} "
|
|
f"cannot have keyword-only parameters"
|
|
)
|
|
|
|
# Check that supports_args has the same total number of parameters
|
|
op_params = op._py_signature.parameters
|
|
if len(params) != len(op_params):
|
|
raise ValueError(
|
|
f"supports_args for provider {provider} must have the same number "
|
|
f"of parameters ({len(params)}) as the native implementation "
|
|
f"({len(op_params)})"
|
|
)
|
|
|
|
# Check that names and defaults match for supports_args
|
|
for p, op_p in zip(params.values(), op_params.values()):
|
|
if p.name != op_p.name:
|
|
raise ValueError(
|
|
f"supports_args for provider {provider} has parameter "
|
|
f"'{p.name}' which does not match native parameter "
|
|
f"'{op_p.name}'"
|
|
)
|
|
if p.default != op_p.default:
|
|
raise ValueError(
|
|
f"supports_args for provider {provider} has parameter "
|
|
f"'{p.name}' with default {p.default} which does not match "
|
|
f"native default {op_p.default}'"
|
|
)
|
|
|
|
if inplace:
|
|
assert op.allow_inplace, (
|
|
f"Inplace implementation cannot be registered for op {op.name}"
|
|
f" that does not allow inplace."
|
|
)
|
|
|
|
self.op = op
|
|
self.provider = provider
|
|
self.impl_fn = impl_fn
|
|
self.supported = supported
|
|
self._supports_args = supports_args
|
|
self.inplace = inplace
|
|
self._registration_stack = registration_stack or []
|
|
|
|
@property
|
|
def supports_all_args(self) -> bool:
|
|
"""Check if this implementation supports all args unconditionally."""
|
|
return self._supports_args is None
|
|
|
|
def supports_args(self, *args, **kwargs) -> bool:
|
|
if self._supports_args is None:
|
|
return True
|
|
|
|
return self._supports_args(*args, **kwargs)
|
|
|
|
@weak_cache
|
|
def uuid(self):
|
|
"""
|
|
Compile-time hash to uniquely determine whether the implementation has changed.
|
|
Used by vllm-compile hash mechanism and torch.compile lowering pass uuid to
|
|
control the vLLM compile cache and AOTAutograd/Inductor caches respectively.
|
|
|
|
Source file contents do not change so we cache uuid.
|
|
TODO(luka): Cache the file hash as multiple impls are likely in the same file.
|
|
"""
|
|
sources = [Path(inspect.getfile(self.impl_fn))]
|
|
return hash_source(*sources)
|
|
|
|
def func_impl_fn(self, *args, **kwargs) -> Any:
|
|
"""
|
|
Copy any inputs in activations if this is an inplace impl,
|
|
to ensure functional semantics.
|
|
"""
|
|
if not self.inplace:
|
|
return self.impl_fn(*args, **kwargs)
|
|
|
|
# copy activations to ensure functional semantics
|
|
new_args = list(args)
|
|
for i in self.op.activation_indices:
|
|
assert isinstance(args[i], torch.Tensor)
|
|
new_args[i] = args[i].clone()
|
|
|
|
return self.impl_fn(*new_args, **kwargs)
|