fix custom ops env set error (#4675)

Move Custom ops register to correct place to make CI happy

- vLLM version: 86e178f7c4d8c3b0eaf3c8e3f810a83f63b90e24
- vLLM main:
86e178f7c4

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>

vllm_ascend/platform.py

@@ -46,6 +46,8 @@ else:
     VllmConfig = None
     FlexibleArgumentParser = None
 
+CUSTOM_OP_REGISTERED = False
+
 
 class NPUPlatform(Platform):
 
@@ -341,7 +343,22 @@ class NPUPlatform(Platform):
         # TODO: when the above issue is fixed, we can uncomment the following lines.
         # from vllm_ascend.utils import enable_custom_op
         # enable_custom_op()
-        pass
+        # set custom ops path
+        global CUSTOM_OP_REGISTERED
+        if CUSTOM_OP_REGISTERED:
+            return
+        CUR_DIR = os.path.dirname(os.path.realpath(__file__))
+        CUSTOM_OPP_PATH = os.path.join(CUR_DIR, "_cann_ops_custom", "vendors",
+                                       "vllm-ascend")
+        if os.path.exists(CUSTOM_OPP_PATH):
+            current_cust_opp_path = os.environ.get("ASCEND_CUSTOM_OPP_PATH",
+                                                   "")
+            if current_cust_opp_path:
+                os.environ[
+                    "ASCEND_CUSTOM_OPP_PATH"] = f"{CUSTOM_OPP_PATH}:{current_cust_opp_path}"
+            else:
+                os.environ["ASCEND_CUSTOM_OPP_PATH"] = CUSTOM_OPP_PATH
+        CUSTOM_OP_REGISTERED = True
 
     @classmethod
     def get_attn_backend_cls(

vllm_ascend/utils.py

@@ -248,18 +248,6 @@ def enable_custom_op():
     """
     global _CUSTOM_OP_ENABLED
 
-    # set custom ops path
-    CUR_DIR = os.path.dirname(os.path.realpath(__file__))
-    CUSTOM_OPP_PATH = os.path.join(CUR_DIR, "_cann_ops_custom", "vendors",
-                                   "vllm-ascend")
-    if os.path.exists(CUSTOM_OPP_PATH):
-        current_cust_opp_path = os.environ.get("ASCEND_CUSTOM_OPP_PATH", "")
-        if current_cust_opp_path:
-            os.environ[
-                "ASCEND_CUSTOM_OPP_PATH"] = f"{CUSTOM_OPP_PATH}:{current_cust_opp_path}"
-        else:
-            os.environ["ASCEND_CUSTOM_OPP_PATH"] = CUSTOM_OPP_PATH
-
     if _CUSTOM_OP_ENABLED is not None:
         return _CUSTOM_OP_ENABLED
     try: