xx

Author: wangxiaoxin-sherie

tests/ut/ops/test_ascend_forwad_context.py

@@ -0,0 +1,21 @@
+import unittest
+import os
+from unittest import mock
+
+from vllm_ascend.ascend_forward_context import get_dispatcher_name
+
+
+class TestGetDispatcherName(unittest.TestCase):
+
+    def test_get_dispatcher_name(self):
+        result = get_dispatcher_name(1, False)
+        assert result == "TokenDispatcherWithAllGather"
+        result = get_dispatcher_name(4, False)
+        assert result == "TokenDispatcherWithAll2AllV"
+        result = get_dispatcher_name(16, True)
+        assert result == "TokenDispatcherWithAll2AllV"
+        result = get_dispatcher_name(16, False)
+        assert result == "TokenDispatcherWithMC2"
+        with mock.patch.dict(os.environ, {"VLLM_ENABLE_FUSED_EXPERTS_ALLGATHER_EP": "1"}):
+            result = get_dispatcher_name(16, False)
+            assert result == "TokenDispatcherWithAllGather"