rm RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES & add lazy init for vllm_ascend_C

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

Author: zhuo97

tests/singlecard/ops/test_rotary_embedding.py

@@ -10,7 +10,9 @@
 import torch
 import torch.nn as nn
 
-import vllm_ascend.platform  # noqa: F401
+import vllm_ascend.platform as pf
+
+pf.CUSTOM_OP_ENABLED = True
 
 # Only Neox style true scenario is supported for now
 IS_NEOX_STYLE = [True]

vllm_ascend/attention/attention.py

@@ -36,7 +36,7 @@
 from vllm.utils import async_tensor_h2d, make_tensor_with_pad
 
 from vllm_ascend.ops.cache import concat_and_cache_mla
-from vllm_ascend.platform import CUSTOM_OP_ENABLED
+from vllm_ascend.utils import enable_custom_op
 from vllm_ascend.worker.model_runner import (
     ModelInputForNPUBuilder, ModelInputForNPUWithSamplingMetadata)
 
@@ -460,7 +460,7 @@ def advance_step(self,
         for i in range(num_queries):
             self.seq_lens[i] += 1
         self.max_decode_seq_len = max(self.seq_lens)
-        if CUSTOM_OP_ENABLED:
+        if enable_custom_op():
             #advance a step on NPU for existing inputs for a multi-step runner if custom ops is enabled
             torch.ops._C.advance_step_flashattn_ascendc(
                 num_seqs=num_seqs,

vllm_ascend/ops/rotary_embedding.py

@@ -22,11 +22,12 @@
 from vllm.model_executor.layers.rotary_embedding import (
     DeepseekScalingRotaryEmbedding, RotaryEmbedding)
 
-from vllm_ascend.platform import CUSTOM_OP_ENABLED
+from vllm_ascend.utils import enable_custom_op
 
 
 def custom_rotary_embedding_enabled(query, neox_style, head_size):
-    return query.dtype == torch.float16 and neox_style and head_size % 32 == 0 and CUSTOM_OP_ENABLED
+    return query.dtype == torch.float16 and neox_style and head_size % 32 == 0 and enable_custom_op(
+    )
 
 
 def rope_forward_oot(

vllm_ascend/platform.py

@@ -15,8 +15,6 @@
 # This file is a part of the vllm-ascend project.
 #
 
-import logging
-import os
 from typing import TYPE_CHECKING, Optional, Tuple
 
 import torch
@@ -27,14 +25,6 @@
 from vllm_ascend.utils import ASCEND_QUATIZATION_METHOD, update_aclgraph_sizes
 
 CUSTOM_OP_ENABLED = False
-try:
-    # register custom ops into torch_library here
-    import vllm_ascend.vllm_ascend_C  # type: ignore  # noqa: F401
-    CUSTOM_OP_ENABLED = True
-except ImportError as e:
-    logging.warning(
-        "Failed to import 'vllm_ascend.vllm_ascend_C': %s. All custom ops will be disabled. ",
-        e)
 
 if TYPE_CHECKING:
     from vllm.config import ModelConfig, VllmConfig
@@ -44,8 +34,6 @@
     VllmConfig = None
     FlexibleArgumentParser = None
 
-os.environ["RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES"] = "1"
-
 
 class NPUPlatform(Platform):
 

vllm_ascend/utils.py

@@ -54,6 +54,22 @@ def try_register_lib(lib_name: str, lib_info: str = ""):
         pass
 
 
+def enable_custom_op():
+    CUSTOM_OP_ENABLED = False
+    try:
+        # register custom ops into torch_library here
+        import vllm_ascend.vllm_ascend_C  # type: ignore  # noqa: F401
+
+    except ImportError:
+        logger.warning(
+            "Warning: Failed to register custom ops, all custom ops will be disabled"
+        )
+    else:
+        CUSTOM_OP_ENABLED = True
+
+    return CUSTOM_OP_ENABLED
+
+
 def find_hccl_library() -> str:
     """
     We either use the library file specified by the `HCCL_SO_PATH`