@@ -383,25 +383,22 @@ def test_apply_without_expert_map(self, moe_method, mock_dist_env,
383383 else :
384384 assert result .shape == x .shape
385385
386- @patch ('torch_npu.npu_moe_gating_top_k' )
387386 @pytest .mark .parametrize ("others_param" , [16 , 1 , 4 ])
388387 def test_apply_with_expert_map (self , moe_method , mock_dist_env ,
389- mock_moe_env , others_param , mock_topk ):
388+ mock_moe_env , others_param ):
390389 """
391390 1 test use_select_experts and use fused_expters_with_mc2
392391 2 test use_select_experts and fused_experts_with_all2all_buffer
393392 3 test use_select_experts and fused_experts_with_all2all
394393 4 test use_select_experts and fused_experts
395394 """
396- mock_topk .return_value = (torch .randn (8 ,
397- 2 ), torch .randint (0 , 8 ,
398- (8 , 2 )), None )
399395 ep_size = others_param
400396 is_prefill = False
401397 forward_context = MagicMock (
402398 fused_moe_state = _get_fused_moe_state (ep_size , is_prefill , True ))
403399 with patch ("vllm_ascend.torchair.ops.torchair_fused_moe.get_forward_context" , return_value = forward_context ), \
404- patch ("vllm_ascend.torchair.ops.torchair_fused_moe.get_ascend_soc_version" , return_value = AscendSocVersion .A3 ):
400+ patch ("vllm_ascend.torchair.ops.torchair_fused_moe.get_ascend_soc_version" , return_value = AscendSocVersion .A3 ), \
401+ patch ('torch_npu.npu_moe_gating_top_k' , return_value = (torch .randn (8 , 2 ),torch .randint (0 , 8 , (8 , 2 )),None )):
405402 expert_map = torch .tensor ([0 , 1 , 2 , - 1 , - 1 , - 1 , - 1 , - 1 ])
406403 moe_method .ep_size = ep_size
407404 x = torch .randn (8 , 2 , 2 )
0 commit comments