qwen3_moe/qwen25 support torchair graph

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

Author: NicholasTao

tests/ut/test_ascend_config.py

@@ -232,7 +232,7 @@ def test_check_ascend_config_wrong_case(self):
 
     def test_check_torchair_supported(self):
         test_cases = [('deepseek_v3', True), ('PanguProMoE', True),
-                      ('qwen', False), ('llama', False)]
+                      ('qwen', True), ('llama', False)]
         for model_type, expected_output in test_cases:
             self.assertEqual(_check_torchair_supported(model_type),
                              expected_output)

vllm_ascend/ascend_config.py

@@ -17,7 +17,7 @@
 
 from vllm.logger import logger
 
-TORCHAIR_MODEL_LIST = ["deepseek", "pangu", "kimi_k2"]
+TORCHAIR_MODEL_LIST = ["deepseek", "pangu", "kimi_k2", "qwen"]
 
 
 def _check_torchair_supported(model_type: str):
@@ -162,7 +162,7 @@ def check_ascend_config(vllm_config, enforce_eager):
     else:
         # torchair_graph case
         if ascend_config.torchair_graph_config.enabled:
-            # torchair_graph is supported for deepseek/pangu model only.
+            # torchair_graph is supported for deepseek/pangu/qwen model only.
             if vllm_config.model_config:
                 model_type = vllm_config.model_config.hf_config.model_type
                 if not _check_torchair_supported(model_type):

vllm_ascend/ops/rotary_embedding.py

@@ -19,6 +19,8 @@
 from typing import Optional, Tuple
 
 import torch
+import torch.nn.functional as F
+import torch_npu
 from vllm.model_executor.layers.rotary_embedding import (
     DeepseekScalingRotaryEmbedding, RotaryEmbedding)
 
@@ -37,17 +39,18 @@ def rope_forward_oot(
     query: torch.Tensor,
     key: torch.Tensor,
     offsets: Optional[torch.Tensor] = None,
-    is_neox_style_override: Optional[bool] = None
+    is_neox_style_override: Optional[bool] = None,
+    is_qwen_torchair: Optional[bool] = False,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
-    if get_ascend_config().torchair_graph_config.enabled:
+    if get_ascend_config(
+    ).torchair_graph_config.enabled and not is_qwen_torchair:
         return self.forward_native(
             positions,
             query,
             key,
             offsets,
         )
 
-    import torch_npu
     query_shape, key_shape = query.shape, key.shape
     if self.cos_sin_cache.device != query.device:
         self.cos_sin_cache = self.cos_sin_cache.to(query.device)
@@ -246,6 +249,98 @@ def _set_cos_sin_cache(self, seq_len, device, dtype):
     self.register_buffer("sin_cached", sin_cached, persistent=False)
 
 
+def __set_cos_sin_cache(self, seq_len, device, dtype):
+    inv_freq = 1.0 / (self.base**(torch.arange(
+        0, self.rotary_dim, 2, device=device, dtype=torch.float32) *
+                                  (1 / self.rotary_dim)))
+    self.register_buffer("inv_freq", inv_freq)
+
+    t = torch.arange(self.max_position_embeddings,
+                     device=self.inv_freq.device,
+                     dtype=torch.float32)
+    freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+
+    emb = torch.cat((freqs, freqs), dim=-1)
+    self.register_buffer("cos", emb.cos().to(dtype=dtype), persistent=False)
+    self.register_buffer("sin", emb.sin().to(dtype=dtype), persistent=False)
+    self.embed = F.embedding
+
+
+def qwen_rope_init_func(
+    self,
+    head_size: int,
+    rotary_dim: int,
+    max_position_embeddings: int,
+    base: float,
+    is_neox_style: bool,
+    dtype: torch.dtype,
+) -> None:
+    super(RotaryEmbedding, self).__init__()
+    self.head_size = head_size
+    self.rotary_dim = rotary_dim
+    self.max_position_embeddings = max_position_embeddings
+    self.base = base
+    self.is_neox_style = is_neox_style
+    self.dtype = dtype
+
+    cache = self._compute_cos_sin_cache()
+    cache = cache.to(dtype)
+    self.cos_sin_cache: torch.Tensor  # type: ignore[misc]
+    self.register_buffer("cos_sin_cache", cache, persistent=False)
+    if get_ascend_config().torchair_graph_config.enabled:
+        __set_cos_sin_cache(self,
+                            seq_len=max_position_embeddings,
+                            device="npu",
+                            dtype=dtype)
+
+
+def rope_forward(
+    self,
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    offsets: Optional[torch.Tensor] = None,
+    is_neox_style_override: Optional[bool] = None,
+    max_seq_len: Optional[int] = None,
+    is_prefill: Optional[bool] = True,
+    is_qwen_torchair: Optional[bool] = False,
+):
+    if (not get_ascend_config().torchair_graph_config.enabled
+            or not is_qwen_torchair or is_prefill):
+        return rope_forward_oot(self, positions, query, key, offsets,
+                                is_neox_style_override,
+                                is_qwen_torchair)  # type: ignore
+
+    if max_seq_len is not None and torch.gt(max_seq_len,
+                                            self.max_position_embeddings):
+        __set_cos_sin_cache(self,
+                            seq_len=max_seq_len,
+                            device=query.device,
+                            dtype=torch.float32)
+
+    # bsnd/bnsd
+    if positions is not None:
+        cos = self.embed(positions, self.cos)
+        sin = self.embed(positions, self.sin)
+        self.cos_embed = cos
+        self.sin_embed = sin
+    else:
+        cos = self.cos_embed
+        sin = self.sin_embed
+
+    query = query.view(*query.shape[:-1], -1, self.head_size).contiguous()
+    key = key.view(*key.shape[:-1], -1, self.head_size).contiguous()
+
+    cos = cos.unsqueeze(-2).unsqueeze(-2)
+    sin = sin.unsqueeze(-2).unsqueeze(-2)
+
+    query = query.unsqueeze(1)
+    key = key.unsqueeze(1)
+
+    q_embed, k_embed = torch_npu.npu_apply_rotary_pos_emb(query, key, cos, sin)
+    return q_embed.flatten(-2), k_embed.flatten(-2)
+
+
 def deepseek_rope_init_func(
     self,
     head_size: int,
@@ -283,7 +378,8 @@ def deepseek_rope_init_func(
                        device="npu")
 
 
-RotaryEmbedding.forward_oot = rope_forward_oot
+RotaryEmbedding.__init__ = qwen_rope_init_func
+RotaryEmbedding.forward_oot = rope_forward
 
 # Note: we adopt the native huggingface deepseek rope initialization code from
 # https://huggingface.co/deepseek-ai/DeepSeek-V3-0324/blob/main/modeling_deepseek.py for