| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import torch |
| 5 | from torch import Tensor |
| 6 | from torch.fx.experimental.proxy_tensor import make_fx |
| 7 | |
| 8 | import vllm.ir.op |
| 9 | from vllm.ir.op import IrOp, IrOpInplaceOverload |
| 10 | |
| 11 | pytestmark = pytest.mark.skip_global_cleanup |
| 12 | |
| 13 | |
| 14 | @vllm.ir.register_op(allow_inplace=True) |
| 15 | def _custom_mm2(x: Tensor, w: Tensor) -> Tensor: |
| 16 | return x @ w |
| 17 | |
| 18 | |
| 19 | @_custom_mm2.register_impl("regular") |
| 20 | def _custom_mm2_regular(x: Tensor, w: Tensor) -> Tensor: |
| 21 | return x @ w + 1 |
| 22 | |
| 23 | |
| 24 | @_custom_mm2.register_impl("inplace", inplace=True) |
| 25 | def _custom_mm2_inplace(x: Tensor, w: Tensor) -> Tensor: |
| 26 | x.copy_(x @ w + 2) |
| 27 | return x |
| 28 | |
| 29 | |
| 30 | class TestInplaceOp: |
| 31 | def test_registration(self): |
| 32 | # Test that the inplace op is registered correctly. |
| 33 | assert "_custom_mm2" in IrOp.registry |
| 34 | assert IrOp.registry["_custom_mm2"] is _custom_mm2 |
| 35 | assert _custom_mm2.torch_op is torch.ops.vllm_ir._custom_mm2.default |
| 36 | assert isinstance(_custom_mm2.maybe_inplace, IrOpInplaceOverload) |
| 37 | assert ( |
| 38 | _custom_mm2.maybe_inplace.torch_op |
| 39 | is torch.ops.vllm_ir._custom_mm2.maybe_inplace |
| 40 | ) |
| 41 | |
| 42 | def test_inplace_dispatching(self): |
| 43 | # check that the correct implementation is dispatched based on priority, |
| 44 | # and inplace semantics hold |
| 45 | w = torch.randn(3, 3) |
| 46 | x = torch.randn(2, 3) |
| 47 | x1 = x.clone() |
| 48 | |
| 49 | with _custom_mm2.set_priority(["regular"]): |
| 50 | result_regular = _custom_mm2.maybe_inplace(x, w) |
| 51 | |
| 52 | # check that the regular op does not modify x |
| 53 | torch.testing.assert_close(x, x1, atol=0, rtol=0) |
| 54 | |
| 55 | with _custom_mm2.set_priority(["inplace"]): |
| 56 | result_inplace: Tensor = _custom_mm2.maybe_inplace(x, w) |
| 57 | |
| 58 | # check that the inplace op returns x directly |
| 59 | assert result_inplace.data_ptr() == x.data_ptr() |
| 60 | |
| 61 | torch.testing.assert_close(result_inplace, x1 @ w + 2) |
| 62 | torch.testing.assert_close(result_regular, x1 @ w + 1) |
| 63 | |
| 64 | def test_default_dispatching(self): |
| 65 | # check that the correct implementation is dispatched, |
| 66 | # and ops do not modify inputs when using the default overload |
| 67 | w = torch.randn(3, 3) |
| 68 | x = torch.randn(2, 3) |
| 69 | x1 = x.clone() |
| 70 | |
| 71 | with _custom_mm2.set_priority(["regular"]): |
| 72 | result_regular = _custom_mm2(x, w) |
| 73 | |
| 74 | with _custom_mm2.set_priority(["inplace"]): |
| 75 | result_inplace = _custom_mm2(x, w) |
| 76 | |
| 77 | # check that x was not modified by either impl |
| 78 | torch.testing.assert_close(x, x1, atol=0, rtol=0) |
| 79 | |
| 80 | torch.testing.assert_close(result_inplace, x1 @ w + 2) |
| 81 | torch.testing.assert_close(result_regular, x1 @ w + 1) |
| 82 | |
| 83 | def test_trace(self): |
| 84 | # Test that the inplace op can be used in a graph. |
| 85 | def func(x: Tensor, y: Tensor) -> Tensor: |
| 86 | return _custom_mm2.maybe_inplace(x, y) |
| 87 | |
| 88 | x = torch.randn(2, 3) |
| 89 | y = torch.randn(3, 4) |
| 90 | graph = make_fx(func)(x, y) |
| 91 | assert any( |
| 92 | node.target == torch.ops.vllm_ir._custom_mm2.maybe_inplace |
| 93 | for node in graph.graph.nodes |
| 94 | ) |