| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import importlib.util |
| 4 | import logging |
| 5 | from pathlib import Path |
| 6 | from typing import Any |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | from torch import fx |
| 11 | from torch.fx.experimental.proxy_tensor import make_fx |
| 12 | |
| 13 | import vllm.ir.op |
| 14 | from vllm.ir.op import RESERVED_PROVIDERS, IrOp, IrOpImpl |
| 15 | |
| 16 | pytestmark = pytest.mark.skip_global_cleanup |
| 17 | |
| 18 | |
| 19 | class CustomError(Exception): |
| 20 | pass |
| 21 | |
| 22 | |
| 23 | @pytest.fixture |
| 24 | def custom_add_op(fake_vllm_ir): |
| 25 | """Register ``_custom_add`` plus impl_a, impl_b, impl_even for this test.""" |
| 26 | |
| 27 | @vllm.ir.register_op(allow_inplace=True) |
| 28 | def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 29 | return x + y |
| 30 | |
| 31 | @_custom_add.register_impl("impl_a") |
| 32 | def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 33 | return x + y + 10 |
| 34 | |
| 35 | @_custom_add.register_impl("impl_b", inplace=True) |
| 36 | def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 37 | """Computes x+y+20.""" |
| 38 | x.add_(y) |
| 39 | x.add_(20) |
| 40 | return x |
| 41 | |
| 42 | @_custom_add.register_impl( |
| 43 | "impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0 |
| 44 | ) |
| 45 | def impl_even(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 46 | return x + y + 50 |
| 47 | |
| 48 | return _custom_add |
| 49 | |
| 50 | |
| 51 | def test_registration_overloads(fake_vllm_ir): |
| 52 | assert all( |
| 53 | n not in IrOp.registry for n in ["_custom_sub", "_custom_mul", "_custom_div"] |
| 54 | ) |
| 55 | |
| 56 | # Calling with decorator |
| 57 | @vllm.ir.register_op() |
| 58 | def _custom_sub(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 59 | return x - y |
| 60 | |
| 61 | assert _custom_sub.name == "_custom_sub" |
| 62 | assert _custom_sub is IrOp.registry["_custom_sub"] |
| 63 | |
| 64 | # Custom name |
| 65 | @vllm.ir.register_op(name="_custom_mul") |
| 66 | def custom_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 67 | return x * y |
| 68 | |
| 69 | assert custom_mul.name == "_custom_mul" |
| 70 | assert custom_mul is IrOp.registry["_custom_mul"] |
| 71 | |
| 72 | # Direct construction does not register directly |
| 73 | def _custom_div(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 74 | return x / y |
| 75 | |
| 76 | custom_div = IrOp("_custom_div", _custom_div) |
| 77 | assert custom_div.name == "_custom_div" |
| 78 | assert "_custom_div" not in IrOp.registry |
| 79 | |
| 80 | # Duplicate op registration not allowed |
| 81 | with pytest.raises(AssertionError): |
| 82 | |
| 83 | @vllm.ir.register_op |
| 84 | def _custom_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 85 | return x * y - 100 |
| 86 | |
| 87 | |
| 88 | def test_no_kw_only_args(fake_vllm_ir): |
| 89 | # kw-only args not supported |
| 90 | with pytest.raises(ValueError, match="keyword-only arguments"): |
| 91 | |
| 92 | @vllm.ir.register_op |
| 93 | def _custom_kwarg_op( |
| 94 | x: torch.Tensor, y: torch.Tensor, *, kwarg: int = 0 |
| 95 | ) -> torch.Tensor: |
| 96 | return x + y + kwarg |
| 97 | |
| 98 | assert "_custom_kwarg_op" not in IrOp.registry |
| 99 | |
| 100 | |
| 101 | class TestIrOpCustomAdd: |
| 102 | # Registration invariants |
| 103 | def test_decorated_object(self, custom_add_op): |
| 104 | """Make sure that referring directly to an op is correct.""" |
| 105 | _custom_add = custom_add_op |
| 106 | assert isinstance(_custom_add, IrOp) |
| 107 | assert "_custom_add" in IrOp.registry |
| 108 | assert _custom_add is IrOp.registry["_custom_add"] |
| 109 | |
| 110 | def test_torch_op_is_registered(self, custom_add_op): |
| 111 | _custom_add = custom_add_op |
| 112 | torch_ops = getattr(torch.ops, vllm.ir.op.vllm_ir_torch_lib.ns) |
| 113 | assert hasattr(torch_ops, "_custom_add") |
| 114 | assert callable(torch_ops._custom_add.default) |
| 115 | assert _custom_add.torch_op is torch_ops._custom_add.default |
| 116 | |
| 117 | # Semantic correctness |
| 118 | def test_semantics_match_native(self, custom_add_op): |
| 119 | _custom_add = custom_add_op |
| 120 | x = torch.randn(4, 5) |
| 121 | y = torch.randn(4, 5) |
| 122 | |
| 123 | # Calls native by default |
| 124 | out = _custom_add(x, y) |
| 125 | ref = x + y |
| 126 | |
| 127 | torch.testing.assert_close(out, ref) |
| 128 | |
| 129 | # ------------------------- |
| 130 | # Implementation registration |
| 131 | # ------------------------- |
| 132 | |
| 133 | def test_register_impl_is_non_intrusive(self, custom_add_op): |
| 134 | _custom_add = custom_add_op |
| 135 | |
| 136 | @_custom_add.register_impl("dummy_provider") |
| 137 | def dummy_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 138 | return x + y + 123 |
| 139 | |
| 140 | assert "dummy_provider" in _custom_add.impls |
| 141 | assert isinstance(_custom_add.impls["dummy_provider"], IrOpImpl) |
| 142 | |
| 143 | x = torch.ones(2, 2) |
| 144 | y = torch.ones(2, 2) |
| 145 | |
| 146 | # Native semantics must still hold |
| 147 | torch.testing.assert_close(_custom_add(x, y), x + y) |
| 148 | |
| 149 | def test_schema_contains_tensor_signature(self, custom_add_op): |
| 150 | _custom_add = custom_add_op |
| 151 | schema = _custom_add._schema_str |
| 152 | |
| 153 | assert "Tensor" in schema |
| 154 | assert "-> Tensor" in schema |
| 155 | |
| 156 | # ------------------------- |
| 157 | # FX visibility |
| 158 | # ------------------------- |
| 159 | |
| 160 | @pytest.mark.parametrize("enable_torch_wrap", [True, False]) |
| 161 | @pytest.mark.parametrize("symbolic_trace", [True, False]) |
| 162 | @pytest.mark.parametrize("overload", ["default", "maybe_inplace"]) |
| 163 | def test_trace_sees_single_custom_op( |
| 164 | self, |
| 165 | custom_add_op, |
| 166 | symbolic_trace: bool, |
| 167 | enable_torch_wrap: bool, |
| 168 | overload: str, |
| 169 | ): |
| 170 | _custom_add = custom_add_op |
| 171 | op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace |
| 172 | torch_op = ( |
| 173 | _custom_add.torch_op |
| 174 | if overload == "default" |
| 175 | else _custom_add.maybe_inplace.torch_op |
| 176 | ) |
| 177 | |
| 178 | def fn(x, y): |
| 179 | return op_fn(x, y) |
| 180 | |
| 181 | def find_fn(target: Any, gm: fx.GraphModule): |
| 182 | return gm.graph.find_nodes(op="call_function", target=target) |
| 183 | |
| 184 | with pytest.raises(CustomError), vllm.ir.enable_torch_wrap(enable_torch_wrap): |
| 185 | if symbolic_trace: |
| 186 | gm = torch.fx.symbolic_trace(fn) |
| 187 | else: |
| 188 | gm = make_fx(fn)(torch.randn(2, 2), torch.randn(2, 2)) |
| 189 | |
| 190 | x1, y1 = torch.rand(5, 4), torch.rand(5, 4) |
| 191 | out_fx = gm(x1, y1) |
| 192 | out_eager = fn(x1, y1) |
| 193 | |
| 194 | # raise error to check enable_torch_wrap context restored correctly |
| 195 | raise CustomError |
| 196 | |
| 197 | # check behavior matches eager in all cases |
| 198 | torch.testing.assert_close(out_fx, out_eager) |
| 199 | |
| 200 | # check that IR nodes only appear if enable_torch_wrap=True |
| 201 | ir_nodes = find_fn(torch_op, gm) |
| 202 | if enable_torch_wrap: |
| 203 | assert len(ir_nodes) == 1, gm.code |
| 204 | else: |
| 205 | assert len(ir_nodes) == 0, gm.code |
| 206 | |
| 207 | # with torch wrapping enabled (default), IR nodes appear |
| 208 | if symbolic_trace: |
| 209 | gm = torch.fx.symbolic_trace(fn) |
| 210 | else: |
| 211 | gm = make_fx(fn)(torch.randn(2, 2), torch.randn(2, 2)) |
| 212 | |
| 213 | ir_nodes = find_fn(torch_op, gm) |
| 214 | assert len(ir_nodes) == 1, gm.code |
| 215 | |
| 216 | |
| 217 | class TestIrOpImplDispatch: |
| 218 | def test_register_impl(self, custom_add_op): |
| 219 | _custom_add = custom_add_op |
| 220 | assert "impl_a" in _custom_add.impls |
| 221 | impl = _custom_add.impls["impl_a"] |
| 222 | |
| 223 | assert impl is _custom_add.impls["impl_a"] |
| 224 | assert impl.op is _custom_add |
| 225 | assert impl.provider == "impl_a" |
| 226 | assert callable(impl.impl_fn) |
| 227 | |
| 228 | # Test duplicate registration rejected |
| 229 | with pytest.raises(AssertionError): |
| 230 | |
| 231 | @_custom_add.register_impl("impl_a") |
| 232 | def impl_a_dup(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 233 | return x + y + 30 |
| 234 | |
| 235 | # Check the original impl is still intact |
| 236 | assert _custom_add.impls["impl_a"] is impl |
| 237 | |
| 238 | # Check support all args |
| 239 | assert _custom_add.impls["impl_a"].supports_all_args |
| 240 | assert _custom_add.impls["impl_b"].supports_all_args |
| 241 | assert not _custom_add.impls["impl_even"].supports_all_args |
| 242 | |
| 243 | def test_reserved_provider_rejected(self, custom_add_op): |
| 244 | _custom_add = custom_add_op |
| 245 | for provider in RESERVED_PROVIDERS: |
| 246 | with pytest.raises(AssertionError): |
| 247 | |
| 248 | @_custom_add.register_impl(provider) |
| 249 | def bad_impl(x, y): |
| 250 | return x + y |
| 251 | |
| 252 | def test_set_priority_scoped(self, custom_add_op): |
| 253 | _custom_add = custom_add_op |
| 254 | assert _custom_add.get_priority() == [] |
| 255 | |
| 256 | with _custom_add.set_priority(["impl_even", "impl_b"]): |
| 257 | assert _custom_add.get_priority() == ["impl_even", "impl_b"] |
| 258 | |
| 259 | # Check nesting |
| 260 | with _custom_add.set_priority(["impl_b"]): |
| 261 | assert _custom_add.get_priority() == ["impl_b"] |
| 262 | |
| 263 | # Restored |
| 264 | assert _custom_add.get_priority() == ["impl_even", "impl_b"] |
| 265 | |
| 266 | # Check that exception restores priority |
| 267 | with pytest.raises(CustomError), _custom_add.set_priority(["impl_a"]): |
| 268 | assert _custom_add.get_priority() == ["impl_a"] |
| 269 | raise CustomError |
| 270 | |
| 271 | # Restored again |
| 272 | assert _custom_add.get_priority() == ["impl_even", "impl_b"] |
| 273 | |
| 274 | # Restored to empty |
| 275 | assert _custom_add.get_priority() == [] |
| 276 | |
| 277 | @pytest.mark.parametrize( |
| 278 | "default,override", |
| 279 | [ |
| 280 | (["impl_even", "impl_b"], ["impl_a"]), |
| 281 | (["impl_a"], ["impl_even", "impl_b"]), |
| 282 | ], |
| 283 | ) |
| 284 | def test_set_default_priority( |
| 285 | self, custom_add_op, default: list[str], override: list[str] |
| 286 | ): |
| 287 | _custom_add = custom_add_op |
| 288 | assert _custom_add.get_priority() == [] |
| 289 | |
| 290 | _custom_add.set_default(default) |
| 291 | assert _custom_add.get_priority() == default |
| 292 | |
| 293 | # Priority doesn't change after exiting the set_priority context. |
| 294 | with _custom_add.set_priority(override): |
| 295 | assert _custom_add.get_priority() == override |
| 296 | assert _custom_add.get_priority() == default |
| 297 | |
| 298 | # Should override the previous default. |
| 299 | _custom_add.set_default(override) |
| 300 | assert _custom_add.get_priority() == override |
| 301 | |
| 302 | @pytest.mark.parametrize("overload", ["default", "maybe_inplace"]) |
| 303 | def test_dispatch_priority_order(self, custom_add_op, overload: str): |
| 304 | _custom_add = custom_add_op |
| 305 | op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace |
| 306 | torch_op = ( |
| 307 | _custom_add.torch_op |
| 308 | if overload == "default" |
| 309 | else _custom_add.maybe_inplace.torch_op |
| 310 | ) |
| 311 | |
| 312 | x = torch.tensor(1, dtype=torch.int32) |
| 313 | y = torch.tensor(2, dtype=torch.int32) |
| 314 | |
| 315 | with _custom_add.set_priority(["impl_b", "impl_a"]): |
| 316 | assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_b"] |
| 317 | out1 = op_fn(x.clone(), y) |
| 318 | out2 = torch_op(x.clone(), y) |
| 319 | |
| 320 | with _custom_add.set_priority(["impl_a"]): |
| 321 | assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_a"] |
| 322 | out3 = op_fn(x.clone(), y) |
| 323 | out4 = torch_op(x.clone(), y) |
| 324 | |
| 325 | # impl_b |
| 326 | assert out1.item() == 1 + 2 + 20 |
| 327 | assert out2.item() == 1 + 2 + 20 |
| 328 | # impl_a |
| 329 | assert out3.item() == 1 + 2 + 10 |
| 330 | assert out4.item() == 1 + 2 + 10 |
| 331 | |
| 332 | def test_unsupported_impl_filtered(self, custom_add_op): |
| 333 | _custom_add = custom_add_op |
| 334 | |
| 335 | @_custom_add.register_impl("impl_unsupported", supported=False) |
| 336 | def impl_unsupported(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 337 | return x + y + 999 |
| 338 | |
| 339 | x = torch.tensor(1, dtype=torch.int32) |
| 340 | y = torch.tensor(2, dtype=torch.int32) |
| 341 | |
| 342 | with _custom_add.set_priority(["impl_unsupported", "impl_a"]): |
| 343 | assert _custom_add.get_priority() == ["impl_a"] |
| 344 | out = _custom_add(x, y) |
| 345 | |
| 346 | # impl_unsupported skipped → impl_a |
| 347 | assert out.item() == 1 + 2 + 10 |
| 348 | |
| 349 | def test_supports_args_runtime_dispatch_and_warning( |
| 350 | self, custom_add_op, caplog_vllm: pytest.LogCaptureFixture |
| 351 | ): |
| 352 | _custom_add = custom_add_op |
| 353 | x1 = torch.ones((2, 2), dtype=torch.int32) |
| 354 | y1 = torch.full((2, 2), 2, dtype=torch.int32) |
| 355 | |
| 356 | x2 = torch.ones((2, 3), dtype=torch.int32) |
| 357 | y2 = torch.full((2, 3), 2, dtype=torch.int32) |
| 358 | |
| 359 | with ( |
| 360 | caplog_vllm.at_level(logging.WARNING), |
| 361 | _custom_add.set_priority(["impl_even"]), |
| 362 | ): |
| 363 | # Test the warning about native fallback is logged (before even dispatching) |
| 364 | assert len(caplog_vllm.records) == 1 |
| 365 | message = caplog_vllm.records[0].message |
| 366 | assert "_custom_add" in message |
| 367 | assert "fallback to native" in message |
| 368 | assert "priority" in message |
| 369 | |
| 370 | # Check dispatching |
| 371 | assert _custom_add.get_priority() == ["impl_even", "native"] |
| 372 | assert _custom_add.dispatch(x1, y1) is _custom_add.impls["impl_even"] |
| 373 | assert _custom_add.dispatch(x2, y2) is _custom_add.impls["native"] |
| 374 | |
| 375 | out1 = _custom_add(x1, y1) # size(1) == 2 → impl_even |
| 376 | out2 = _custom_add(x2, y2) # size(1) == 3 → native fallback |
| 377 | |
| 378 | # no other warnings |
| 379 | assert len(caplog_vllm.records) == 1 |
| 380 | assert torch.all(out1 == 1 + 2 + 50) |
| 381 | assert torch.all(out2 == 1 + 2) |
| 382 | |
| 383 | def test_default_priority( |
| 384 | self, |
| 385 | custom_add_op, |
| 386 | caplog_vllm: pytest.LogCaptureFixture, |
| 387 | disable_log_dedup, |
| 388 | ): |
| 389 | _custom_add = custom_add_op |
| 390 | # Make sure logs are not deduplicated to properly test the warning |
| 391 | x = torch.tensor([3], dtype=torch.int32) |
| 392 | y = torch.tensor([4], dtype=torch.int32) |
| 393 | |
| 394 | # No priority set → falls back to native |
| 395 | assert _custom_add.get_priority() == [] |
| 396 | with caplog_vllm.at_level(logging.WARNING): |
| 397 | # Native by default |
| 398 | assert _custom_add.dispatch(x, y) is _custom_add.impls["native"] |
| 399 | out = _custom_add(x, y) |
| 400 | |
| 401 | # Check dispatching to native by default |
| 402 | assert out.item() == 3 + 4 |
| 403 | |
| 404 | # Check warning |
| 405 | assert len(caplog_vllm.records) == 2 |
| 406 | message = caplog_vllm.records[0].message.lower() |
| 407 | assert "_custom_add" in message |
| 408 | assert "priority not set" in message |
| 409 | |
| 410 | |
| 411 | @pytest.mark.parametrize("default", [True, False]) |
| 412 | def test_set_default_torch_wrap(default: bool): |
| 413 | """set_default_torch_wrap permanently flips the global flag.""" |
| 414 | original = vllm.ir.op._ENABLE_TORCH_WRAP |
| 415 | try: |
| 416 | vllm.ir.set_default_torch_wrap(default) |
| 417 | assert vllm.ir.op._ENABLE_TORCH_WRAP is default |
| 418 | |
| 419 | # Flag doesn't change after exiting the enable_torch_wrap context. |
| 420 | with vllm.ir.enable_torch_wrap(not default): |
| 421 | assert vllm.ir.op._ENABLE_TORCH_WRAP is (not default) |
| 422 | assert vllm.ir.op._ENABLE_TORCH_WRAP is default |
| 423 | |
| 424 | # Should override the previous default. |
| 425 | vllm.ir.set_default_torch_wrap(not default) |
| 426 | assert vllm.ir.op._ENABLE_TORCH_WRAP is (not default) |
| 427 | finally: |
| 428 | vllm.ir.op._ENABLE_TORCH_WRAP = original |
| 429 | |
| 430 | |
| 431 | @pytest.fixture |
| 432 | def custom_mm_op(fake_vllm_ir): |
| 433 | """Fixture that registers ``_custom_mm`` (isolated by ``fake_vllm_ir``).""" |
| 434 | |
| 435 | @vllm.ir.register_op |
| 436 | def _custom_mm( |
| 437 | x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None |
| 438 | ) -> torch.Tensor: |
| 439 | tmp = x @ y |
| 440 | return tmp if bias is None else tmp + bias |
| 441 | |
| 442 | return _custom_mm |
| 443 | |
| 444 | |
| 445 | def test_default_args(custom_mm_op): |
| 446 | _custom_mm = custom_mm_op |
| 447 | |
| 448 | # Test that default args are properly applied when dispatching and calling |
| 449 | @_custom_mm.register_impl("impl_mm", supports_args=lambda x, y, bias=None: True) |
| 450 | def impl_mm( |
| 451 | x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None |
| 452 | ) -> torch.Tensor: |
| 453 | tmp = x @ y |
| 454 | return tmp + 50 if bias is None else tmp + bias + 100 |
| 455 | |
| 456 | x1 = torch.tensor([1, 2], dtype=torch.int32) |
| 457 | x2 = torch.tensor([3, 4], dtype=torch.int32) |
| 458 | |
| 459 | # Test that supports_args receives the defaulted args |
| 460 | assert impl_mm.supports_args(x1, x2) |
| 461 | with _custom_mm.set_priority(["impl_mm", "native"]): |
| 462 | assert _custom_mm.dispatch(x1, x2) is impl_mm |
| 463 | |
| 464 | |
| 465 | def test_bad_impl_registrations(custom_mm_op): |
| 466 | _custom_mm = custom_mm_op |
| 467 | # Check bad schema |
| 468 | with pytest.raises(ValueError, match="does not match native schema"): |
| 469 | |
| 470 | @_custom_mm.register_impl("impl_mm_bad_schema") |
| 471 | def impl_mm_bad_schema(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 472 | return x @ y - 1 |
| 473 | |
| 474 | with pytest.raises(ValueError, match="does not match native schema"): |
| 475 | |
| 476 | @_custom_mm.register_impl("impl_mm_bad_schema_2") |
| 477 | def impl_mm_bad_schema_2( |
| 478 | x: torch.Tensor, y: torch.Tensor, b: torch.Tensor | None = None |
| 479 | ) -> torch.Tensor: |
| 480 | return x @ y + b - 2 |
| 481 | |
| 482 | with pytest.raises(ValueError, match="does not match native schema"): |
| 483 | |
| 484 | @_custom_mm.register_impl("impl_mm_bad_schema_3") |
| 485 | def impl_mm_bad_schema_3( |
| 486 | x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor |
| 487 | ) -> torch.Tensor: |
| 488 | return x @ y + bias - 5 |
| 489 | |
| 490 | # check supports_args with incorrect params |
| 491 | with pytest.raises(ValueError, match="supports_args must be a callable"): |
| 492 | |
| 493 | @_custom_mm.register_impl("impl_mm_bad_supports_args", supports_args=True) |
| 494 | def impl_mm_bad_supports_args( |
| 495 | x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None |
| 496 | ) -> torch.Tensor: |
| 497 | return x @ y + 10 |
| 498 | |
| 499 | with pytest.raises(ValueError, match="number of parameters"): |
| 500 | |
| 501 | @_custom_mm.register_impl( |
| 502 | "impl_mm_bad_supports_args_2", supports_args=lambda x, y: True |
| 503 | ) |
| 504 | def impl_mm_bad_supports_args( |
| 505 | x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None |
| 506 | ) -> torch.Tensor: |
| 507 | return x @ y + 10 |
| 508 | |
| 509 | with pytest.raises(ValueError, match="keyword-only parameters"): |
| 510 | |
| 511 | @_custom_mm.register_impl( |
| 512 | "impl_mm_bad_supports_args_3", supports_args=lambda x, y, *, b: True |
| 513 | ) |
| 514 | def impl_mm_bad_supports_args_2( |
| 515 | x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None |
| 516 | ) -> torch.Tensor: |
| 517 | return x @ y + 20 |
| 518 | |
| 519 | with pytest.raises(ValueError, match="does not match native parameter"): |
| 520 | |
| 521 | @_custom_mm.register_impl( |
| 522 | "impl_mm_bad_supports_args_4", supports_args=lambda x, y, b: True |
| 523 | ) |
| 524 | def impl_mm_bad_supports_args_4( |
| 525 | x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None |
| 526 | ) -> torch.Tensor: |
| 527 | return x @ y + 30 |
| 528 | |
| 529 | with pytest.raises(ValueError, match="does not match native default"): |
| 530 | |
| 531 | @_custom_mm.register_impl( |
| 532 | "impl_mm_bad_supports_args_5", supports_args=lambda x, y, bias=1: True |
| 533 | ) |
| 534 | def impl_mm_bad_supports_args_5( |
| 535 | x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None |
| 536 | ) -> torch.Tensor: |
| 537 | return x @ y + 40 |
| 538 | |
| 539 | # With fixture, each test gets a fresh op with only "native" impl |
| 540 | assert set(_custom_mm.impls.keys()) == {"native"} |
| 541 | |
| 542 | |
| 543 | IMPL_OOT_SRC = """ |
| 544 | import torch |
| 545 | |
| 546 | @_custom_mm.register_impl("impl_mm_oot") |
| 547 | def impl_mm_oot( |
| 548 | x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None |
| 549 | ) -> torch.Tensor: |
| 550 | return x @ y - 99 |
| 551 | """ |
| 552 | |
| 553 | |
| 554 | def load_custom_mm_module(file_path: Path, custom_mm_op): |
| 555 | spec = importlib.util.spec_from_file_location("_custom_mm_oot", file_path) |
| 556 | assert spec is not None |
| 557 | module = importlib.util.module_from_spec(spec) |
| 558 | |
| 559 | # Inject the variable into the module's global namespace |
| 560 | # This allows the @_custom_mm.register_impl decorator to work |
| 561 | module._custom_mm = custom_mm_op # type: ignore[attr-defined] |
| 562 | |
| 563 | # Execute the file; this triggers the decorator |
| 564 | assert spec.loader is not None |
| 565 | spec.loader.exec_module(module) |
| 566 | return module |
| 567 | |
| 568 | |
| 569 | def test_uuid_and_oot(custom_mm_op, tmp_path: Path): |
| 570 | _custom_mm = custom_mm_op |
| 571 | file_path = tmp_path / "_custom_mm_oot.py" |
| 572 | file_path.write_text(IMPL_OOT_SRC) |
| 573 | |
| 574 | assert "impl_mm_oot" not in _custom_mm.impls |
| 575 | _ = load_custom_mm_module(file_path, _custom_mm) |
| 576 | assert "impl_mm_oot" in _custom_mm.impls |
| 577 | |
| 578 | uuid = _custom_mm.impls["impl_mm_oot"].uuid() |
| 579 | del _custom_mm.impls["impl_mm_oot"] |
| 580 | |
| 581 | # Replace file source |
| 582 | file_path.write_text(IMPL_OOT_SRC + " # added file source") |
| 583 | assert "impl_mm_oot" not in _custom_mm.impls |
| 584 | _ = load_custom_mm_module(file_path, _custom_mm) |
| 585 | assert "impl_mm_oot" in _custom_mm.impls |
| 586 | |
| 587 | uuid1 = _custom_mm.impls["impl_mm_oot"].uuid() |
| 588 | assert uuid1 != uuid |
| 589 | del _custom_mm.impls["impl_mm_oot"] |
| 590 | |
| 591 | # Back to original |
| 592 | file_path.write_text(IMPL_OOT_SRC) |
| 593 | assert "impl_mm_oot" not in _custom_mm.impls |
| 594 | _ = load_custom_mm_module(file_path, _custom_mm) |
| 595 | assert "impl_mm_oot" in _custom_mm.impls |
| 596 | |
| 597 | uuid2 = _custom_mm.impls["impl_mm_oot"].uuid() |
| 598 | assert uuid2 == uuid |
| 599 | assert uuid2 != uuid1 |
| 600 | del _custom_mm.impls["impl_mm_oot"] |
| 601 | |
| 602 | |
| 603 | def _test_native(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 604 | return x + y |
| 605 | |
| 606 | |
| 607 | def _make_op_with_generator(name: str = "_ig_test"): |
| 608 | op = IrOp(name, _test_native) |
| 609 | |
| 610 | @op.register_input_generator |
| 611 | def _gen(n: int = 4): |
| 612 | x = torch.randn(n, 3) |
| 613 | y = torch.randn(n, 3) |
| 614 | return x, y |
| 615 | |
| 616 | return op |
| 617 | |
| 618 | |
| 619 | def _test_native_single(x: torch.Tensor) -> torch.Tensor: |
| 620 | return x |
| 621 | |
| 622 | |
| 623 | class TestInputGenerator: |
| 624 | def test_no_input_generator_by_default(self): |
| 625 | op = IrOp("_ig_test_no_gen", _test_native_single) |
| 626 | assert not op.has_input_generator |
| 627 | |
| 628 | def test_register_input_generator(self): |
| 629 | op = _make_op_with_generator("_ig_test_reg") |
| 630 | assert op.has_input_generator |
| 631 | |
| 632 | def test_generate_inputs_returns_tuple(self): |
| 633 | op = _make_op_with_generator("_ig_test_tuple") |
| 634 | result = op.generate_inputs(n=2) |
| 635 | assert isinstance(result, tuple) |
| 636 | assert len(result) == 2 |
| 637 | assert result[0].shape == (2, 3) |
| 638 | assert result[1].shape == (2, 3) |
| 639 | |
| 640 | def test_generate_inputs_default_kwargs(self): |
| 641 | op = _make_op_with_generator("_ig_test_default") |
| 642 | result = op.generate_inputs() |
| 643 | assert result[0].shape == (4, 3) |
| 644 | |
| 645 | def test_generate_inputs_without_registration_raises(self): |
| 646 | op = IrOp("_ig_test_no_gen_raises", _test_native_single) |
| 647 | with pytest.raises(RuntimeError, match="No input generator"): |
| 648 | op.generate_inputs() |
| 649 | |
| 650 | |
| 651 | class TestTolerance: |
| 652 | def test_override_and_get_tolerance(self): |
| 653 | op = IrOp("_tol_test", _test_native) |
| 654 | |
| 655 | tol = op.get_tolerance(torch.float32) |
| 656 | assert tol == {"atol": 1e-5, "rtol": 1.3e-6} |
| 657 | |
| 658 | op.override_tolerance(torch.float32, atol=0.1, rtol=0.2) |
| 659 | assert op.get_tolerance(torch.float32) == {"atol": 0.1, "rtol": 0.2} |
| 660 | assert op.get_tolerance(torch.float16) == {"atol": 1e-3, "rtol": 1e-3} |
| 661 | |
| 662 | def test_get_tolerance_raises_for_unknown_dtype(self): |
| 663 | op = IrOp("_tol_test_unknown", _test_native) |
| 664 | with pytest.raises(ValueError, match="No tolerance defined"): |
| 665 | op.get_tolerance(torch.complex64) |
| 666 | |
| 667 | |
| 668 | def test_naming_validation(fake_vllm_ir): |
| 669 | """Test that op and provider names are validated ([a-z_][a-z_0-9]*).""" |
| 670 | |
| 671 | # Valid op and provider names |
| 672 | @vllm.ir.register_op |
| 673 | def _valid_name_123(x: torch.Tensor) -> torch.Tensor: |
| 674 | return x |
| 675 | |
| 676 | @_valid_name_123.register_impl("valid_provider_123") |
| 677 | def valid_impl(x: torch.Tensor) -> torch.Tensor: |
| 678 | return x + 1 |
| 679 | |
| 680 | # Invalid op names should fail |
| 681 | with pytest.raises(ValueError, match="name.*invalid"): |
| 682 | |
| 683 | @vllm.ir.register_op |
| 684 | def InvalidName(x: torch.Tensor) -> torch.Tensor: |
| 685 | return x |
| 686 | |
| 687 | with pytest.raises(ValueError, match="name.*invalid"): |
| 688 | |
| 689 | @vllm.ir.register_op(name="123invalid") |
| 690 | def some_func(x: torch.Tensor) -> torch.Tensor: |
| 691 | return x |
| 692 | |
| 693 | # Invalid provider names should fail |
| 694 | with pytest.raises(ValueError, match="name.*invalid"): |
| 695 | |
| 696 | @_valid_name_123.register_impl("Invalid-Provider") |
| 697 | def invalid_impl(x: torch.Tensor) -> torch.Tensor: |
| 698 | return x + 1 |
| 699 | |
| 700 | |
| 701 | def test_registration_stack_traces(fake_vllm_ir): |
| 702 | """Test that stack traces are captured for ops and impls.""" |
| 703 | |
| 704 | @vllm.ir.register_op |
| 705 | def _test_stack(x: torch.Tensor) -> torch.Tensor: |
| 706 | return x |
| 707 | |
| 708 | @_test_stack.register_impl("test_provider") |
| 709 | def test_impl(x: torch.Tensor) -> torch.Tensor: |
| 710 | return x + 1 |
| 711 | |
| 712 | # Verify op stack trace |
| 713 | assert hasattr(_test_stack, "_registration_stack") |
| 714 | assert len(_test_stack._registration_stack) > 0 |
| 715 | op_stack_str = "".join(_test_stack._registration_stack) |
| 716 | assert "test_op.py" in op_stack_str |
| 717 | # Last frame should be the decorator in user code, not internal decorator logic |
| 718 | assert "@vllm.ir.register_op" in _test_stack._registration_stack[-1] |
| 719 | assert "return decorator(f)" not in op_stack_str |
| 720 | |
| 721 | # Verify impl stack trace |
| 722 | impl = _test_stack.impls["test_provider"] |
| 723 | assert hasattr(impl, "_registration_stack") |
| 724 | assert len(impl._registration_stack) > 0 |
| 725 | impl_stack_str = "".join(impl._registration_stack) |
| 726 | assert "test_op.py" in impl_stack_str |
| 727 | # Last frame should be the decorator in user code |
| 728 | assert '@_test_stack.register_impl("test_provider")' in impl._registration_stack[-1] |
| 729 | |
| 730 | |
| 731 | def test_op_repr_uses_docstring(fake_vllm_ir): |
| 732 | """Test that __str__ uses the function's docstring and __repr__ is simple.""" |
| 733 | |
| 734 | @vllm.ir.register_op |
| 735 | def _test_repr_with_doc(x: torch.Tensor) -> torch.Tensor: |
| 736 | """First line of docstring. |
| 737 | |
| 738 | Additional details here. |
| 739 | """ |
| 740 | return x |
| 741 | |
| 742 | @vllm.ir.register_op |
| 743 | def _test_repr_no_doc(x: torch.Tensor) -> torch.Tensor: |
| 744 | return x |
| 745 | |
| 746 | # __str__ with docstring: uses first line only |
| 747 | str_with = str(_test_repr_with_doc) |
| 748 | assert "IrOp('_test_repr_with_doc')" in str_with |
| 749 | assert "First line of docstring." in str_with |
| 750 | assert "Additional details" not in str_with |
| 751 | |
| 752 | # __str__ without docstring: simple format |
| 753 | assert str(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')" |
| 754 | |
| 755 | # __repr__ should be simple for both |
| 756 | assert repr(_test_repr_with_doc) == "IrOp('_test_repr_with_doc')" |
| 757 | assert repr(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')" |
| 758 | |
| 759 | |
| 760 | def test_vllm_ir_fixture(fake_vllm_ir): |
| 761 | """Test that the fake_vllm_ir fixture provides test isolation.""" |
| 762 | |
| 763 | @vllm.ir.register_op |
| 764 | def _test_fixture(x: torch.Tensor) -> torch.Tensor: |
| 765 | return x |
| 766 | |
| 767 | assert "_test_fixture" in IrOp.registry |
| 768 | # Fixture will automatically clean up after test |