init qwen torchair graph mode

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

Author: NicholasTao

tests/e2e/multicard/test_torchair_graph_mode.py

@@ -162,3 +162,65 @@ def test_e2e_pangu_with_torchair():
         },
     }
     _pangu_torchair_test_fixture(additional_config)
+
+
+def _qwen_torchair_test_fixture(
+    model,
+    tp,
+    enable_expert_parallel,
+):
+    # The current access control does not support 16 cards,
+    # so the MC2 operator in Qwen's graph mode cannot run.
+    # Once 16-card support is available,
+    # this e2e can be switched to graph mode.
+    example_prompts = [
+        "Hello, my name is",
+        "The president of the United States is",
+        "The capital of France is",
+        "The future of AI is",
+    ]
+
+    additional_config = {
+        "torchair_graph_config": {
+            "enabled": False,
+        },
+        "ascend_scheduler_config": {
+            "enabled": True,
+        },
+        "refresh": True,
+    }
+
+    with VllmRunner(
+            model,
+            dtype="half",
+            tensor_parallel_size=tp,
+            distributed_executor_backend="mp",
+            enforce_eager=True,
+            additional_config=additional_config,
+            enable_expert_parallel=enable_expert_parallel,
+    ) as vllm_model:
+        # use greedy sampler to make sure the generated results are fix
+        vllm_output = vllm_model.generate_greedy(example_prompts, 5)
+
+    # NOTE: vllm-ascend/pangu-pro-moe-pruing is only part of PanguProMoE
+    # with 2 hidden layers, thus the golden results seems inaccurate.
+    # This will only change if accuracy changes with the official weights
+    # of PanguProMoE.
+    golden_results = [
+        'Hello, my name is Remempondeprecatedmiot忱',
+        'The president of the United States is Remem下的一个 rever ceremoni Segnali',
+        'The capital of France is Rememvoud administrativ Remem投',
+        'The future of AI isotope Segnali Zoeken精细化 supus',
+    ]
+
+    assert len(golden_results) == len(vllm_output)
+    for i in range(len(vllm_output)):
+        print(f"Generated text: {vllm_output[i][1]!r}")
+
+
+def test_e2e_qwen2_with_torchair():
+    _qwen_torchair_test_fixture("Qwen/Qwen2.5-0.5B-Instruct", 2, False)
+
+
+def test_e2e_qwen3_moe_with_torchair():
+    _qwen_torchair_test_fixture("Qwen/Qwen3-30B-A3B", 2, True)

tests/ut/models/test_qwen3_moe.py

@@ -12,11 +12,15 @@
 # limitations under the License.
 # This file is a part of the vllm-ascend project.
 #
+import math
+import unittest
 
 import pytest
+import torch
 from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM
 
 from vllm_ascend.models.qwen3_moe import CustomQwen3MoeForCausalLM
+from vllm_ascend.torchair.models.qwen3_moe import CustomQwen3MoeAttention
 
 
 class TestCustomQwen3MoeForCausalLM:
@@ -44,3 +48,51 @@ def test_packed_modules_mapping_structure(self):
             ]
         }
         assert CustomQwen3MoeForCausalLM.packed_modules_mapping == expected_mapping
+
+
+class DummyRMSNorm:
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        self.dim = dim
+        self.eps = eps
+
+    def __call__(self, x):
+        mean_sq = x.pow(2).mean(dim=-1, keepdim=True)
+        denom = (mean_sq + self.eps).sqrt()
+        return x / denom
+
+
+class TestCustomQwen3MoeAttention(unittest.TestCase):
+
+    def setUp(self):
+        self.batch = 2
+        self.seq_len = 3
+        self.q_size = 8
+        self.kv_size = 8
+        self.head_dim = 4
+        self.rms_eps = 1e-6
+
+        total_dim = self.q_size + 2 * self.kv_size
+
+        self.qkv = torch.arange(self.batch * self.seq_len * total_dim,
+                                dtype=torch.float32).reshape(
+                                    self.batch, self.seq_len, total_dim)
+
+    def test_constant_input_normalization(self):
+        ones_qkv = torch.ones((1, 1, self.q_size + 2 * self.kv_size),
+                              dtype=torch.float32)
+
+        q_norm = DummyRMSNorm(self.head_dim, self.rms_eps)
+        k_norm = DummyRMSNorm(self.head_dim, self.rms_eps)
+        q, k, v = CustomQwen3MoeAttention.normalize_qkv(
+            ones_qkv, self.q_size, self.kv_size, self.head_dim, q_norm, k_norm)
+
+        norm_val = 1.0 / math.sqrt(1.0 + self.rms_eps)
+
+        expected_q = torch.full((1, 1, self.q_size), norm_val)
+        expected_k = torch.full((1, 1, self.kv_size), norm_val)
+        expected_v = torch.ones((1, 1, self.kv_size), dtype=torch.float32)
+
+        self.assertTrue(torch.allclose(q, expected_q, atol=1e-6))
+        self.assertTrue(torch.allclose(k, expected_k, atol=1e-6))
+        self.assertTrue(torch.equal(v, expected_v))

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):