@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 289cee0473

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 20:15 UTC

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

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

Avishek Goswami committed May 12, 2026 at 20:28 UTC 289cee047390ffec8c61600e887232fc8b339e2c

3 files changed +328 -89

tests/conftest.py

+41

@@ -1686,3 +1686,44 @@ def disable_log_dedup(monkeypatch):

1686 logger._print_warning_once = original_print_warning_once

1687 logger._print_info_once = original_print_info_once

1688 logger._print_debug_once = original_print_debug_once

1689 +

1690 +

1691 +@pytest.fixture(scope="function")

1692 +def fake_vllm_ir(monkeypatch):

1693 + """

1694 + Pytest fixture to allow isolated IR op registration in tests.

1695 +

1696 + Replaces IrOp.registry with an empty dict and swaps ``vllm_ir_torch_lib`` for a

1697 + fresh ``Library`` with a unique namespace per test (see ``Library.ns``).

1698 +

1699 + Torch keeps registrations for the process lifetime; reusing the fragment

1700 + name ``vllm_ir`` and defining the same op string again can segfault. A

1701 + random library name keeps each fixture run on a disjoint namespace.

1702 +

1703 + The test Library is kept alive until after monkeypatch teardown so PyTorch's

1704 + C++ state is not freed while references may still exist.

1705 +

1706 + Usage:

1707 + def test_my_ir_op(fake_vllm_ir):

1708 + @vllm.ir.register_op

1709 + def my_test_op(x: torch.Tensor) -> torch.Tensor:

1710 + return x * 2

1711 +

1712 + result = my_test_op(torch.tensor([1, 2, 3]))

1713 + # Registry and library cleaned up automatically after the test

1714 + """

1715 + import secrets

1716 +

1717 + from torch.library import Library

1718 + from vllm.ir.op import IrOp

1719 +

1720 + monkeypatch.setattr(IrOp, "registry", {})

1721 +

1722 + # Keep a local reference so the Library is not GC'd before monkeypatch

1723 + # teardown restores the original reference.

1724 + test_lib = Library(f"vllm_ir_{secrets.token_hex(8)}", "FRAGMENT")

1725 + monkeypatch.setattr("vllm.ir.op.vllm_ir_torch_lib", test_lib)

1726 +

1727 + yield

1728 +

1729 + del test_lib

tests/ir/test_op.py

+214 -66

@@ -13,20 +13,40 @@ from torch.fx.experimental.proxy_tensor import make_fx

13 import vllm.ir.op

14 from vllm.ir.op import RESERVED_PROVIDERS, IrOp, IrOpImpl

16 -# This should not exist

17 -assert "_custom_add" not in IrOp.registry

18 -

17 class CustomError(Exception):

18 pass

24 -@vllm.ir.register_op(allow_inplace=True)

25 -def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

26 - return x + y

21 +@pytest.fixture

22 +def custom_add_op(fake_vllm_ir):

23 + """Register ``_custom_add`` plus impl_a, impl_b, impl_even for this test."""

24 +

25 + @vllm.ir.register_op(allow_inplace=True)

26 + def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

27 + return x + y

28 +

29 + @_custom_add.register_impl("impl_a")

30 + def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

31 + return x + y + 10

32 +

33 + @_custom_add.register_impl("impl_b", inplace=True)

34 + def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

35 + """Computes x+y+20"""

36 + x.add_(y)

37 + x.add_(20)

38 + return x

39 +

40 + @_custom_add.register_impl(

41 + "impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0

42 + )

43 + def impl_even(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

44 + return x + y + 50

46 + return _custom_add

29 -def test_registration_overloads():

48 +

49 +def test_registration_overloads(fake_vllm_ir):

50 assert all(

51 n not in IrOp.registry for n in ["_custom_sub", "_custom_mul", "_custom_div"]

52 )

@@ -63,7 +83,7 @@ def test_registration_overloads():

83 return x * y - 100

66 -def test_no_kw_only_args():

86 +def test_no_kw_only_args(fake_vllm_ir):

87 # kw-only args not supported

88 with pytest.raises(ValueError, match="keyword-only arguments"):

@@ -78,18 +98,23 @@ def test_no_kw_only_args():

99 class TestIrOpCustomAdd:

100 # Registration invariants

81 - def test_decorated_object(self):

101 + def test_decorated_object(self, custom_add_op):

102 """Make sure that referring directly to an op is correct"""

103 + _custom_add = custom_add_op

104 assert isinstance(_custom_add, IrOp)

105 assert "_custom_add" in IrOp.registry

106 assert _custom_add is IrOp.registry["_custom_add"]

107

87 - def test_torch_op_is_registered(self):

88 - assert hasattr(torch.ops.vllm_ir, "_custom_add")

89 - assert callable(torch.ops.vllm_ir._custom_add.default)

108 + def test_torch_op_is_registered(self, custom_add_op):

109 + _custom_add = custom_add_op

110 + torch_ops = getattr(torch.ops, vllm.ir.op.vllm_ir_torch_lib.ns)

111 + assert hasattr(torch_ops, "_custom_add")

112 + assert callable(torch_ops._custom_add.default)

113 + assert _custom_add.torch_op is torch_ops._custom_add.default

114

115 # Semantic correctness

92 - def test_semantics_match_native(self):

116 + def test_semantics_match_native(self, custom_add_op):

117 + _custom_add = custom_add_op

118 x = torch.randn(4, 5)

119 y = torch.randn(4, 5)

120

@@ -103,7 +128,9 @@ class TestIrOpCustomAdd:

128 # Implementation registration

129 # -------------------------

130

106 - def test_register_impl_is_non_intrusive(self):

131 + def test_register_impl_is_non_intrusive(self, custom_add_op):

132 + _custom_add = custom_add_op

133 +

134 @_custom_add.register_impl("dummy_provider")

135 def dummy_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

136 return x + y + 123

@@ -117,7 +144,8 @@ class TestIrOpCustomAdd:

144 # Native semantics must still hold

145 torch.testing.assert_close(_custom_add(x, y), x + y)

146

120 - def test_schema_contains_tensor_signature(self):

147 + def test_schema_contains_tensor_signature(self, custom_add_op):

148 + _custom_add = custom_add_op

149 schema = _custom_add._schema_str

150

151 assert "Tensor" in schema

@@ -131,10 +159,19 @@ class TestIrOpCustomAdd:

159 @pytest.mark.parametrize("symbolic_trace", [True, False])

160 @pytest.mark.parametrize("overload", ["default", "maybe_inplace"])

161 def test_trace_sees_single_custom_op(

134 - self, symbolic_trace: bool, enable_torch_wrap: bool, overload: str

162 + self,

163 + custom_add_op,

164 + symbolic_trace: bool,

165 + enable_torch_wrap: bool,

166 + overload: str,

167 ):

168 + _custom_add = custom_add_op

169 op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace

137 - torch_op = getattr(torch.ops.vllm_ir._custom_add, overload)

170 + torch_op = (

171 + _custom_add.torch_op

172 + if overload == "default"

173 + else _custom_add.maybe_inplace.torch_op

174 + )

175

176 def fn(x, y):

177 return op_fn(x, y)

@@ -175,30 +212,13 @@ class TestIrOpCustomAdd:

212 assert len(ir_nodes) == 1, gm.code

213

214

178 -@_custom_add.register_impl("impl_a")

179 -def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

180 - return x + y + 10

181 -

182 -

183 -@_custom_add.register_impl("impl_b", inplace=True)

184 -def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

185 - """Computes x+y+20"""

186 - x.add_(y)

187 - x.add_(20)

188 - return x

189 -

190 -

191 -@_custom_add.register_impl("impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0)

192 -def impl_even(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

193 - return x + y + 50

194 -

195 -

215 class TestIrOpImplDispatch:

197 - def test_register_impl(self):

216 + def test_register_impl(self, custom_add_op):

217 + _custom_add = custom_add_op

218 assert "impl_a" in _custom_add.impls

219 impl = _custom_add.impls["impl_a"]

220

201 - assert impl is impl_a

221 + assert impl is _custom_add.impls["impl_a"]

222 assert impl.op is _custom_add

223 assert impl.provider == "impl_a"

224 assert callable(impl.impl_fn)

@@ -211,14 +231,15 @@ class TestIrOpImplDispatch:

231 return x + y + 30

232

233 # Check the original impl is still intact

214 - assert _custom_add.impls["impl_a"] is impl_a

234 + assert _custom_add.impls["impl_a"] is impl

235

236 # Check support all args

217 - assert impl_a.supports_all_args

218 - assert impl_b.supports_all_args

219 - assert not impl_even.supports_all_args

237 + assert _custom_add.impls["impl_a"].supports_all_args

238 + assert _custom_add.impls["impl_b"].supports_all_args

239 + assert not _custom_add.impls["impl_even"].supports_all_args

240

221 - def test_reserved_provider_rejected(self):

241 + def test_reserved_provider_rejected(self, custom_add_op):

242 + _custom_add = custom_add_op

243 for provider in RESERVED_PROVIDERS:

244 with pytest.raises(AssertionError):

245

@@ -226,7 +247,8 @@ class TestIrOpImplDispatch:

247 def bad_impl(x, y):

248 return x + y

249

229 - def test_set_priority_scoped(self):

250 + def test_set_priority_scoped(self, custom_add_op):

251 + _custom_add = custom_add_op

252 assert _custom_add.get_priority() == []

253

254 with _custom_add.set_priority(["impl_even", "impl_b"]):

@@ -251,20 +273,25 @@ class TestIrOpImplDispatch:

273 assert _custom_add.get_priority() == []

274

275 @pytest.mark.parametrize("overload", ["default", "maybe_inplace"])

254 - def test_dispatch_priority_order(self, overload: str):

276 + def test_dispatch_priority_order(self, custom_add_op, overload: str):

277 + _custom_add = custom_add_op

278 op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace

256 - torch_op = getattr(torch.ops.vllm_ir._custom_add, overload)

279 + torch_op = (

280 + _custom_add.torch_op

281 + if overload == "default"

282 + else _custom_add.maybe_inplace.torch_op

283 + )

284

285 x = torch.tensor(1, dtype=torch.int32)

286 y = torch.tensor(2, dtype=torch.int32)

287

288 with _custom_add.set_priority(["impl_b", "impl_a"]):

262 - assert _custom_add.dispatch(x, y) is impl_b

289 + assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_b"]

290 out1 = op_fn(x.clone(), y)

291 out2 = torch_op(x.clone(), y)

292

293 with _custom_add.set_priority(["impl_a"]):

267 - assert _custom_add.dispatch(x, y) is impl_a

294 + assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_a"]

295 out3 = op_fn(x.clone(), y)

296 out4 = torch_op(x.clone(), y)

297

@@ -275,7 +302,9 @@ class TestIrOpImplDispatch:

302 assert out3.item() == 1 + 2 + 10

303 assert out4.item() == 1 + 2 + 10

304

278 - def test_unsupported_impl_filtered(self):

305 + def test_unsupported_impl_filtered(self, custom_add_op):

306 + _custom_add = custom_add_op

307 +

308 @_custom_add.register_impl("impl_unsupported", supported=False)

309 def impl_unsupported(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:

310 return x + y + 999

@@ -291,8 +320,9 @@ class TestIrOpImplDispatch:

320 assert out.item() == 1 + 2 + 10

321

322 def test_supports_args_runtime_dispatch_and_warning(

294 - self, caplog_vllm: pytest.LogCaptureFixture

323 + self, custom_add_op, caplog_vllm: pytest.LogCaptureFixture

324 ):

325 + _custom_add = custom_add_op

326 x1 = torch.ones((2, 2), dtype=torch.int32)

327 y1 = torch.full((2, 2), 2, dtype=torch.int32)

328

@@ -312,7 +342,7 @@ class TestIrOpImplDispatch:

342

343 # Check dispatching

344 assert _custom_add.get_priority() == ["impl_even", "native"]

315 - assert _custom_add.dispatch(x1, y1) is impl_even

345 + assert _custom_add.dispatch(x1, y1) is _custom_add.impls["impl_even"]

346 assert _custom_add.dispatch(x2, y2) is _custom_add.impls["native"]

347

348 out1 = _custom_add(x1, y1) # size(1) == 2 → impl_even

@@ -324,8 +354,12 @@ class TestIrOpImplDispatch:

354 assert torch.all(out2 == 1 + 2)

355

356 def test_default_priority(

327 - self, caplog_vllm: pytest.LogCaptureFixture, disable_log_dedup

357 + self,

358 + custom_add_op,

359 + caplog_vllm: pytest.LogCaptureFixture,

360 + disable_log_dedup,

361 ):

362 + _custom_add = custom_add_op

363 # Make sure logs are not deduplicated to properly test the warning

364 x = torch.tensor([3], dtype=torch.int32)

365 y = torch.tensor([4], dtype=torch.int32)

@@ -347,15 +381,23 @@ class TestIrOpImplDispatch:

381 assert "priority not set" in message

382

383

350 -@vllm.ir.register_op

351 -def _custom_mm(

352 - x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None

353 -) -> torch.Tensor:

354 - tmp = x @ y

355 - return tmp if bias is None else tmp + bias

384 +@pytest.fixture

385 +def custom_mm_op(fake_vllm_ir):

386 + """Fixture that registers ``_custom_mm`` (isolated by ``fake_vllm_ir``)."""

387

388 + @vllm.ir.register_op

389 + def _custom_mm(

390 + x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None

391 + ) -> torch.Tensor:

392 + tmp = x @ y

393 + return tmp if bias is None else tmp + bias

394 +

395 + return _custom_mm

396 +

397 +

398 +def test_default_args(custom_mm_op):

399 + _custom_mm = custom_mm_op

400

358 -def test_default_args():

401 # Test that default args are properly applied when dispatching and calling

402 @_custom_mm.register_impl("impl_mm", supports_args=lambda x, y, bias=None: True)

403 def impl_mm(

@@ -373,7 +415,8 @@ def test_default_args():

415 assert _custom_mm.dispatch(x1, x2) is impl_mm

416

417

376 -def test_bad_impl_registrations():

418 +def test_bad_impl_registrations(custom_mm_op):

419 + _custom_mm = custom_mm_op

420 # Check bad schema

421 with pytest.raises(ValueError, match="does not match native schema"):

422

@@ -446,7 +489,8 @@ def test_bad_impl_registrations():

489 ) -> torch.Tensor:

490 return x @ y + 40

491

449 - assert set(_custom_mm.impls.keys()) == {"impl_mm", "native"}

492 + # With fixture, each test gets a fresh op with only "native" impl

493 + assert set(_custom_mm.impls.keys()) == {"native"}

494

495

496 IMPL_OOT_SRC = """

@@ -460,14 +504,14 @@ def impl_mm_oot(

504 """

505

506

463 -def load_custom_mm_module(file_path: Path):

507 +def load_custom_mm_module(file_path: Path, custom_mm_op):

508 spec = importlib.util.spec_from_file_location("_custom_mm_oot", file_path)

509 assert spec is not None

510 module = importlib.util.module_from_spec(spec)

511

512 # Inject the variable into the module's global namespace

513 # This allows the @_custom_mm.register_impl decorator to work

470 - module._custom_mm = _custom_mm # type: ignore[attr-defined]

514 + module._custom_mm = custom_mm_op # type: ignore[attr-defined]

515

516 # Execute the file; this triggers the decorator

517 assert spec.loader is not None

@@ -475,12 +519,13 @@ def load_custom_mm_module(file_path: Path):

519 return module

520

521

478 -def test_uuid_and_oot(tmp_path: Path):

522 +def test_uuid_and_oot(custom_mm_op, tmp_path: Path):

523 + _custom_mm = custom_mm_op

524 file_path = tmp_path / "_custom_mm_oot.py"

525 file_path.write_text(IMPL_OOT_SRC)

526

527 assert "impl_mm_oot" not in _custom_mm.impls

483 - _ = load_custom_mm_module(file_path)

528 + _ = load_custom_mm_module(file_path, _custom_mm)

529 assert "impl_mm_oot" in _custom_mm.impls

530

531 uuid = _custom_mm.impls["impl_mm_oot"].uuid()

@@ -489,7 +534,7 @@ def test_uuid_and_oot(tmp_path: Path):

534 # Replace file source

535 file_path.write_text(IMPL_OOT_SRC + " # added file source")

536 assert "impl_mm_oot" not in _custom_mm.impls

492 - _ = load_custom_mm_module(file_path)

537 + _ = load_custom_mm_module(file_path, _custom_mm)

538 assert "impl_mm_oot" in _custom_mm.impls

539

540 uuid1 = _custom_mm.impls["impl_mm_oot"].uuid()

@@ -499,7 +544,7 @@ def test_uuid_and_oot(tmp_path: Path):

544 # Back to original

545 file_path.write_text(IMPL_OOT_SRC)

546 assert "impl_mm_oot" not in _custom_mm.impls

502 - _ = load_custom_mm_module(file_path)

547 + _ = load_custom_mm_module(file_path, _custom_mm)

548 assert "impl_mm_oot" in _custom_mm.impls

549

550 uuid2 = _custom_mm.impls["impl_mm_oot"].uuid()

@@ -571,3 +616,106 @@ class TestTolerance:

616 op = IrOp("_tol_test_unknown", _test_native)

617 with pytest.raises(ValueError, match="No tolerance defined"):

618 op.get_tolerance(torch.complex64)

619 +

620 +

621 +def test_naming_validation(fake_vllm_ir):

622 + """Test that op and provider names are validated ([a-z_][a-z_0-9]*)."""

623 +

624 + # Valid op and provider names

625 + @vllm.ir.register_op

626 + def _valid_name_123(x: torch.Tensor) -> torch.Tensor:

627 + return x

628 +

629 + @_valid_name_123.register_impl("valid_provider_123")

630 + def valid_impl(x: torch.Tensor) -> torch.Tensor:

631 + return x + 1

632 +

633 + # Invalid op names should fail

634 + with pytest.raises(ValueError, match="name.*invalid"):

635 +

636 + @vllm.ir.register_op

637 + def InvalidName(x: torch.Tensor) -> torch.Tensor:

638 + return x

639 +

640 + with pytest.raises(ValueError, match="name.*invalid"):

641 +

642 + @vllm.ir.register_op(name="123invalid")

643 + def some_func(x: torch.Tensor) -> torch.Tensor:

644 + return x

645 +

646 + # Invalid provider names should fail

647 + with pytest.raises(ValueError, match="name.*invalid"):

648 +

649 + @_valid_name_123.register_impl("Invalid-Provider")

650 + def invalid_impl(x: torch.Tensor) -> torch.Tensor:

651 + return x + 1

652 +

653 +

654 +def test_registration_stack_traces(fake_vllm_ir):

655 + """Test that stack traces are captured for ops and impls."""

656 +

657 + @vllm.ir.register_op

658 + def _test_stack(x: torch.Tensor) -> torch.Tensor:

659 + return x

660 +

661 + @_test_stack.register_impl("test_provider")

662 + def test_impl(x: torch.Tensor) -> torch.Tensor:

663 + return x + 1

664 +

665 + # Verify op stack trace

666 + assert hasattr(_test_stack, "_registration_stack")

667 + assert len(_test_stack._registration_stack) > 0

668 + op_stack_str = "".join(_test_stack._registration_stack)

669 + assert "test_op.py" in op_stack_str

670 + # Last frame should be the decorator in user code, not internal decorator logic

671 + assert "@vllm.ir.register_op" in _test_stack._registration_stack[-1]

672 + assert "return decorator(f)" not in op_stack_str

673 +

674 + # Verify impl stack trace

675 + impl = _test_stack.impls["test_provider"]

676 + assert hasattr(impl, "_registration_stack")

677 + assert len(impl._registration_stack) > 0

678 + impl_stack_str = "".join(impl._registration_stack)

679 + assert "test_op.py" in impl_stack_str

680 + # Last frame should be the decorator in user code

681 + assert '@_test_stack.register_impl("test_provider")' in impl._registration_stack[-1]

682 +

683 +

684 +def test_op_repr_uses_docstring(fake_vllm_ir):

685 + """Test that __str__ uses the function's docstring and __repr__ is simple."""

686 +

687 + @vllm.ir.register_op

688 + def _test_repr_with_doc(x: torch.Tensor) -> torch.Tensor:

689 + """First line of docstring.

690 +

691 + Additional details here.

692 + """

693 + return x

694 +

695 + @vllm.ir.register_op

696 + def _test_repr_no_doc(x: torch.Tensor) -> torch.Tensor:

697 + return x

698 +

699 + # __str__ with docstring: uses first line only

700 + str_with = str(_test_repr_with_doc)

701 + assert "IrOp('_test_repr_with_doc')" in str_with

702 + assert "First line of docstring." in str_with

703 + assert "Additional details" not in str_with

704 +

705 + # __str__ without docstring: simple format

706 + assert str(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')"

707 +

708 + # __repr__ should be simple for both

709 + assert repr(_test_repr_with_doc) == "IrOp('_test_repr_with_doc')"

710 + assert repr(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')"

711 +

712 +

713 +def test_vllm_ir_fixture(fake_vllm_ir):

714 + """Test that the fake_vllm_ir fixture provides test isolation."""

715 +

716 + @vllm.ir.register_op

717 + def _test_fixture(x: torch.Tensor) -> torch.Tensor:

718 + return x

719 +

720 + assert "_test_fixture" in IrOp.registry

721 + # Fixture will automatically clean up after test

vllm/ir/op.py

+73 -23

@@ -2,10 +2,12 @@

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

10 +import regex as re

11 import torch

12 from torch.library import Library, infer_schema

@@ -16,13 +18,35 @@ from vllm.logging_utils import lazy, tensors_str_no_data

19 InputGenerator = Callable[..., tuple[Any, ...]]

19 -vllm_ir_lib = Library("vllm_ir", "FRAGMENT")

21 +vllm_ir_torch_lib = Library("vllm_ir", "FRAGMENT") # IR op lib; monkeypatch in tests.

23 logger = init_logger(__name__)

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."""

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."""

@@ -100,11 +124,14 @@ def register_op(

124

125 def decorator(_f: Callable):

126 op_name: str = _f.__name__ if name is None else name

103 - assert op_name not in IrOp.registry

127 + _validate_name(op_name, "Op")

128 + assert op_name not in IrOp.registry, f"Op '{op_name}' is already registered."

129 + # Slice out the decorator function frames from the stack

130 + stack = traceback.format_stack()[:-2]

131 if allow_inplace:

105 - op: IrOp = IrOpInplace(op_name, _f, activations)

132 + op: IrOp = IrOpInplace(op_name, _f, activations, stack)

133 else:

107 - op = IrOp(op_name, _f, activations)

134 + op = IrOp(op_name, _f, activations, stack)

135 IrOp.registry[op_name] = op

136 return op

137

@@ -126,6 +153,7 @@ class IrOp:

153 name: str,

154 native_impl: Callable,

155 activations: list[str] | None = None,

156 + registration_stack: list[str] | None = None,

157 ):

158 self._py_signature = inspect.signature(native_impl)

159 if any(

@@ -146,6 +174,8 @@ class IrOp:

174 ]

175

176 self.name = name

177 + self._docstring = inspect.getdoc(native_impl) or ""

178 + self._registration_stack = registration_stack or []

179 self.impls: dict[str, IrOpImpl] = {}

180 self.activations = activations

181 self.activation_indices = [

@@ -166,22 +196,23 @@ class IrOp:

196 # always supported

197 supported=True,

198 supports_args=None,

199 + registration_stack=self._registration_stack,

200 )

201

202 # By default, fake routes directly to native,

203 # can be overridden by register_fake

204 self._fake_fn = native_impl

205

175 - # torch registration

176 - vllm_ir_lib.define(self.name + self._schema_str)

206 + # torch registration (resolve ``torch.ops`` subtree from ``lib.ns``)

207 + lib = vllm_ir_torch_lib

208 + lib.define(self.name + self._schema_str)

209 # CompositeExplicitAutograd is not decomposed

210 # by ATen IR normalization in AOTAutograd

179 - vllm_ir_lib.impl(

180 - self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd"

181 - )

182 - vllm_ir_lib._register_fake(self.name, self._fake_call)

183 - assert hasattr(torch.ops.vllm_ir, name)

184 - self.torch_op: torch._ops.OpOverload = getattr(torch.ops.vllm_ir, name).default

211 + lib.impl(self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd")

212 + lib._register_fake(self.name, self._fake_call)

213 + torch_ops = _torch_ops_subtree(lib)

214 + assert hasattr(torch_ops, name)

215 + self.torch_op: torch._ops.OpOverload = getattr(torch_ops, name).default

216

217 def register_fake(self, fn: Callable) -> Callable:

218 """

@@ -237,9 +268,12 @@ class IrOp:

268 assert provider not in RESERVED_PROVIDERS, (

269 f"Provider name {provider} is reserved."

270 )

271 + _validate_name(provider, "Provider")

272

273 def _register_impl(f: Callable):

242 - impl = IrOpImpl(self, provider, f, supported, supports_args, inplace)

274 + # Slice out the decorator function from the stack

275 + stack = traceback.format_stack()[:-1]

276 + impl = IrOpImpl(self, provider, f, supported, supports_args, inplace, stack)

277 self.impls[provider] = impl

278

279 if self.get_priority():

@@ -323,6 +357,17 @@ class IrOp:

357

358 return self.torch_op(*args, **kwargs)

359

360 + def __repr__(self) -> str:

361 + """Return unambiguous string representation."""

362 + return f"IrOp('{self.name}')"

363 +

364 + def __str__(self) -> str:

365 + """Return human-readable string representation using docstring."""

366 + if not self._docstring:

367 + return f"IrOp('{self.name}')"

368 + first_line = self._docstring.split("\n")[0].strip()

369 + return f"IrOp('{self.name}') - {first_line}"

370 +

371 def get_priority(self) -> list[str]:

372 """Get the current dispatch priority for implementations for this op."""

373 return [p.provider for p in self._priority_impls]

@@ -419,8 +464,9 @@ class IrOpInplace(IrOp):

464 name: str,

465 native_impl: Callable,

466 activations: list[str] | None = None,

467 + registration_stack: list[str] | None = None,

468 ):

423 - super().__init__(name, native_impl, activations)

469 + super().__init__(name, native_impl, activations, registration_stack)

470

471 # Create the inplace overload

472 self.maybe_inplace = IrOpInplaceOverload(self)

@@ -445,16 +491,16 @@ class IrOpInplaceOverload:

491 op.impls["native"].impl_fn, mutates_args=op.activations

492 )

493

448 - # torch registration

449 - vllm_ir_lib.define(self.name + self._schema_str)

450 - vllm_ir_lib.impl(

451 - self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd"

452 - )

494 + # torch registration (resolve ``torch.ops`` subtree from ``lib.ns``)

495 + lib = vllm_ir_torch_lib

496 + lib.define(self.name + self._schema_str)

497 + lib.impl(self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd")

498 # fake goes to default overload for now

454 - vllm_ir_lib._register_fake(self.name, self.op._fake_call)

499 + lib._register_fake(self.name, self.op._fake_call)

500

456 - assert hasattr(getattr(torch.ops.vllm_ir, self.op.name), "maybe_inplace")

457 - self.torch_op = getattr(torch.ops.vllm_ir, self.op.name).maybe_inplace

501 + torch_ops = _torch_ops_subtree(lib)

502 + assert hasattr(getattr(torch_ops, self.op.name), "maybe_inplace")

503 + self.torch_op = getattr(torch_ops, self.op.name).maybe_inplace

504

505 def __call__(self, *args, **kwargs) -> Any:

506 if not _ENABLE_TORCH_WRAP:

@@ -477,12 +523,15 @@ class IrOpImpl:

523 supported: bool,

524 supports_args: Callable[..., bool] | None,

525 inplace: bool = False,

526 + registration_stack: list[str] | None = None,

527 ):

528 assert provider not in op.impls, (

529 f"Implementation for provider {provider} already registered."

530 )

531 # Native also uses this path, so we allow it here.

485 - assert provider == "native" or provider not in RESERVED_PROVIDERS

532 + assert provider == "native" or provider not in RESERVED_PROVIDERS, (

533 + f"Provider name {provider} is reserved."

534 + )

535

536 # Enforce the exact same schema as the native implementation.

537 # This takes care of names, types, and defaults.

@@ -547,6 +596,7 @@ class IrOpImpl:

596 self.supported = supported

597 self._supports_args = supports_args

598 self.inplace = inplace

599 + self._registration_stack = registration_stack or []

600

601 @property

602 def supports_all_args(self) -> bool: