Add a switch to custom_allgather import (#705)

Co-authored-by: baishihao <[email protected]>
ModelTC · Jan 16, 2025 · bd54e43 · bd54e43
1 parent 0fcedc0
commit bd54e43
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/lightllm/common/vllm_kernel/_custom_ops.py b/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")