correctly use multistream in acl graph

Signed-off-by: daishixun <[email protected]>

Author: dsxsteven

vllm_ascend/ops/fused_moe/fused_moe.py

@@ -392,11 +392,14 @@ def forward_impl(self, hidden_states: torch.Tensor,
         if isinstance(final_hidden_states, tuple):
             final_hidden_states, group_list_type, expert_tokens = final_hidden_states
             if self.dynamic_eplb:
-                with npu_stream_switch(moe_load_async_stream()):
-                    moe_load_async_stream().wait_stream(
-                        torch.npu.current_stream(device=expert_tokens.device))
+                moe_load_stream = moe_load_async_stream()
+                cur_stream = torch.npu.current_stream()
+
+                moe_load_stream.wait_stream(cur_stream)
+                with npu_stream_switch(moe_load_stream):
                     self.moe_load += expert_tokens if group_list_type == 1 else \
                         torch.cat([expert_tokens[:1], expert_tokens[1:] - expert_tokens[:-1]])
+                cur_stream.wait_stream(moe_load_stream)
 
         final_hidden_states = forward_context.moe_comm_method.finalize(
             hidden_states=final_hidden_states,