diff --git a/tests/conftest.py b/tests/conftest.py index 779bd475f3..5ff58f94e9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1686,3 +1686,44 @@ def disable_log_dedup(monkeypatch): logger._print_warning_once = original_print_warning_once logger._print_info_once = original_print_info_once logger._print_debug_once = original_print_debug_once + + +@pytest.fixture(scope="function") +def fake_vllm_ir(monkeypatch): + """ + Pytest fixture to allow isolated IR op registration in tests. + + Replaces IrOp.registry with an empty dict and swaps ``vllm_ir_torch_lib`` for a + fresh ``Library`` with a unique namespace per test (see ``Library.ns``). + + Torch keeps registrations for the process lifetime; reusing the fragment + name ``vllm_ir`` and defining the same op string again can segfault. A + random library name keeps each fixture run on a disjoint namespace. + + The test Library is kept alive until after monkeypatch teardown so PyTorch's + C++ state is not freed while references may still exist. + + Usage: + def test_my_ir_op(fake_vllm_ir): + @vllm.ir.register_op + def my_test_op(x: torch.Tensor) -> torch.Tensor: + return x * 2 + + result = my_test_op(torch.tensor([1, 2, 3])) + # Registry and library cleaned up automatically after the test + """ + import secrets + + from torch.library import Library + from vllm.ir.op import IrOp + + monkeypatch.setattr(IrOp, "registry", {}) + + # Keep a local reference so the Library is not GC'd before monkeypatch + # teardown restores the original reference. + test_lib = Library(f"vllm_ir_{secrets.token_hex(8)}", "FRAGMENT") + monkeypatch.setattr("vllm.ir.op.vllm_ir_torch_lib", test_lib) + + yield + + del test_lib diff --git a/tests/ir/test_op.py b/tests/ir/test_op.py index 3576e5aef8..38fb7247a5 100644 --- a/tests/ir/test_op.py +++ b/tests/ir/test_op.py @@ -13,20 +13,40 @@ from torch.fx.experimental.proxy_tensor import make_fx import vllm.ir.op from vllm.ir.op import RESERVED_PROVIDERS, IrOp, IrOpImpl -# This should not exist -assert "_custom_add" not in IrOp.registry - class CustomError(Exception): pass -@vllm.ir.register_op(allow_inplace=True) -def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - return x + y +@pytest.fixture +def custom_add_op(fake_vllm_ir): + """Register ``_custom_add`` plus impl_a, impl_b, impl_even for this test.""" + + @vllm.ir.register_op(allow_inplace=True) + def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x + y + + @_custom_add.register_impl("impl_a") + def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x + y + 10 + + @_custom_add.register_impl("impl_b", inplace=True) + def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + """Computes x+y+20""" + x.add_(y) + x.add_(20) + return x + + @_custom_add.register_impl( + "impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0 + ) + def impl_even(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x + y + 50 + + return _custom_add -def test_registration_overloads(): +def test_registration_overloads(fake_vllm_ir): assert all( n not in IrOp.registry for n in ["_custom_sub", "_custom_mul", "_custom_div"] ) @@ -63,7 +83,7 @@ def test_registration_overloads(): return x * y - 100 -def test_no_kw_only_args(): +def test_no_kw_only_args(fake_vllm_ir): # kw-only args not supported with pytest.raises(ValueError, match="keyword-only arguments"): @@ -78,18 +98,23 @@ def test_no_kw_only_args(): class TestIrOpCustomAdd: # Registration invariants - def test_decorated_object(self): + def test_decorated_object(self, custom_add_op): """Make sure that referring directly to an op is correct""" + _custom_add = custom_add_op assert isinstance(_custom_add, IrOp) assert "_custom_add" in IrOp.registry assert _custom_add is IrOp.registry["_custom_add"] - def test_torch_op_is_registered(self): - assert hasattr(torch.ops.vllm_ir, "_custom_add") - assert callable(torch.ops.vllm_ir._custom_add.default) + def test_torch_op_is_registered(self, custom_add_op): + _custom_add = custom_add_op + torch_ops = getattr(torch.ops, vllm.ir.op.vllm_ir_torch_lib.ns) + assert hasattr(torch_ops, "_custom_add") + assert callable(torch_ops._custom_add.default) + assert _custom_add.torch_op is torch_ops._custom_add.default # Semantic correctness - def test_semantics_match_native(self): + def test_semantics_match_native(self, custom_add_op): + _custom_add = custom_add_op x = torch.randn(4, 5) y = torch.randn(4, 5) @@ -103,7 +128,9 @@ class TestIrOpCustomAdd: # Implementation registration # ------------------------- - def test_register_impl_is_non_intrusive(self): + def test_register_impl_is_non_intrusive(self, custom_add_op): + _custom_add = custom_add_op + @_custom_add.register_impl("dummy_provider") def dummy_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return x + y + 123 @@ -117,7 +144,8 @@ class TestIrOpCustomAdd: # Native semantics must still hold torch.testing.assert_close(_custom_add(x, y), x + y) - def test_schema_contains_tensor_signature(self): + def test_schema_contains_tensor_signature(self, custom_add_op): + _custom_add = custom_add_op schema = _custom_add._schema_str assert "Tensor" in schema @@ -131,10 +159,19 @@ class TestIrOpCustomAdd: @pytest.mark.parametrize("symbolic_trace", [True, False]) @pytest.mark.parametrize("overload", ["default", "maybe_inplace"]) def test_trace_sees_single_custom_op( - self, symbolic_trace: bool, enable_torch_wrap: bool, overload: str + self, + custom_add_op, + symbolic_trace: bool, + enable_torch_wrap: bool, + overload: str, ): + _custom_add = custom_add_op op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace - torch_op = getattr(torch.ops.vllm_ir._custom_add, overload) + torch_op = ( + _custom_add.torch_op + if overload == "default" + else _custom_add.maybe_inplace.torch_op + ) def fn(x, y): return op_fn(x, y) @@ -175,30 +212,13 @@ class TestIrOpCustomAdd: assert len(ir_nodes) == 1, gm.code -@_custom_add.register_impl("impl_a") -def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - return x + y + 10 - - -@_custom_add.register_impl("impl_b", inplace=True) -def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - """Computes x+y+20""" - x.add_(y) - x.add_(20) - return x - - -@_custom_add.register_impl("impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0) -def impl_even(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - return x + y + 50 - - class TestIrOpImplDispatch: - def test_register_impl(self): + def test_register_impl(self, custom_add_op): + _custom_add = custom_add_op assert "impl_a" in _custom_add.impls impl = _custom_add.impls["impl_a"] - assert impl is impl_a + assert impl is _custom_add.impls["impl_a"] assert impl.op is _custom_add assert impl.provider == "impl_a" assert callable(impl.impl_fn) @@ -211,14 +231,15 @@ class TestIrOpImplDispatch: return x + y + 30 # Check the original impl is still intact - assert _custom_add.impls["impl_a"] is impl_a + assert _custom_add.impls["impl_a"] is impl # Check support all args - assert impl_a.supports_all_args - assert impl_b.supports_all_args - assert not impl_even.supports_all_args + assert _custom_add.impls["impl_a"].supports_all_args + assert _custom_add.impls["impl_b"].supports_all_args + assert not _custom_add.impls["impl_even"].supports_all_args - def test_reserved_provider_rejected(self): + def test_reserved_provider_rejected(self, custom_add_op): + _custom_add = custom_add_op for provider in RESERVED_PROVIDERS: with pytest.raises(AssertionError): @@ -226,7 +247,8 @@ class TestIrOpImplDispatch: def bad_impl(x, y): return x + y - def test_set_priority_scoped(self): + def test_set_priority_scoped(self, custom_add_op): + _custom_add = custom_add_op assert _custom_add.get_priority() == [] with _custom_add.set_priority(["impl_even", "impl_b"]): @@ -251,20 +273,25 @@ class TestIrOpImplDispatch: assert _custom_add.get_priority() == [] @pytest.mark.parametrize("overload", ["default", "maybe_inplace"]) - def test_dispatch_priority_order(self, overload: str): + def test_dispatch_priority_order(self, custom_add_op, overload: str): + _custom_add = custom_add_op op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace - torch_op = getattr(torch.ops.vllm_ir._custom_add, overload) + torch_op = ( + _custom_add.torch_op + if overload == "default" + else _custom_add.maybe_inplace.torch_op + ) x = torch.tensor(1, dtype=torch.int32) y = torch.tensor(2, dtype=torch.int32) with _custom_add.set_priority(["impl_b", "impl_a"]): - assert _custom_add.dispatch(x, y) is impl_b + assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_b"] out1 = op_fn(x.clone(), y) out2 = torch_op(x.clone(), y) with _custom_add.set_priority(["impl_a"]): - assert _custom_add.dispatch(x, y) is impl_a + assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_a"] out3 = op_fn(x.clone(), y) out4 = torch_op(x.clone(), y) @@ -275,7 +302,9 @@ class TestIrOpImplDispatch: assert out3.item() == 1 + 2 + 10 assert out4.item() == 1 + 2 + 10 - def test_unsupported_impl_filtered(self): + def test_unsupported_impl_filtered(self, custom_add_op): + _custom_add = custom_add_op + @_custom_add.register_impl("impl_unsupported", supported=False) def impl_unsupported(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return x + y + 999 @@ -291,8 +320,9 @@ class TestIrOpImplDispatch: assert out.item() == 1 + 2 + 10 def test_supports_args_runtime_dispatch_and_warning( - self, caplog_vllm: pytest.LogCaptureFixture + self, custom_add_op, caplog_vllm: pytest.LogCaptureFixture ): + _custom_add = custom_add_op x1 = torch.ones((2, 2), dtype=torch.int32) y1 = torch.full((2, 2), 2, dtype=torch.int32) @@ -312,7 +342,7 @@ class TestIrOpImplDispatch: # Check dispatching assert _custom_add.get_priority() == ["impl_even", "native"] - assert _custom_add.dispatch(x1, y1) is impl_even + assert _custom_add.dispatch(x1, y1) is _custom_add.impls["impl_even"] assert _custom_add.dispatch(x2, y2) is _custom_add.impls["native"] out1 = _custom_add(x1, y1) # size(1) == 2 → impl_even @@ -324,8 +354,12 @@ class TestIrOpImplDispatch: assert torch.all(out2 == 1 + 2) def test_default_priority( - self, caplog_vllm: pytest.LogCaptureFixture, disable_log_dedup + self, + custom_add_op, + caplog_vllm: pytest.LogCaptureFixture, + disable_log_dedup, ): + _custom_add = custom_add_op # Make sure logs are not deduplicated to properly test the warning x = torch.tensor([3], dtype=torch.int32) y = torch.tensor([4], dtype=torch.int32) @@ -347,15 +381,23 @@ class TestIrOpImplDispatch: assert "priority not set" in message -@vllm.ir.register_op -def _custom_mm( - x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None -) -> torch.Tensor: - tmp = x @ y - return tmp if bias is None else tmp + bias +@pytest.fixture +def custom_mm_op(fake_vllm_ir): + """Fixture that registers ``_custom_mm`` (isolated by ``fake_vllm_ir``).""" + + @vllm.ir.register_op + def _custom_mm( + x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None + ) -> torch.Tensor: + tmp = x @ y + return tmp if bias is None else tmp + bias + + return _custom_mm -def test_default_args(): +def test_default_args(custom_mm_op): + _custom_mm = custom_mm_op + # Test that default args are properly applied when dispatching and calling @_custom_mm.register_impl("impl_mm", supports_args=lambda x, y, bias=None: True) def impl_mm( @@ -373,7 +415,8 @@ def test_default_args(): assert _custom_mm.dispatch(x1, x2) is impl_mm -def test_bad_impl_registrations(): +def test_bad_impl_registrations(custom_mm_op): + _custom_mm = custom_mm_op # Check bad schema with pytest.raises(ValueError, match="does not match native schema"): @@ -446,7 +489,8 @@ def test_bad_impl_registrations(): ) -> torch.Tensor: return x @ y + 40 - assert set(_custom_mm.impls.keys()) == {"impl_mm", "native"} + # With fixture, each test gets a fresh op with only "native" impl + assert set(_custom_mm.impls.keys()) == {"native"} IMPL_OOT_SRC = """ @@ -460,14 +504,14 @@ def impl_mm_oot( """ -def load_custom_mm_module(file_path: Path): +def load_custom_mm_module(file_path: Path, custom_mm_op): spec = importlib.util.spec_from_file_location("_custom_mm_oot", file_path) assert spec is not None module = importlib.util.module_from_spec(spec) # Inject the variable into the module's global namespace # This allows the @_custom_mm.register_impl decorator to work - module._custom_mm = _custom_mm # type: ignore[attr-defined] + module._custom_mm = custom_mm_op # type: ignore[attr-defined] # Execute the file; this triggers the decorator assert spec.loader is not None @@ -475,12 +519,13 @@ def load_custom_mm_module(file_path: Path): return module -def test_uuid_and_oot(tmp_path: Path): +def test_uuid_and_oot(custom_mm_op, tmp_path: Path): + _custom_mm = custom_mm_op file_path = tmp_path / "_custom_mm_oot.py" file_path.write_text(IMPL_OOT_SRC) assert "impl_mm_oot" not in _custom_mm.impls - _ = load_custom_mm_module(file_path) + _ = load_custom_mm_module(file_path, _custom_mm) assert "impl_mm_oot" in _custom_mm.impls uuid = _custom_mm.impls["impl_mm_oot"].uuid() @@ -489,7 +534,7 @@ def test_uuid_and_oot(tmp_path: Path): # Replace file source file_path.write_text(IMPL_OOT_SRC + " # added file source") assert "impl_mm_oot" not in _custom_mm.impls - _ = load_custom_mm_module(file_path) + _ = load_custom_mm_module(file_path, _custom_mm) assert "impl_mm_oot" in _custom_mm.impls uuid1 = _custom_mm.impls["impl_mm_oot"].uuid() @@ -499,7 +544,7 @@ def test_uuid_and_oot(tmp_path: Path): # Back to original file_path.write_text(IMPL_OOT_SRC) assert "impl_mm_oot" not in _custom_mm.impls - _ = load_custom_mm_module(file_path) + _ = load_custom_mm_module(file_path, _custom_mm) assert "impl_mm_oot" in _custom_mm.impls uuid2 = _custom_mm.impls["impl_mm_oot"].uuid() @@ -571,3 +616,106 @@ class TestTolerance: op = IrOp("_tol_test_unknown", _test_native) with pytest.raises(ValueError, match="No tolerance defined"): op.get_tolerance(torch.complex64) + + +def test_naming_validation(fake_vllm_ir): + """Test that op and provider names are validated ([a-z_][a-z_0-9]*).""" + + # Valid op and provider names + @vllm.ir.register_op + def _valid_name_123(x: torch.Tensor) -> torch.Tensor: + return x + + @_valid_name_123.register_impl("valid_provider_123") + def valid_impl(x: torch.Tensor) -> torch.Tensor: + return x + 1 + + # Invalid op names should fail + with pytest.raises(ValueError, match="name.*invalid"): + + @vllm.ir.register_op + def InvalidName(x: torch.Tensor) -> torch.Tensor: + return x + + with pytest.raises(ValueError, match="name.*invalid"): + + @vllm.ir.register_op(name="123invalid") + def some_func(x: torch.Tensor) -> torch.Tensor: + return x + + # Invalid provider names should fail + with pytest.raises(ValueError, match="name.*invalid"): + + @_valid_name_123.register_impl("Invalid-Provider") + def invalid_impl(x: torch.Tensor) -> torch.Tensor: + return x + 1 + + +def test_registration_stack_traces(fake_vllm_ir): + """Test that stack traces are captured for ops and impls.""" + + @vllm.ir.register_op + def _test_stack(x: torch.Tensor) -> torch.Tensor: + return x + + @_test_stack.register_impl("test_provider") + def test_impl(x: torch.Tensor) -> torch.Tensor: + return x + 1 + + # Verify op stack trace + assert hasattr(_test_stack, "_registration_stack") + assert len(_test_stack._registration_stack) > 0 + op_stack_str = "".join(_test_stack._registration_stack) + assert "test_op.py" in op_stack_str + # Last frame should be the decorator in user code, not internal decorator logic + assert "@vllm.ir.register_op" in _test_stack._registration_stack[-1] + assert "return decorator(f)" not in op_stack_str + + # Verify impl stack trace + impl = _test_stack.impls["test_provider"] + assert hasattr(impl, "_registration_stack") + assert len(impl._registration_stack) > 0 + impl_stack_str = "".join(impl._registration_stack) + assert "test_op.py" in impl_stack_str + # Last frame should be the decorator in user code + assert '@_test_stack.register_impl("test_provider")' in impl._registration_stack[-1] + + +def test_op_repr_uses_docstring(fake_vllm_ir): + """Test that __str__ uses the function's docstring and __repr__ is simple.""" + + @vllm.ir.register_op + def _test_repr_with_doc(x: torch.Tensor) -> torch.Tensor: + """First line of docstring. + + Additional details here. + """ + return x + + @vllm.ir.register_op + def _test_repr_no_doc(x: torch.Tensor) -> torch.Tensor: + return x + + # __str__ with docstring: uses first line only + str_with = str(_test_repr_with_doc) + assert "IrOp('_test_repr_with_doc')" in str_with + assert "First line of docstring." in str_with + assert "Additional details" not in str_with + + # __str__ without docstring: simple format + assert str(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')" + + # __repr__ should be simple for both + assert repr(_test_repr_with_doc) == "IrOp('_test_repr_with_doc')" + assert repr(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')" + + +def test_vllm_ir_fixture(fake_vllm_ir): + """Test that the fake_vllm_ir fixture provides test isolation.""" + + @vllm.ir.register_op + def _test_fixture(x: torch.Tensor) -> torch.Tensor: + return x + + assert "_test_fixture" in IrOp.registry + # Fixture will automatically clean up after test diff --git a/vllm/ir/op.py b/vllm/ir/op.py index 841df0f9ad..b21c8b994e 100644 --- a/vllm/ir/op.py +++ b/vllm/ir/op.py @@ -2,10 +2,12 @@ # 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 @@ -16,13 +18,35 @@ from vllm.logging_utils import lazy, tensors_str_no_data InputGenerator = Callable[..., tuple[Any, ...]] -vllm_ir_lib = Library("vllm_ir", "FRAGMENT") +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.""" @@ -100,11 +124,14 @@ def register_op( def decorator(_f: Callable): op_name: str = _f.__name__ if name is None else name - assert op_name not in IrOp.registry + _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) + op: IrOp = IrOpInplace(op_name, _f, activations, stack) else: - op = IrOp(op_name, _f, activations) + op = IrOp(op_name, _f, activations, stack) IrOp.registry[op_name] = op return op @@ -126,6 +153,7 @@ class IrOp: 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( @@ -146,6 +174,8 @@ class IrOp: ] 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 = [ @@ -166,22 +196,23 @@ class IrOp: # 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 - vllm_ir_lib.define(self.name + self._schema_str) + # 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 - vllm_ir_lib.impl( - self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd" - ) - vllm_ir_lib._register_fake(self.name, self._fake_call) - assert hasattr(torch.ops.vllm_ir, name) - self.torch_op: torch._ops.OpOverload = getattr(torch.ops.vllm_ir, name).default + 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: """ @@ -237,9 +268,12 @@ class IrOp: assert provider not in RESERVED_PROVIDERS, ( f"Provider name {provider} is reserved." ) + _validate_name(provider, "Provider") def _register_impl(f: Callable): - impl = IrOpImpl(self, provider, f, supported, supports_args, inplace) + # 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(): @@ -323,6 +357,17 @@ class IrOp: 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] @@ -419,8 +464,9 @@ class IrOpInplace(IrOp): name: str, native_impl: Callable, activations: list[str] | None = None, + registration_stack: list[str] | None = None, ): - super().__init__(name, native_impl, activations) + super().__init__(name, native_impl, activations, registration_stack) # Create the inplace overload self.maybe_inplace = IrOpInplaceOverload(self) @@ -445,16 +491,16 @@ class IrOpInplaceOverload: op.impls["native"].impl_fn, mutates_args=op.activations ) - # torch registration - vllm_ir_lib.define(self.name + self._schema_str) - vllm_ir_lib.impl( - self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd" - ) + # 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 - vllm_ir_lib._register_fake(self.name, self.op._fake_call) + lib._register_fake(self.name, self.op._fake_call) - assert hasattr(getattr(torch.ops.vllm_ir, self.op.name), "maybe_inplace") - self.torch_op = getattr(torch.ops.vllm_ir, self.op.name).maybe_inplace + 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: @@ -477,12 +523,15 @@ class IrOpImpl: 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 + 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. @@ -547,6 +596,7 @@ class IrOpImpl: self.supported = supported self._supports_args = supports_args self.inplace = inplace + self._registration_stack = registration_stack or [] @property def supports_all_args(self) -> bool: