[vLLM IR] Minor improvements (#39362) (#39558)

Signed-off-by: Avishek Goswami <avishek.goswami@ibm.com>
Co-authored-by: Avishek Goswami <avishek.goswami@ibm.com>

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

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

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: