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,9 +165,10 @@ 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:
             return self.forward_cuda(x, residual) 
+        elif envs.VLLM_USE_OPT_OP:
+            return self.forward_cuda(x, residual)  
         else:
             orig_dtype = x.dtype
             x = x.to(torch.float32)