fix ut

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

Author: NicholasTao

tests/ut/models/test_qwen3_moe.py

@@ -50,6 +50,18 @@ 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):
@@ -70,8 +82,10 @@ 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, self.rms_eps)
+            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)
 

tests/ut/ops/test_rotary_embedding.py

@@ -4,8 +4,9 @@
 import torch
 
 from tests.ut.base import TestBase
-from vllm_ascend.ops.rotary_embedding import (__set_cos_sin_cache,
-                                              custom_rotary_embedding_enabled,
+from vllm_ascend.ops.rotary_embedding import \
+    __set_cos_sin_cache as raw__set_cos_sin_cache
+from vllm_ascend.ops.rotary_embedding import (custom_rotary_embedding_enabled,
                                               native_rope_deepseek_forward,
                                               rope_forward_oot, rotate_half,
                                               yarn_find_correction_dim,
@@ -327,7 +328,7 @@ def __init__(self, base, rotary_dim, max_position_embeddings):
         self.max_position_embeddings = max_position_embeddings
 
     def _set_cos_sin_cache(self, seq_len, device, dtype):
-        return __set_cos_sin_cache(self, seq_len, device, dtype)
+        return raw__set_cos_sin_cache(self, seq_len, device, dtype)
 
 
 class TestSetCosSinCache(TestBase):

vllm_ascend/models/qwen3_moe.py

@@ -207,15 +207,15 @@ def __init__(
 
     @staticmethod
     def normalize_qkv(qkv: torch.Tensor, q_size: int, kv_size: int,
-                      head_dim: int, rms_norm_eps: float):
+                      head_dim: int, q_norm, k_norm):
         q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
 
         q_by_head = q.view(*q.shape[:-1], q.shape[-1] // head_dim, head_dim)
-        q_by_head = RMSNorm(head_dim, eps=rms_norm_eps)(q_by_head)
+        q_by_head = q_norm(q_by_head)
         q = q_by_head.view(q.shape)
 
         k_by_head = k.view(*k.shape[:-1], k.shape[-1] // head_dim, head_dim)
-        k_by_head = RMSNorm(head_dim, eps=rms_norm_eps)(k_by_head)
+        k_by_head = k_norm(k_by_head)
         k = k_by_head.view(k.shape)
 
         return q, k, v
@@ -228,7 +228,7 @@ def forward(
             attn_metadata: Optional[AttentionMetadata] = None) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = self.normalize_qkv(qkv, self.q_size, self.kv_size,
-                                     self.head_dim, self.rms_norm_eps)
+                                     self.head_dim, self.q_norm, self.k_norm)
 
         if (self.torchair_graph_enabled and attn_metadata is not None and
                 attn_metadata.attn_state == AscendAttentionState.DecodeOnly):