Add a switch to custom_allgather import

Author: shihaobai

lightllm/common/vllm_kernel/_custom_ops.py

@@ -16,11 +16,13 @@
         logger.error("vllm or lightllm_kernel is not installed, you can't use custom ops")
 
 try:
-    from lightllm.common.vllm_kernel._ops import init_custom_gather_ar
-    from lightllm.common.vllm_kernel._ops import all_gather
-    from lightllm.common.vllm_kernel._ops import allgather_dispose
-    from lightllm.common.vllm_kernel._ops import allgather_register_buffer
-    from lightllm.common.vllm_kernel._ops import allgather_get_graph_buffer_ipc_meta
-    from lightllm.common.vllm_kernel._ops import allgather_register_graph_buffers
+    ENABLE_CUSTOM_GATHER = os.getenv("ENABLE_CUSTOM_GATHER", "False").upper() in ["ON", "TRUE", "1"]
+    if ENABLE_CUSTOM_GATHER:
+        from lightllm.common.vllm_kernel._ops import init_custom_gather_ar
+        from lightllm.common.vllm_kernel._ops import all_gather
+        from lightllm.common.vllm_kernel._ops import allgather_dispose
+        from lightllm.common.vllm_kernel._ops import allgather_register_buffer
+        from lightllm.common.vllm_kernel._ops import allgather_get_graph_buffer_ipc_meta
+        from lightllm.common.vllm_kernel._ops import allgather_register_graph_buffers
 except ImportError:
     logger.error("lightllm_kernel is not installed, you can't use custom allgather")