Some fixes

  * register in setup entrypoints
  * add ut for activation

Signed-off-by: MengqingCao <cmq0113@163.com>

Author: MengqingCao

setup.py

@@ -391,7 +391,9 @@ def _read_requirements(filename: str) -> List[str]:
     extras_require={},
     entry_points={
         "vllm.platform_plugins": ["ascend = vllm_ascend:register"],
-        "vllm.general_plugins":
-        ["ascend_enhanced_model = vllm_ascend:register_model"],
+        "vllm.general_plugins": [
+            "ascend_enhanced_model = vllm_ascend:register_model",
+            "dummy_custom_ops = vllm_ascend:register_ops"
+        ],
     },
 )

tests/ut/ops/test_activation.py

@@ -0,0 +1,48 @@
+from unittest.mock import patch
+
+import pytest
+import torch
+from vllm.model_executor.layers.activation import QuickGELU, SiluAndMul
+
+import vllm_ascend.patch.worker.patch_common.patch_utils  # noqa: F401
+
+
+@pytest.fixture
+def dummy_tensor():
+    return torch.randn(4, 8, dtype=torch.float16)
+
+
+@patch("torch_npu.npu_fast_gelu", side_effect=lambda x: x + 1)
+def test_QuickGELU_forward(mock_gelu, dummy_tensor):
+    layer = QuickGELU()
+    out = layer.forward(dummy_tensor)
+
+    expected_out = dummy_tensor + 1
+    assert torch.allclose(out, expected_out)
+
+    mock_gelu.assert_called_once()
+
+
+@pytest.mark.parametrize("is_310p_return", [True, False])
+@patch("torch_npu.npu_swiglu", side_effect=lambda x: x + 1)
+def test_SiluAndMul_forward(mock_swiglu, is_310p_return, dummy_tensor):
+
+    with patch("vllm_ascend.utils.is_310p", return_value=is_310p_return):
+        layer = SiluAndMul()
+        out = layer.forward(dummy_tensor)
+
+        if is_310p_return:
+            expected_arg = dummy_tensor.to(torch.float32)
+        else:
+            expected_arg = dummy_tensor
+
+        # assert mock_swiglu.call_count == 1
+        mock_swiglu.assert_called_once()
+
+        actual_arg = mock_swiglu.call_args[0][0]
+        assert torch.allclose(
+            actual_arg,
+            expected_arg), "npu_swiglu called with unexpected input"
+
+        expected_out = dummy_tensor + 1
+        assert torch.allclose(out, expected_out)

vllm_ascend/__init__.py

@@ -29,3 +29,7 @@ def register_model():
 
     from .models import register_model
     register_model()
+
+
+def register_ops():
+    import vllm_ascend.ops  # noqa: F401

vllm_ascend/ops/activation.py

@@ -18,8 +18,6 @@
 import torch
 from vllm.model_executor.layers.activation import QuickGELU, SiluAndMul
 
-from vllm_ascend.utils import is_310p
-
 
 @QuickGELU.register_oot
 class AscendQuickGELU(QuickGELU):
@@ -37,6 +35,8 @@ class AscendSiluAndMul(SiluAndMul):
     def forward_oot(self, x: torch.Tensor) -> torch.Tensor:
         import torch_npu
 
+        from vllm_ascend.utils import is_310p
+
         if is_310p():
             out = torch_npu.npu_swiglu(x.to(torch.float32)).to(torch.float16)
         else: