diff --git a/vllm_ascend/ops/fused_moe/fused_moe.py b/vllm_ascend/ops/fused_moe/fused_moe.py index f94871fc..12d6a63e 100644 --- a/vllm_ascend/ops/fused_moe/fused_moe.py +++ b/vllm_ascend/ops/fused_moe/fused_moe.py @@ -382,9 +382,9 @@ class AscendFusedMoE(FusedMoE): group_list_type = fused_experts_results.group_list_type assert expert_tokens is not None and group_list_type is not None, \ "expert_tokens and group_list_type should not be None when dynamic_eplb is enabled." - self.moe_load += expert_tokens if group_list_type == 1 else \ + local_load = expert_tokens if group_list_type == 1 else \ torch.cat([expert_tokens[:1], expert_tokens[1:] - expert_tokens[:-1]]) - + self.moe_load.add_(local_load) routed_out = forward_context.moe_comm_method.finalize( hidden_states=fused_experts_results.routed_out, reduce_results=self.reduce_results,