From 6ece6660ecf1317043d7275110a7d73fa19f5ecd Mon Sep 17 00:00:00 2001
From: wangxiyuan <wangxiyuan1007@gmail.com>
Date: Wed, 3 Dec 2025 19:27:38 +0800
Subject: [PATCH] fix custom ops env set error (#4675)

Move Custom ops register to correct place to make CI happy

- vLLM version: 86e178f7c4d8c3b0eaf3c8e3f810a83f63b90e24
- vLLM main:
https://github.com/vllm-project/vllm/commit/86e178f7c4d8c3b0eaf3c8e3f810a83f63b90e24

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
---
 vllm_ascend/platform.py | 19 ++++++++++++++++++-
 vllm_ascend/utils.py    | 12 ------------
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/vllm_ascend/platform.py b/vllm_ascend/platform.py
index 6337f8f9..f59d1ed1 100644
--- a/vllm_ascend/platform.py
+++ b/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(
diff --git a/vllm_ascend/utils.py b/vllm_ascend/utils.py
index 97eeafa8..fafd6bb1 100644
--- a/vllm_ascend/utils.py
+++ b/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: