add st for moe token dispatcher

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>

Author: weijinqian_v1

tests/ut/moe_util.py renamed to tests/ut/test_moe_util.py

@@ -4,9 +4,9 @@
 import torch
 import pytest
 import math
+import vllm_ascend.patch.worker.patch_common.patch_utils
 
-from vllm_ascend.ops.moe_dispatcher.moe_utils import permute, get_capacity, topk_softmax_with_capacity, \
-    group_limited_topk, unpermute, sort_chunks_by_idxs
+from vllm_ascend.ops.moe_dispatcher.moe_utils import permute, get_capacity, topk_softmax_with_capacity, group_limited_topk, unpermute, sort_chunks_by_idxs
 
 
 class TestMoeUtils:
@@ -22,6 +22,7 @@ def setup(self):
         self.num_groups = 2
         self.scaling_factor = 1.0
 
+
     def test_group_limited_topk(self, setup):
         # Test group-limited topk routing
         scores = torch.randn(self.num_tokens, self.num_experts)
@@ -38,42 +39,33 @@ def test_group_limited_topk(self, setup):
         assert indices.shape == (self.num_tokens, self.topk)
         assert torch.all(indices < self.num_experts)
 
-    def test_topk_softmax_with_capacity(self, setup):
+
+    @pytest.mark.parametrize("score_function", ["softmax"])
+    def test_topk_softmax_with_capacity(self, setup, score_function):
         # Test topk softmax with capacity
         logits = torch.randn(self.num_tokens, self.num_experts)
 
         # Test without capacity
         probs, routing_map, tokens_per_expert, top_indices = topk_softmax_with_capacity(
             logits,
-            topk=self.topk
+            topk=self.topk,
+            score_function=score_function
         )
         assert probs.shape == (self.num_tokens, self.num_experts)
         assert routing_map.shape == (self.num_tokens, self.num_experts)
         assert tokens_per_expert.shape == (self.num_experts,)
 
-        # Test with capacity
-        probs, routing_map, tokens_per_expert, top_indices = topk_softmax_with_capacity(
-            logits,
-            topk=self.topk,
-            capacity_factor=self.capacity_factor,
-            pad_to_capacity=True
-        )
-        expert_capacity = get_capacity(
-            num_tokens=self.num_tokens * self.topk,
-            num_experts=self.num_experts,
-            capacity_factor=self.capacity_factor
-        )
-        assert tokens_per_expert.max() <= expert_capacity
-
         # Test with group routing
         probs, routing_map, tokens_per_expert, top_indices = topk_softmax_with_capacity(
             logits,
             topk=self.topk,
             num_groups=self.num_groups,
-            group_topk=self.group_topk
+            group_topk=self.group_topk,
+            score_function=score_function
         )
         assert probs.shape == (self.num_tokens, self.num_experts)
 
+
     def test_get_capacity(self, setup):
         # Test capacity calculation
         capacity = get_capacity(
@@ -94,6 +86,7 @@ def test_get_capacity(self, setup):
         )
         assert capacity == min_capacity
 
+
     def test_permute(self, setup):
         # Test token permutation
         tokens = torch.randn(self.num_tokens, self.hidden_size)
@@ -120,6 +113,7 @@ def test_permute(self, setup):
         assert permuted_tokens.shape[0] == num_out_tokens
         assert sorted_indices.shape[0] == num_out_tokens
 
+
     def test_unpermute(self, setup):
         # Test token unpermutation
         tokens = torch.randn(self.num_tokens, self.hidden_size)
@@ -162,6 +156,7 @@ def test_unpermute(self, setup):
         )
         assert restored_tokens.shape == tokens.shape
 
+
     def test_sort_chunks_by_idxs(self, setup):
         # Test chunk sorting
         input_tensor = torch.randn(10, self.hidden_size)
@@ -173,10 +168,10 @@ def test_sort_chunks_by_idxs(self, setup):
 
         # Verify the order is correct
         expected = torch.cat([input_tensor[5:], input_tensor[0: 3], input_tensor[3: 5]])
-        assert torch.allclose(output, expected) \
- \
-               @ pytest.mark.parametrize("score_function", ["softmax", "sigmoid"])
+        assert torch.allclose(output, expected)
 
+
+    @pytest.mark.parametrize("score_function", ["softmax"])
     def test_score_functions(self, setup, score_function):
         # Test different score functions
         logits = torch.randn(self.num_tokens, self.num_experts)
@@ -190,28 +185,4 @@ def test_score_functions(self, setup, score_function):
         )
         assert probs.shape == (self.num_tokens, self.num_experts)
         assert routing_map.shape == (self.num_tokens, self.num_experts)
-        assert tokens_per_expert.shape == (self.num_experts,)
-
-    def test_edge_cases(self, setup):
-        # Test empty input
-        empty_logits = torch.randn(0, self.num_experts)
-        with pytest.raises(AssertionError):
-            topk_softmax_with_capacity(empty_logits, topk=self.topk)
-
-        # Test invalid score function
-        logits = torch.randn(self.num_tokens, self.num_experts)
-        with pytest.raises(ValueError):
-            topk_softmax_with_capacity(
-                logits,
-                topk=self.topk,
-                score_function="invalid"
-            )
-
-        # Test invalid drop policy
-        with pytest.raises(ValueError):
-            topk_softmax_with_capacity(
-                logits,
-                topk=self.topk,
-                capacity_factor=1.0,
-                drop_policy="invalid"
-            )
+        assert tokens_per_expert.shape == (self.num_experts,)

tests/ut/test_token_dispatcher.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+
+import torch
+import pytest
+from pytest_mock import MockerFixture
+import vllm_ascend.patch.worker.patch_common.patch_utils
+from vllm_ascend.utils import adapt_patch # noqa E402
+
+from vllm_ascend.ops.moe_dispatcher.token_dispatcher import MoeDispatcherConfig, MoEAlltoAllSeqOverLapDispatcher
+
+adapt_patch(True)
+
+class TestMoEAlltoAllSeqOverLapDispatcher:
+
+    @pytest.fixture
+    def config(self):
+        config = MoeDispatcherConfig()
+        config.set_num_local_experts(2)
+        config.set_num_moe_experts(4)
+        config.set_moe_pad_expert_input_to_capacity(False)
+        config.set_moe_expert_capacity_factor(None)
+        config.set_moe_router_topk(2)
+        config.set_moe_grouped_gemm(False)
+        config.set_group_topk(0)
+        config.set_num_groups(1)
+        config.set_is_fused(False)
+        return config.build()
+
+    def mock_ep_group(self, mocker):
+        mock_group = mocker.MagicMock()
+        mock_group.rank_in_group = 0
+        mock_group.world_size = 2
+        mock_group.device_group = "mock_group"
+        return mock_group
+
+    @pytest.fixture
+    def dispatcher(self, config, mocker: MockerFixture):
+        mocker.patch("vllm_ascend.ops.moe_dispatcher.token_dispatcher.get_ep_group",
+                     return_value=self.mock_ep_group(mocker))
+        return MoEAlltoAllSeqOverLapDispatcher(config)
+
+    def test_initialization(self, dispatcher, config):
+        assert dispatcher.num_local_experts == config.num_local_experts
+        assert dispatcher.num_experts == config.num_moe_experts
+        assert dispatcher.local_expert_indices == [0, 1]
+        assert dispatcher.ep_rank == 0
+        assert dispatcher.ep_size == 2
+        assert dispatcher.overlap_stream is not None
+
+    def test_routing(self, dispatcher):
+        probs = torch.randn(4, 4)  # 4 tokens, 4 experts
+        scores, routing_map = dispatcher.routing(probs)
+        assert scores.shape == (4, 4)  # topk=2
+        assert routing_map.shape == (4, 4)