replace triton_ of rms and act_and_mul

vllm/model_executor/layers/activation.py

@@ -75,7 +75,9 @@ class SiluAndMul(CustomOp):
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""
-        if not torch.compiler.is_compiling() and envs.VLLM_ENABLE_TBO:
+        if not torch.compiler.is_compiling() and envs.envs.VLLM_ENABLE_TBO:
+            return self.forward_cuda(x)
+        elif envs.VLLM_USE_OPT_OP:
             return self.forward_cuda(x)
         else:
             d = x.shape[-1] // 2

vllm/model_executor/layers/layernorm.py

@@ -165,8 +165,9 @@ class RMSNorm(CustomOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
-        
-        if not torch.compiler.is_compiling() and envs.VLLM_ENABLE_TBO:  
+        if not torch.compiler.is_compiling() and envs.VLLM_ENABLE_TBO:
+            return self.forward_cuda(x, residual) 
+        elif envs.VLLM_USE_OPT_OP:
             return self.forward_cuda(x, residual)  
         else:
             orig_dtype = x.dtype
@@ -203,7 +204,7 @@ class RMSNorm(CustomOp):
         self,
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
-    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+    ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: 
         if self.variance_size_override is not None:
             return self.forward_native(x, residual)