update VLLM_USE_FUSE_SILU_AND_MUL

vllm/envs.py

@@ -186,6 +186,7 @@ if TYPE_CHECKING:
     VLLM_USE_CUDA_GRAPH_SIZES: bool = False
     VLLM_USE_CAT_MLA: bool = False
     VLLM_REJECT_SAMPLE_OPT: bool = False
+    VLLM_USE_FUSE_SILU_AND_MUL: bool = True
 
 def get_default_cache_root():
     return os.getenv(
@@ -1206,6 +1207,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_REJECT_SAMPLE_OPT":
         lambda: (os.getenv('VLLM_REJECT_SAMPLE_OPT', 'False').lower() in
                  ("true", "1")),      
+
+    # vLLM will use fused silu+mul kernel
+    "VLLM_USE_FUSE_SILU_AND_MUL":
+        lambda: (os.environ.get("VLLM_USE_FUSE_SILU_AND_MUL", "False").lower() in
+                 ("true", "1")),
 }
 
 # --8<-- [end:env-vars-definition]

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -1875,8 +1875,12 @@ def fused_experts_impl(
                                 use_nn_moe=use_nn_moe)
 
         if activation == "silu":
-            torch.ops._C.silu_and_mul(intermediate_cache2,
-                                      intermediate_cache1.view(-1, N))
+            if envs.VLLM_USE_FUSE_SILU_AND_MUL:
+                from lightop import fuse_silu_and_mul
+                fuse_silu_and_mul(intermediate_cache1.view(-1, N),intermediate_cache2)    
+            else:
+                torch.ops._C.silu_and_mul(intermediate_cache2,
+                                        intermediate_cache1.view(-1, N))
         elif activation == "gelu":
             torch.ops._C.gelu_and_mul(intermediate_cache2,
                                       intermediate_cache1.view(-1, N))

vllm/model_executor/model_loader/utils.py

@@ -258,7 +258,8 @@ def get_model_architecture(
             else:
                 if not envs.is_set("VLLM_USE_PD_SPLIT"):
                     os.environ['VLLM_USE_PD_SPLIT'] = '1'
-                
+                if not envs.is_set("VLLM_USE_FUSE_SILU_AND_MUL"):
+                    os.environ['VLLM_USE_FUSE_SILU_AND_MUL'] = '1'
                 
             if os.getenv('GEMM_PAD') != '1': 
                 os.environ['GEMM_PAD'] = '0'
@@ -279,6 +280,8 @@ def get_model_architecture(
             else:
                 if not envs.is_set("VLLM_USE_PD_SPLIT"):
                     os.environ['VLLM_USE_PD_SPLIT'] = '1'
+                if not envs.is_set("VLLM_USE_FUSE_SILU_AND_MUL"):
+                    os.environ['VLLM_USE_FUSE_SILU_AND_MUL'] = '1'
                     
         # awq相关配置
         try: