[Feature] qwen3_moe qk norm support multi_stream.

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>

Author: weijinqian_v1

vllm_ascend/models/qwen3.py

@@ -28,6 +28,7 @@
 
 from vllm_ascend import envs
 from vllm_ascend.ops.layernorm import AddRMSNormW8A8Quant
+from vllm_ascend.utils import npu_stream_switch
 
 
 def pad(tensor, x):
@@ -130,6 +131,7 @@ def __init__(self,
         self.tp_size = get_tensor_model_parallel_world_size()
         self.tp_rank = get_tensor_model_parallel_rank()
         self.enable_fc = envs.VLLM_ASCEND_ENABLE_FLASHCOMM
+        self.alt_stream = torch.npu.Stream()
         if self.enable_fc == 2:
             self.o_proj = ReplicatedLinear(
                 self.total_num_heads * self.head_dim,
@@ -156,15 +158,17 @@ def forward(
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        # Add qk-norm
-        q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim,
-                           self.head_dim)
-        q_by_head = self.q_norm(q_by_head)
-        q = q_by_head.view(q.shape)
+        with npu_stream_switch(self.alt_stream):
+            # Add qk-norm
+            q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim,
+                               self.head_dim)
+            q_by_head = self.q_norm(q_by_head)
+            q = q_by_head.view(q.shape)
         k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim,
                            self.head_dim)
         k_by_head = self.k_norm(k_by_head)
         k = k_by_head.view(k.shape)
+        torch.npu.current_stream().wait_stream(self.alt_stream)
         q, k = self.rotary_emb(positions,
                                q,
                                k,

vllm_ascend/utils.py

@@ -454,3 +454,16 @@ def delete_torchair_cache_file():
         shutil.rmtree(torch_air_abs_path)
     except FileNotFoundError:
         pass
+
+
+def npu_stream_switch(target_stream: torch.npu.Stream,
+                      *,
+                      enabled: bool = True):
+    """
+    Switch to the target stream if enabled is True.
+    Otherwise, do nothing.
+    """
+    if not enabled:
+        return nullcontext()
+    assert target_stream is not None
+    return torch.npu.stream(target_stream)