# # Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved. # Copyright 2023 The vLLM team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # This file is a part of the vllm-ascend project. from unittest.mock import MagicMock, PropertyMock, patch import numpy as np import pytest import torch from tests.ut.base import TestBase from vllm_ascend.ops.fused_moe.moe_runtime_args import ( MoEAllGatherCombineMetadata, MoEAllToAllCombineMetadata, MoEMC2CombineMetadata, MoEQuantParams, MoERoutingParams, MoETokenDispatchInput, ) from vllm_ascend.ops.fused_moe.token_dispatcher import ( # isort: skip AscendDeviceType, EXPERT_TOKEN_NUMS_TYPE_COUNT, EXPERT_TOKEN_NUMS_TYPE_CUMSUM, TokenDispatcherWithAll2AllV, TokenDispatcherWithAllGather, TokenDispatcherWithMC2, ) from vllm_ascend.ops.fused_moe.moe_stage_params import MoEMxfpParams from vllm_ascend.quantization.quant_type import QuantType MXFP4_TEST_DTYPE = getattr(torch, "float4_e2m1fn_x2", torch.float16) def build_token_dispatch_input_fixture( *, hidden_states: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, expert_map: torch.Tensor | None = None, global_redundant_expert_num: int = 0, apply_router_weight_on_input: bool = False, pertoken_scale: torch.Tensor | None = None, quant_type: QuantType = QuantType.NONE, comm_quant_mode: int | None = None, act_quant_type: torch.dtype | None = None, is_per_channel_weight: bool = False, mc2_mask: torch.Tensor | None = None, ) -> MoETokenDispatchInput: mxfp_spec = None if quant_type in (QuantType.MXFP8, QuantType.MXFP4): mxfp_spec = MoEMxfpParams(act_quant_type=act_quant_type) return MoETokenDispatchInput( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, routing=MoERoutingParams( expert_map=expert_map, global_redundant_expert_num=global_redundant_expert_num, mc2_mask=mc2_mask, apply_router_weight_on_input=apply_router_weight_on_input, pertoken_scale=pertoken_scale, ), quant=MoEQuantParams( quant_type=quant_type, comm_quant_mode=comm_quant_mode, mxfp=mxfp_spec, is_per_channel_weight=is_per_channel_weight, ), ) class TestTokenDispatcherWithMC2(TestBase): def setUp(self): self.config_patcher = patch("vllm_ascend.ops.fused_moe.token_dispatcher.get_current_vllm_config") self.mock_get_config = self.config_patcher.start() mock_config = MagicMock() mock_config.scheduler_config.max_num_seqs = 256 mock_config.compilation_config.custom_ops = ["all"] mock_config.speculative_config = None mock_config.parallel_config.tensor_parallel_size = 1 self.mock_get_config.return_value = mock_config self.mc2_tokens_capacity = 128 self.mc2_capacity_patch = patch( "vllm_ascend.ops.fused_moe.token_dispatcher.get_mc2_tokens_capacity", return_value=self.mc2_tokens_capacity, ) self.mock_get_mc2_tokens_capacity = self.mc2_capacity_patch.start() self.mc2_group = MagicMock() self.mc2_group.device_group.return_value._get_backend.return_value.get_hccl_comm_name.return_value = "hccl_123" self.mc2_group.rank_in_group = 0 self.mc2_group.world_size = 8 self.mc2_group_patch = patch( "vllm_ascend.ops.fused_moe.token_dispatcher.get_mc2_group", return_value=self.mc2_group ) self.mc2_group_patch.start() self.rank_group_patch = patch("torch.distributed.get_rank", return_value=0) self.rank_group_patch.start() # Mock get_forward_context().mc2_mask self.forward_context = MagicMock() self.forward_context.mc2_mask = torch.tensor([1, 0, 1]) self.forward_context_patch = patch( "vllm.forward_context.get_forward_context", return_value=self.forward_context ) self.forward_context_patch.start() # Mock get_ascend_device_type() self.ascend_soc_version_patch = patch( "vllm_ascend.ops.fused_moe.token_dispatcher.get_ascend_device_type", return_value=AscendDeviceType.A3 ) self.ascend_soc_version_patch.start() # Mock get_ascend_config() and is_hierarchical_communication_enabled() mock_ascend_config = MagicMock() mock_ascend_config.enable_mc2_hierarchy_comm = False mock_ascend_config.eplb_config = MagicMock() mock_ascend_config.eplb_config.dynamic_eplb = False self.ascend_config_patch = patch( "vllm_ascend.ops.fused_moe.token_dispatcher.get_ascend_config", return_value=mock_ascend_config ) self.ascend_config_patch.start() self.ascend_config_utils_patch = patch("vllm_ascend.utils.get_ascend_config", return_value=mock_ascend_config) self.ascend_config_utils_patch.start() self.hier_comm_patch = patch( "vllm_ascend.ops.fused_moe.token_dispatcher.is_hierarchical_communication_enabled", return_value=False ) self.hier_comm_patch.start() self.skip_allreduce_patch = patch( "vllm_ascend.ops.fused_moe.token_dispatcher.should_skip_allreduce_across_dp_group", return_value=False ) self.mock_skip_allreduce = self.skip_allreduce_patch.start() kwargs = {"with_quant": False, "top_k": 8, "num_experts": 128} self.dispatcher = TokenDispatcherWithMC2(**kwargs) def tearDown(self): self.config_patcher.stop() self.mc2_capacity_patch.stop() self.mc2_group_patch.stop() self.rank_group_patch.stop() self.forward_context_patch.stop() self.ascend_soc_version_patch.stop() self.ascend_config_patch.stop() self.ascend_config_utils_patch.stop() self.hier_comm_patch.stop() self.skip_allreduce_patch.stop() def test_init(self): self.assertEqual(self.dispatcher.ep_rank_id, 0) self.assertEqual(self.dispatcher.ep_world_size, 8) self.assertTrue(self.dispatcher.enable_dispatch_v2) self.assertTrue(self.dispatcher.need_extra_args) self.assertEqual(self.dispatcher.global_bs, 0) def test_init_uses_mc2_capacity_for_non_uniform_global_bs(self): self.mock_get_config.return_value.parallel_config.tensor_parallel_size = 4 self.mock_skip_allreduce.return_value = True dispatcher = TokenDispatcherWithMC2(with_quant=False, top_k=8, num_experts=128) self.assertEqual(dispatcher.global_bs, 256) def test_get_dispatch_mc2_kwargs_with_skip_allreduce_omits_mc2_mask(self): self.mock_get_config.return_value.parallel_config.tensor_parallel_size = 4 self.mock_skip_allreduce.return_value = True dispatcher = TokenDispatcherWithMC2(with_quant=False, top_k=8, num_experts=128) hidden_states = torch.randn(10, 128) topk_ids = torch.randint(0, 8, (10, 1)) topk_weights = torch.randn(10, 1) expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]) mc2_mask = torch.tensor([True, False, True, False]) token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, mc2_mask=mc2_mask, ) kwargs = dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input) self.assertEqual(kwargs["global_bs"], 256) self.assertNotIn("x_active_mask", kwargs) def test_get_dispatch_mc2_kwargs_without_skip_allreduce_keeps_mc2_mask(self): hidden_states = torch.randn(10, 128) topk_ids = torch.randint(0, 8, (10, 1)) topk_weights = torch.randn(10, 1) expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]) mc2_mask = torch.tensor([True, False, True, False]) token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, mc2_mask=mc2_mask, ) kwargs = self.dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input) self.assertEqual(kwargs["global_bs"], 0) self.assertIs(kwargs["x_active_mask"], mc2_mask) def test_get_dispatch_mc2_kwargs_without_quant(self): hidden_states = torch.randn(10, 128) topk_ids = torch.randint(0, 8, (10, 1)) topk_weights = torch.randn(10, 1) expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]) token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, global_redundant_expert_num=0, apply_router_weight_on_input=False, pertoken_scale=None, ) kwargs = self.dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input) self.assertIn("x", kwargs) self.assertIn("expert_ids", kwargs) self.assertEqual(kwargs["moe_expert_num"], 8) def test_token_permutation_dispatch(self): hidden_states = torch.randn(10, 128) topk_weights = torch.randn(10, 1) topk_ids = torch.randint(0, 8, (10, 1)) expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]) with patch( "torch_npu.npu_moe_distribute_dispatch_v2", return_value=(torch.randn(10, 128),) * 5 + (None, None) ) as mock_dispatch: token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, ) output = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input) mock_dispatch.assert_called_once() self.assertEqual(output.group_list_type, 0) # group_list_type == 0 self.assertIsInstance(output.combine_metadata, MoEMC2CombineMetadata) def test_w4a8_per_channel_dispatch_uses_count_group_list(self): hidden_states = torch.randn(10, 128) topk_weights = torch.randn(10, 1) topk_ids = torch.randint(0, 8, (10, 1)) expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]) self.dispatcher.enable_dispatch_v2 = True token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, quant_type=QuantType.W4A8, is_per_channel_weight=True, ) with patch( "torch_npu.npu_moe_distribute_dispatch_v2", return_value=(torch.randn(10, 128),) * 5 + (None, None) ) as mock_dispatch: output = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input) mock_dispatch.assert_called_once() self.assertEqual(mock_dispatch.call_args.kwargs["expert_token_nums_type"], EXPERT_TOKEN_NUMS_TYPE_COUNT) self.assertEqual(output.group_list_type, EXPERT_TOKEN_NUMS_TYPE_COUNT) def test_w4a8_group_dispatch_keeps_prefix_sum_group_list(self): hidden_states = torch.randn(10, 128) topk_weights = torch.randn(10, 1) topk_ids = torch.randint(0, 8, (10, 1)) expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]) token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, quant_type=QuantType.W4A8, is_per_channel_weight=False, ) kwargs = self.dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input) self.assertEqual(kwargs["expert_token_nums_type"], EXPERT_TOKEN_NUMS_TYPE_CUMSUM) def test_get_combine_mc_kwargs_with_quant(self): hidden_states = torch.randn(10, 128) topk_ids = torch.randint(0, 8, (10, 1)) topk_weights = torch.randn(10, 1) expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]) ep_recv_counts = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]) tp_recv_counts = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]) assist_info_for_combine = torch.arange(10) token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, ) combine_metadata = MoEMC2CombineMetadata( topk_ids=topk_ids, topk_weights=topk_weights, expert_map=expert_map, ep_recv_counts=ep_recv_counts, tp_recv_counts=tp_recv_counts, assist_info_for_combine=assist_info_for_combine, expand_scales=None, quant=token_dispatch_input.quant, ) self.dispatcher.need_extra_args = True self.dispatcher.enable_dispatch_v2 = True self.dispatcher.moe_expert_num = len(expert_map) kwargs = self.dispatcher.get_combine_mc_kwargs(hidden_states, combine_metadata) self.assertIn("tp_send_counts", kwargs) def test_get_dispatch_mc2_kwargs_with_mxfp8_quant(self): hidden_states = torch.randn(10, 128) topk_ids = torch.randint(0, 8, (10, 1)) topk_weights = torch.randn(10, 1) expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]) self.dispatcher.a5_need_extra_args = True token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, quant_type=QuantType.MXFP8, act_quant_type=torch.float8_e4m3fn, ) kwargs = self.dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input) self.assertTrue(token_dispatch_input.quant.dispatch_with_quant) self.assertEqual(kwargs["quant_mode"], 4) self.assertEqual(kwargs["y_dtype"], torch.float8_e4m3fn) def test_get_dispatch_mc2_kwargs_with_mxfp4_quant(self): hidden_states = torch.randn(10, 128) topk_weights = torch.randn(10, 1) topk_ids = torch.randint(0, 8, (10, 1)) expert_map = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7]) self.dispatcher.a5_need_extra_args = True token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, quant_type=QuantType.MXFP4, act_quant_type=MXFP4_TEST_DTYPE, ) kwargs = self.dispatcher.get_dispatch_mc2_kwargs(token_dispatch_input) self.assertTrue(token_dispatch_input.quant.dispatch_with_quant) self.assertEqual(kwargs["quant_mode"], 4) self.assertIn("y_dtype", kwargs) self.assertNotEqual(kwargs["y_dtype"], torch.float8_e4m3fn) with patch( "torch_npu.npu_moe_distribute_dispatch_v2", return_value=( torch.randn(10, 128), torch.randn(10, 1), torch.arange(10, dtype=torch.int32), torch.tensor([10], dtype=torch.int64), torch.tensor([10], dtype=torch.int64), torch.tensor([10], dtype=torch.int64), torch.randn(10, 1), ), ) as mock_dispatch: output = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input) mock_dispatch.assert_called_once() self.assertIsNotNone(output.dynamic_scale) self.assertTrue(output.combine_metadata.quant.dispatch_with_quant) def test_allgather_token_dispatch_quant_mode_without_dynamic_scale(): dispatcher = TokenDispatcherWithAllGather(top_k=2, num_experts=128) hidden_states = torch.randn(3, 128) topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]]) topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]], dtype=torch.int32) init_routing_output = ( torch.randn(6, 128), torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.int32), torch.tensor([2, 2, 2], dtype=torch.int32), torch.randn(6, 4), ) cases = [ { "quant_type": QuantType.MXFP8, "act_quant_type": torch.float8_e4m3fn, "expected_quant_mode": 3, "expected_act_quant_type": torch.float8_e4m3fn, "expect_dynamic_scale": True, }, { "quant_type": QuantType.MXFP4, "act_quant_type": MXFP4_TEST_DTYPE, "expected_quant_mode": -1, "expected_act_quant_type": None, "expect_dynamic_scale": False, }, ] for case in cases: token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, quant_type=case["quant_type"], act_quant_type=case["act_quant_type"], ) with patch( "vllm_ascend.ops.fused_moe.token_dispatcher.DeviceOperator.npu_moe_init_routing", return_value=init_routing_output, ) as mock_init_routing: output = dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input) init_kwargs = mock_init_routing.call_args.kwargs assert init_kwargs["quant_mode"] == case["expected_quant_mode"] assert init_kwargs["act_quant_type"] == case["expected_act_quant_type"] assert (output.dynamic_scale is not None) == case["expect_dynamic_scale"] def test_allgather_token_dispatch_mxfp4_keeps_prequantized_scale(): dispatcher = TokenDispatcherWithAllGather(top_k=2, num_experts=128) hidden_states = torch.randn(3, 128) topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]]) topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]], dtype=torch.int32) pertoken_scale = torch.randn(3, 4) returned_scale = torch.randn(6, 4) init_routing_output = ( torch.randn(6, 128), torch.tensor([0, 1, 2, 3, 4, 5], dtype=torch.int32), torch.tensor([2, 2, 2], dtype=torch.int32), returned_scale, ) token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, pertoken_scale=pertoken_scale, quant_type=QuantType.MXFP4, act_quant_type=MXFP4_TEST_DTYPE, ) with patch( "vllm_ascend.ops.fused_moe.token_dispatcher.DeviceOperator.npu_moe_init_routing", return_value=init_routing_output, ) as mock_init_routing: output = dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input) init_kwargs = mock_init_routing.call_args.kwargs assert init_kwargs["scale"] is pertoken_scale assert init_kwargs["quant_mode"] == -1 assert init_kwargs["act_quant_type"] == MXFP4_TEST_DTYPE assert output.dynamic_scale is returned_scale class TestTokenDispatcherWithAllGather(TestBase): def setUp(self): # Mock dependencies kwargs = { "apply_router_weight_on_input": False, "top_k": 2, "max_num_tokens": 100, "ep_size": 2, "num_experts": 128, "with_quant": False, } self.dispatcher = TokenDispatcherWithAllGather(**kwargs) # Mock NPU functions self.patcher_npu_moe_init_routing_custom = patch("torch.ops._C_ascend.npu_moe_init_routing_custom") self.mock_npu_moe_init_routing_custom = self.patcher_npu_moe_init_routing_custom.start() self.mock_npu_moe_init_routing_custom.return_value = ( torch.randn(6, 128), # sorted_hidden_states torch.tensor([0, 1, 2, 3, 4, 5]), # expanded_row_idx torch.tensor([0, 1, 0, 1, 0, 1]), # expanded_expert_idx torch.tensor([0, 1, 0, 1, 0, 1]), ) self.patcher_npu_moe_token_unpermute = patch("torch_npu.npu_moe_token_unpermute") self.mock_npu_moe_token_unpermute = self.patcher_npu_moe_token_unpermute.start() self.mock_npu_moe_token_unpermute.return_value = torch.randn(6, 128) def tearDown(self): self.patcher_npu_moe_init_routing_custom.stop() self.patcher_npu_moe_token_unpermute.stop() @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.") def test_token_dispatch_without_expert_map(self): hidden_states = torch.randn(3, 128) topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]]) topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]]) token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, ) results = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input) # Verify npu_moe_init_routing is called self.mock_npu_moe_init_routing_custom.assert_called_once() args, kwargs = self.mock_npu_moe_init_routing_custom.call_args self.assertEqual(results.group_list_type, 1) self.assertIsInstance(results.combine_metadata, MoEAllGatherCombineMetadata) @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.") def test_token_dispatch_with_expert_map(self): self.dispatcher.expert_map = torch.tensor([0, 1, 2, 3]) hidden_states = torch.randn(3, 128) topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]]) topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]]) token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, ) results = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input) # Verify npu_moe_init_routing is called self.mock_npu_moe_init_routing_custom.assert_called_once() args, kwargs = self.mock_npu_moe_init_routing_custom.call_args self.assertEqual(results.group_list_type, 1) self.assertIsInstance(results.combine_metadata, MoEAllGatherCombineMetadata) @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.") def test_token_dispatch_without_quant(self): kwargs = { "apply_router_weight_on_input": False, "top_k": 2, "max_num_tokens": 100, "ep_size": 2, "num_experts": 128, } self.dispatcher_quant = TokenDispatcherWithAllGather(**kwargs) hidden_states = torch.randn(3, 128) topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]]) topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]]) token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, ) results = self.dispatcher_quant.token_dispatch(token_dispatch_input=token_dispatch_input) self.assertEqual(results.group_list_type, 1) @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.") def test_token_dispatch_with_quant(self): kwargs = { "apply_router_weight_on_input": False, "top_k": 2, "max_num_tokens": 100, "ep_size": 2, "num_experts": 128, } self.dispatcher_quant = TokenDispatcherWithAllGather(**kwargs) hidden_states = torch.randn(3, 128) topk_weights = torch.tensor([[0.7, 0.3], [0.6, 0.4], [0.5, 0.5]]) topk_ids = torch.tensor([[0, 1], [1, 2], [2, 3]]) token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, quant_type=QuantType.W8A8, ) results = self.dispatcher_quant.token_dispatch(token_dispatch_input=token_dispatch_input) self.assertIsNotNone(results.hidden_states) self.assertIsNotNone(results.group_list) self.assertIsNotNone(results.dynamic_scale) self.assertEqual(results.group_list_type, 1) @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.") def test_token_combine_with_expert_map(self): hidden_states = torch.randn(6, 128) combine_metadata = MoEAllGatherCombineMetadata( expanded_row_idx=torch.tensor([0, 1, 1, 1, 1, 1]), topk_weights=torch.tensor([0.5, 0.5, 0.5, 0.5, 0.5, 0.5]), restore_shape=torch.Size([6, 128]), ) final_hidden_states = self.dispatcher.token_combine(hidden_states, combine_metadata) self.assertEqual(final_hidden_states.shape, (6, 128)) @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.") def test_token_combine_without_expert_map(self): hidden_states = torch.randn(6, 128) combine_metadata = MoEAllGatherCombineMetadata( expanded_row_idx=torch.tensor([0, 1, 1, 1, 1, 1]), topk_weights=torch.tensor([0.5, 0.5, 0.5, 0.5, 0.5, 0.5]), restore_shape=torch.Size([6, 128]), ) final_hidden_states = self.dispatcher.token_combine(hidden_states, combine_metadata) self.mock_npu_moe_token_unpermute.assert_called_once() self.assertEqual(final_hidden_states.shape, (6, 128)) @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.") def test_token_dispatch_with_router_weight(self): hidden_states = torch.randn(3, 128) topk_weights = torch.tensor([[0.7], [0.6], [0.5]]) # topk=1 topk_ids = torch.tensor([[0], [1], [2]]) token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, apply_router_weight_on_input=True, ) results = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input) self.assertEqual(results.hidden_states.shape, (6, 128)) self.assertIsInstance(results.combine_metadata, MoEAllGatherCombineMetadata) class TestTokenDispatcherWithAll2AllV(TestBase): def setUp(self): # Patch properties patcher1 = patch.object( TokenDispatcherWithAll2AllV, "ep_group", new_callable=PropertyMock, return_value=MagicMock() ) patcher2 = patch.object(TokenDispatcherWithAll2AllV, "ep_rank", new_callable=PropertyMock, return_value=0) patcher3 = patch.object(TokenDispatcherWithAll2AllV, "ep_size", new_callable=PropertyMock, return_value=2) self.addCleanup(patcher1.stop) self.addCleanup(patcher2.stop) self.addCleanup(patcher3.stop) self.mock_ep_group_prop = patcher1.start() self.mock_ep_rank_prop = patcher2.start() self.mock_ep_size_prop = patcher3.start() # Mock torch_npu.npu_moe_token_permute patcher4 = patch("torch_npu.npu_moe_token_permute") self.mock_npu_moe_token_permute = patcher4.start() self.addCleanup(patcher4.stop) self.mock_npu_moe_token_permute.return_value = (torch.randn(16, 16), torch.arange(16)) # Mock torch_npu.npu_moe_token_unpermute patcher5 = patch("torch_npu.npu_moe_token_unpermute") self.mock_npu_moe_token_unpermute = patcher5.start() self.addCleanup(patcher5.stop) self.mock_npu_moe_token_unpermute.return_value = torch.randn(8, 16) # Mock async_all_to_all patcher6 = patch("vllm_ascend.ops.fused_moe.comm_utils.async_all_to_all") self.mock_async_all_to_all = patcher6.start() self.addCleanup(patcher6.stop) self.mock_async_all_to_all.return_value = (None, torch.randn(16, 16), MagicMock()) # Mock gather_from_sequence_parallel_region patcher7 = patch("vllm_ascend.ops.fused_moe.token_dispatcher.gather_from_sequence_parallel_region") self.mock_gather_from_sequence_parallel_region = patcher7.start() self.addCleanup(patcher7.stop) self.mock_gather_from_sequence_parallel_region.return_value = torch.tensor( [[2, 2, 2, 2], [2, 2, 2, 2]], dtype=torch.int64 ) # Mock torch.histc patcher8 = patch("torch.histc") self.mock_histc = patcher8.start() self.addCleanup(patcher8.stop) self.mock_histc.return_value = torch.tensor([2, 2, 2, 2], dtype=torch.int64) # Mock torch.npu.current_device patcher9 = patch("torch.npu.current_device") self.mock_current_device = patcher9.start() self.addCleanup(patcher9.stop) self.mock_current_device.return_value = "cpu" # Mock torch_npu.npu_dynamic_quant patcher10 = patch("torch_npu.npu_dynamic_quant") self.mock_npu_dynamic_quant = patcher10.start() self.addCleanup(patcher10.stop) self.mock_npu_dynamic_quant.return_value = (torch.randn(16, 16), torch.randn(16)) # Mock torch.ops._C_ascend.npu_moe_init_routing_custom patcher11 = patch("torch.ops._C_ascend.npu_moe_init_routing_custom") self.mock_npu_moe_init_routing_custom = patcher11.start() self.addCleanup(patcher11.stop) self.mock_npu_moe_init_routing_custom.return_value = ( torch.randn(16, 16), torch.arange(16), None, torch.randn(16), ) # Mock torch.repeat_interleave patcher12 = patch("torch.repeat_interleave") self.mock_repeat_interleave = patcher12.start() self.addCleanup(patcher12.stop) self.mock_repeat_interleave.return_value = torch.arange(16) self.dispatcher = TokenDispatcherWithAll2AllV(top_k=2, num_experts=4, num_local_experts=2, with_quant=False) @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.") def test_token_dispatch(self): hidden_states = torch.randn(8, 16) topk_weights = torch.rand(8, 4) topk_ids = torch.randint(0, 4, (8, 2)).long() expert_map = torch.tensor([0, 1, 2, 3]) self.dispatcher.expert_ids_per_ep_rank = torch.tensor([0, 1], dtype=torch.int32) self.dispatcher.local_expert_indices = [0, 1] token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, ) result = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input) self.assertIsNotNone(result.hidden_states) self.assertIsNotNone(result.group_list) self.assertEqual(result.group_list_type, 1) self.assertIsInstance(result.combine_metadata, MoEAllToAllCombineMetadata) @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.") def test_token_combine(self): hidden_states = torch.randn(16, 16) combine_metadata = MoEAllToAllCombineMetadata( input_splits=np.array([4, 4]), output_splits=np.array([4, 4]), topk_weights=torch.rand(8, 4), reversed_local_input_permutation_mapping=torch.arange(8), reversed_global_input_permutation_mapping=torch.arange(16), hidden_shape=torch.Size([8, 16]), hidden_shape_before_permute=torch.Size([8, 16]), ) self.dispatcher.expert_ids_per_ep_rank = torch.tensor([0, 1], dtype=torch.int32) self.dispatcher.local_expert_indices = [0, 1] output = self.dispatcher.token_combine(hidden_states, combine_metadata) self.assertIsNotNone(output) self.assertEqual(output.shape, (8, 16)) @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.") def test_token_dispatch_with_quant(self): self.dispatcher = TokenDispatcherWithAll2AllV(top_k=2, num_experts=4, num_local_experts=2) hidden_states = torch.randn(8, 16) topk_weights = torch.rand(8, 4) topk_ids = torch.randint(0, 4, (8, 2)).long() expert_map = torch.tensor([0, 1, 2, 3]) self.dispatcher.expert_ids_per_ep_rank = torch.tensor([0, 1], dtype=torch.int32) self.dispatcher.local_expert_indices = [0, 1] token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, quant_type=QuantType.W8A8, ) result = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input) self.assertIsNotNone(result.hidden_states) self.assertIsNotNone(result.group_list) self.assertIsNotNone(result.dynamic_scale) self.assertEqual(result.group_list_type, 1) self.assertIsInstance(result.combine_metadata, MoEAllToAllCombineMetadata) @pytest.mark.skip("Skip as register_kernels has NPU SocName checking in CANN 8.5.0.") def test_token_dispatch_with_quant_no_active_tokens(self): self.dispatcher = TokenDispatcherWithAll2AllV(top_k=2, num_experts=4, num_local_experts=2) self.mock_repeat_interleave.return_value = torch.tensor([], dtype=torch.long) hidden_states = torch.randn(8, 16) topk_weights = torch.rand(8, 4) topk_ids = torch.randint(0, 4, (8, 2)).long() expert_map = torch.tensor([0, 1, 2, 3]) self.dispatcher.expert_ids_per_ep_rank = torch.tensor([0, 1], dtype=torch.int32) self.dispatcher.local_expert_indices = [0, 1] token_dispatch_input = build_token_dispatch_input_fixture( hidden_states=hidden_states, topk_weights=topk_weights, topk_ids=topk_ids, expert_map=expert_map, quant_type=QuantType.W8A8, ) result = self.dispatcher.token_dispatch(token_dispatch_input=token_dispatch_input) self.assertIsNotNone(result.hidden_states) self.assertIsNotNone(result.group_list) self.assertIsNotNone(result.dynamic_scale) self.assertEqual(result.group_list_type, 1)