From 43a2ed2f471475e6681b9d3bb109fcc55be83661 Mon Sep 17 00:00:00 2001
From: handsomezhuzhu <2658601135@qq.com>
Date: Wed, 11 Feb 2026 00:14:35 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E5=9C=A8=E6=A8=A1=E5=9D=97=E5=8A=A0?=
 =?UTF-8?q?=E8=BD=BD=E6=97=B6=E8=BF=9B=E8=A1=8Cvllm=5Fnpu=5FC=E7=9A=84?=
 =?UTF-8?q?=E5=AF=BC=E5=85=A5=E6=A3=80=E6=9F=A5=EF=BC=8C=E4=BB=A5=E9=81=BF?=
 =?UTF-8?q?=E5=85=8Dtorch.compile/dynamo=E8=B7=9F=E8=B8=AA=E5=A4=B1?=
 =?UTF-8?q?=E8=B4=A5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 vllm_npu/utils.py | 31 ++++++++++++-------------------
 1 file changed, 12 insertions(+), 19 deletions(-)

diff --git a/vllm_npu/utils.py b/vllm_npu/utils.py
index 9b20407..a2d7909 100644
--- a/vllm_npu/utils.py
+++ b/vllm_npu/utils.py
@@ -47,7 +47,15 @@ REGISTERED_ASCEND_OPS = {}
 ACL_FORMAT_FRACTAL_ND = 2
 ACL_FORMAT_FRACTAL_NZ = 29
 
-_CUSTOM_OP_ENABLED = None
+# Pre-check for vllm_npu_C at module load time so that torch.compile/dynamo
+# never encounters a failing import during tracing.
+try:
+    import vllm_npu.vllm_npu_C  # type: ignore  # noqa: F401
+    import vllm_npu.meta_registration  # type: ignore  # noqa: F401
+    _CUSTOM_OP_ENABLED = True
+except (ImportError, ModuleNotFoundError):
+    _CUSTOM_OP_ENABLED = False
+
 _IS_310P = None
 _SLEEP_MODE_ENABLED = None
 _CURRENT_STREAM = None
@@ -219,25 +227,10 @@ def try_register_lib(lib_name: str, lib_info: str = ""):
 
 def enable_custom_op():
     """
-    Enable lazy init for vllm_npu_C to avoid early initialization of CANN's RTS component.
-    Ensure that ASCEND_RT_VISIBLE_DEVICES can be dynamically modified before torch.npu.set_device().
+    Check if vllm_npu_C custom ops are available.
+    The import check is done at module load time to avoid torch.compile/dynamo
+    tracing failures.
     """
-    global _CUSTOM_OP_ENABLED
-    if _CUSTOM_OP_ENABLED is not None:
-        return _CUSTOM_OP_ENABLED
-    try:
-        # isort: off
-        # register custom ops into torch_library here
-        import vllm_npu.vllm_npu_C  # type: ignore  # noqa: F401
-        # register the meta implementation for custom kernel if necessary
-        import vllm_npu.meta_registration  # type: ignore  # noqa: F401
-        # isort: on
-        _CUSTOM_OP_ENABLED = True
-    except ImportError:
-        _CUSTOM_OP_ENABLED = False
-        logger.warning(
-            "Warning: Failed to register custom ops, all custom ops will be disabled"
-        )
     return _CUSTOM_OP_ENABLED