diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py index 62fee47852e9..8fab75d77e26 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py @@ -966,6 +966,7 @@ def forward_fake( output_dtype: Optional[torch.dtype] = None, all_rank_num_tokens: Optional[List[int]] = None, use_dp_padding: Optional[bool] = None, + alltoall_result_do_sum: bool = True, **kwargs, ) -> Union[torch.Tensor, List[torch.Tensor]]: moe_output = super().forward_fake( @@ -976,7 +977,7 @@ def forward_fake( all_rank_num_tokens=all_rank_num_tokens, use_dp_padding=use_dp_padding, **kwargs) - if self.alltoall_method_type == AlltoallMethodType.MNNVL: + if self.alltoall_method_type == AlltoallMethodType.NVLinkTwoSided and not alltoall_result_do_sum: shape = moe_output.shape top_k = self.routing_method.experts_per_token new_shape = [shape[0], top_k, shape[1]] diff --git a/tests/unittest/_torch/modules/test_fused_moe.py b/tests/unittest/_torch/modules/test_fused_moe.py index 14acccf452f6..cea8ff57e299 100644 --- a/tests/unittest/_torch/modules/test_fused_moe.py +++ b/tests/unittest/_torch/modules/test_fused_moe.py @@ -335,10 +335,6 @@ def per_rank_test_fused_moe_alltoall(job_id): ], ids=lambda s: s.name) def test_fused_moe_alltoall_fp4(alltoall_method_type): - - if alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: - pytest.skip("Skipped due to https://nvbugs/5467531") - world_size = 4 dtype = torch.bfloat16 HIDDEN_SIZE = 4096