polish code

Author: cyber-pioneer

vllm_fl/dispatch/backends/flaggems/impl/activation.py

@@ -7,7 +7,8 @@
 from __future__ import annotations
 
 import torch
-from vllm_fl.utils import use_flaggems_vllm
+
+import vllm_fl.envs as fl_envs
 
 
 def silu_and_mul_flaggems(obj, x: torch.Tensor) -> torch.Tensor:
@@ -21,7 +22,7 @@ def silu_and_mul_flaggems(obj, x: torch.Tensor) -> torch.Tensor:
     Returns:
         Output tensor of shape [..., d]
     """
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.silu_and_mul import silu_and_mul
     else:
         from flag_gems import silu_and_mul
@@ -42,7 +43,7 @@ def gelu_and_mul_flaggems(obj, x: torch.Tensor) -> torch.Tensor:
     Returns:
         Output tensor of shape [..., d]
     """
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.gelu_and_mul import gelu_and_mul
     else:
         from flag_gems import gelu_and_mul
@@ -69,7 +70,7 @@ def silu_and_mul_with_clamp_flaggems(x: torch.Tensor, swiglu_limit: torch.Tensor
     Returns:
         Output tensor of shape [..., d]
     """
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.silu_and_mul_with_clamp import silu_and_mul_with_clamp_kernel
     else:
         from flag_gems.fused.silu_and_mul_with_clamp import silu_and_mul_with_clamp_kernel

vllm_fl/dispatch/backends/flaggems/impl/attention.py

@@ -40,9 +40,10 @@
 )
 from vllm.v1.kv_cache_interface import AttentionSpec
 from vllm.platforms.interface import DeviceCapability
-from vllm_fl.utils import use_flaggems_vllm
 
-if use_flaggems_vllm():
+import vllm_fl.envs as fl_envs
+
+if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
     from flaggems_vllm.ops.flash_attn_varlen_func import flash_attn_varlen_func
     from flaggems_vllm.ops.reshape_and_cache_flash import reshape_and_cache_flash
 else:

vllm_fl/dispatch/backends/flaggems/impl/fused_moe.py

@@ -9,7 +9,8 @@
 import torch
 from vllm.triton_utils import triton
 from vllm.utils.math_utils import round_up
-from vllm_fl.utils import use_flaggems_vllm
+
+import vllm_fl.envs as fl_envs
 
 
 def moe_align_block_size_flaggems(
@@ -20,7 +21,7 @@ def moe_align_block_size_flaggems(
     pad_sorted_ids: bool = False,
     ignore_invalid_experts: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.moe_align_block_size import moe_align_block_size_triton
     else:
         from flag_gems import moe_align_block_size_triton
@@ -60,7 +61,7 @@ def moe_align_block_size_flaggems(
 def topk_softmax_flaggems(
     topk_weights, topk_indices, token_expert_indices, gating_output, renormalize=False
 ):
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.topk_softmax import topk_softmax
         topk_softmax(
             topk_weights,
@@ -141,7 +142,7 @@ def _alloc_topk_buffers():
         return topk_weights, topk_ids
 
     elif scoring_func == "sqrtsoftplus":
-        if use_flaggems_vllm():
+        if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
             from flaggems_vllm.ops.topk_softplus_sqrt import topk_softplus_sqrt
         else:
             from flag_gems import topk_softplus_sqrt
@@ -205,7 +206,7 @@ def invoke_fused_moe_triton_kernel_flaggems(
     block_shape=None,
     B_bias=None,
 ):
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.invoke_fused_moe_kernel import invoke_fused_moe_triton_kernel
     else:
         from flag_gems import invoke_fused_moe_triton_kernel
@@ -244,7 +245,7 @@ def grouped_topk_flaggems(
     bias,
     scoring_func=0,
 ):
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.grouped_topk import grouped_topk
     else:
         from flag_gems import grouped_topk
@@ -262,7 +263,7 @@ def grouped_topk_flaggems(
 
 
 def moe_sum_flaggems(inp, out):
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.moe_sum import moe_sum
     else:
         from flag_gems import moe_sum

vllm_fl/dispatch/backends/flaggems/impl/mhc.py

@@ -5,7 +5,8 @@
 """
 
 import torch
-from vllm_fl.utils import use_flaggems_vllm
+
+import vllm_fl.envs as fl_envs
 
 
 def mhc_pre_flaggems(
@@ -23,7 +24,7 @@ def mhc_pre_flaggems(
     """FlagGems native implementation of mhc_pre."""
     import importlib
 
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.mhc.mhc_pre import mhc_pre
         _mhc_pre_mod = importlib.import_module('flaggems_vllm.ops.mhc.mhc_pre')
     else:
@@ -85,7 +86,7 @@ def mhc_post_flaggems(
     comb: torch.Tensor,
 ) -> torch.Tensor:
     """FlagGems native implementation of mhc_post."""
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.mhc_post import mhc_post
     else:
         from flag_gems import mhc_post
@@ -105,7 +106,7 @@ def hc_head_fused_kernel_flaggems(
     hc_mult: int,
 ) -> None:
     """FlagGems native implementation of hc_head_fused_kernel. Mutates `out` in-place."""
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.hc_head_fused_kernel import hc_head_fused_kernel
     else:
         from flag_gems import hc_head_fused_kernel

vllm_fl/dispatch/backends/flaggems/impl/mla.py

@@ -23,9 +23,9 @@
     MLACommonMetadata,
 )
 
-from vllm_fl.utils import use_flaggems_vllm
+import vllm_fl.envs as fl_envs
 
-if use_flaggems_vllm():
+if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
     from flaggems_vllm.ops.flash_attn_varlen_func import flash_attn_varlen_func
     from flaggems_vllm.ops.flash_mla import flash_mla
 else:

vllm_fl/dispatch/backends/flaggems/impl/normalization.py

@@ -9,7 +9,8 @@
 from typing import Optional, Union
 
 import torch
-from vllm_fl.utils import use_flaggems_vllm
+
+import vllm_fl.envs as fl_envs
 
 
 def rms_norm_flaggems(
@@ -28,7 +29,7 @@ def rms_norm_flaggems(
     Returns:
         Normalized tensor, or tuple of (normalized, residual) if residual is provided
     """
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.rms_norm import gems_rms_forward
     else:
         from flag_gems import rms_norm_forward as gems_rms_forward

vllm_fl/dispatch/backends/flaggems/impl/rotary.py

@@ -7,7 +7,8 @@
 from __future__ import annotations
 
 import torch
-from vllm_fl.utils import use_flaggems_vllm
+
+import vllm_fl.envs as fl_envs
 
 
 def rotary_embedding_flaggems(
@@ -36,7 +37,7 @@ def rotary_embedding_flaggems(
     Returns:
         Tuple of (embedded_query, embedded_key)
     """
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.rope import gems_rope_forward
     else:
         from flag_gems import apply_rotary_pos_emb as gems_rope_forward

vllm_fl/dispatch/backends/flaggems/impl/router_gemm.py

@@ -3,13 +3,14 @@
 """FlagGems MoE router GEMM implementation."""
 
 import torch
-from vllm_fl.utils import use_flaggems_vllm
+
+import vllm_fl.envs as fl_envs
 
 
 def router_gemm_bf16_fp32_flaggems(
     x: torch.Tensor, weight: torch.Tensor
 ) -> torch.Tensor:
-    if use_flaggems_vllm():
+    if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
         from flaggems_vllm.ops.router_gemm import router_gemm
     else:
         from flag_gems import router_gemm

vllm_fl/envs.py

@@ -13,6 +13,9 @@
         "FLAGGEMS_ENABLE_OPLIST_PATH", "/tmp/flaggems_enable_oplist.txt"
     ),
     "USE_FLAGGEMS": use_flaggems,
+    "VLLM_FL_USE_FLAGGEMS_VLLM": lambda: (
+        os.environ.get("VLLM_FL_USE_FLAGGEMS_VLLM", "1").lower() in ("1", "true")
+    ),
 }
 
 

vllm_fl/ops/fused_moe/fused_moe_utils.py

@@ -44,6 +44,7 @@
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl
 
+import vllm_fl.envs as fl_envs
 from vllm_fl.dispatch import CachedOp
 from vllm_fl.ops.fused_moe.activation import apply_moe_activation
 from vllm_fl.utils import use_flaggems
@@ -328,9 +329,7 @@ def apply(
         # Fast path (no LoRA, NVIDIA only): let FlagGems own both expert GEMMs
         # for unquantized and W8A16 inputs.
         if self._lora_context is None and current_platform.is_cuda():
-            from vllm_fl.utils import use_flaggems_vllm
-
-            if use_flaggems_vllm():
+            if fl_envs.VLLM_FL_USE_FLAGGEMS_VLLM:
                 import flaggems_vllm
                 fused_experts_impl = flaggems_vllm.fused_experts_impl
             else: