refacotr util registry

Signed-off-by: Icey <1790571317@qq.com>

Author: wxsIcey

tests/ut/test_utils.py

@@ -24,6 +24,7 @@
 
 from tests.ut.base import TestBase
 from vllm_ascend import utils
+from vllm_ascend.utils import REGISTERED_ASCEND_OPS
 
 
 class TestUtils(TestBase):
@@ -302,14 +303,14 @@ def test_register_ascend_customop(self, mock_ascend_rmsnorm,
 
         # ascend custom op is not registered
         utils.register_ascend_customop()
-        # should call register_oot three
-        self.assertEqual(mock_customop.register_oot.call_count, 14)
+        self.assertEqual(mock_customop.register_oot.call_count,
+                         len(REGISTERED_ASCEND_OPS))
         self.assertTrue(utils._ASCEND_CUSTOMOP_IS_REIGISTERED)
 
         # ascend custom op is already registered
         utils.register_ascend_customop()
-        # should not register_oot again, thus only called three in this ut
-        self.assertEqual(mock_customop.register_oot.call_count, 14)
+        self.assertEqual(mock_customop.register_oot.call_count,
+                         len(REGISTERED_ASCEND_OPS))
 
 
 class TestProfileExecuteDuration(TestBase):

vllm_ascend/utils.py

@@ -479,6 +479,37 @@ def get_all_reduce_merge_state(ep_size: int, is_deepseek_v3_r1: bool):
     return False
 
 
+from vllm_ascend.models.layers.mla import AscendMultiHeadLatentAttention
+from vllm_ascend.ops.activation import AscendQuickGELU, AscendSiluAndMul
+from vllm_ascend.ops.common_fused_moe import AscendFusedMoE
+from vllm_ascend.ops.layernorm import AscendRMSNorm
+from vllm_ascend.ops.linear import (AscendColumnParallelLinear,
+                                    AscendMergedColumnParallelLinear,
+                                    AscendQKVParallelLinear,
+                                    AscendRowParallelLinear)
+from vllm_ascend.ops.rotary_embedding import (
+    AscendDeepseekScalingRotaryEmbedding, AscendRotaryEmbedding)
+from vllm_ascend.ops.vocab_parallel_embedding import (
+    AscendLogitsProcessor, AscendParallelLMHead, AscendVocabParallelEmbedding)
+
+REGISTERED_ASCEND_OPS = {
+    "QuickGELU": AscendQuickGELU,
+    "SiluAndMul": AscendSiluAndMul,
+    "RotaryEmbedding": AscendRotaryEmbedding,
+    "ColumnParallelLinear": AscendColumnParallelLinear,
+    "RowParallelLinear": AscendRowParallelLinear,
+    "MergedColumnParallelLinear": AscendMergedColumnParallelLinear,
+    "QKVParallelLinear": AscendQKVParallelLinear,
+    "DeepseekScalingRotaryEmbedding": AscendDeepseekScalingRotaryEmbedding,
+    "VocabParallelEmbedding": AscendVocabParallelEmbedding,
+    "ParallelLMHead": AscendParallelLMHead,
+    "LogitsProcessor": AscendLogitsProcessor,
+    "RMSNorm": AscendRMSNorm,
+    "FusedMoE": AscendFusedMoE,
+    "MultiHeadLatentAttention": AscendMultiHeadLatentAttention,
+}
+
+
 def register_ascend_customop():
     """Register Ascend CustomOP
 
@@ -490,48 +521,8 @@ def register_ascend_customop():
         return
     from vllm.model_executor.custom_op import CustomOp
 
-    from vllm_ascend.ops.activation import AscendQuickGELU, AscendSiluAndMul
-    from vllm_ascend.ops.linear import (AscendColumnParallelLinear,
-                                        AscendMergedColumnParallelLinear,
-                                        AscendQKVParallelLinear,
-                                        AscendRowParallelLinear)
-    from vllm_ascend.ops.rotary_embedding import (
-        AscendDeepseekScalingRotaryEmbedding, AscendRotaryEmbedding)
-    from vllm_ascend.ops.vocab_parallel_embedding import (
-        AscendLogitsProcessor, AscendParallelLMHead,
-        AscendVocabParallelEmbedding)
-    CustomOp.register_oot(_decorated_op_cls=AscendQuickGELU, name="QuickGELU")
-    CustomOp.register_oot(_decorated_op_cls=AscendSiluAndMul,
-                          name="SiluAndMul")
-    CustomOp.register_oot(_decorated_op_cls=AscendRotaryEmbedding,
-                          name="RotaryEmbedding")
-    CustomOp.register_oot(_decorated_op_cls=AscendColumnParallelLinear,
-                          name="ColumnParallelLinear")
-    CustomOp.register_oot(_decorated_op_cls=AscendRowParallelLinear,
-                          name="RowParallelLinear")
-    CustomOp.register_oot(_decorated_op_cls=AscendMergedColumnParallelLinear,
-                          name="MergedColumnParallelLinear")
-    CustomOp.register_oot(_decorated_op_cls=AscendQKVParallelLinear,
-                          name="QKVParallelLinear")
-    CustomOp.register_oot(
-        _decorated_op_cls=AscendDeepseekScalingRotaryEmbedding,
-        name="DeepseekScalingRotaryEmbedding")
-    CustomOp.register_oot(_decorated_op_cls=AscendVocabParallelEmbedding,
-                          name="VocabParallelEmbedding")
-    CustomOp.register_oot(_decorated_op_cls=AscendParallelLMHead,
-                          name="ParallelLMHead")
-    CustomOp.register_oot(_decorated_op_cls=AscendLogitsProcessor,
-                          name="LogitsProcessor")
-
-    from vllm_ascend.ops.layernorm import AscendRMSNorm
-    CustomOp.register_oot(_decorated_op_cls=AscendRMSNorm, name="RMSNorm")
-
-    from vllm_ascend.ops.common_fused_moe import AscendFusedMoE
-    CustomOp.register_oot(_decorated_op_cls=AscendFusedMoE, name="FusedMoE")
-
-    from vllm_ascend.models.layers.mla import AscendMultiHeadLatentAttention
-    CustomOp.register_oot(_decorated_op_cls=AscendMultiHeadLatentAttention,
-                          name="MultiHeadLatentAttention")
+    for name, op_cls in REGISTERED_ASCEND_OPS.items():
+        CustomOp.register_oot(_decorated_op_cls=op_cls, name=name)
 
     # NOTE: Keep this at last to ensure all custom actions are registered
     _ASCEND_CUSTOMOP_IS_REIGISTERED = True