use opt layernorm_kernels

4440e8c0 · zhuwenwen · d8ae62c7 · 4440e8c0
Commit 4440e8c0 authored Aug 22, 2024 by zhuwenwen
Show whitespace changes
Inline Side-by-side

Showing with 26 additions and 26 deletions

vllm/model_executor/layers/layernorm.py vllm/model_executor/layers/layernorm.py +26 -26

No files found.
--- a/vllm/model_executor/layers/layernorm.py
+++ b/vllm/model_executor/layers/layernorm.py
@@ -52,14 +52,14 @@ class RMSNorm(CustomOp):
        from vllm import _custom_ops as ops
        if residual is not None:
-            # if envs.VLLM_USE_OPT_OP:
+            if envs.VLLM_USE_OPT_OP:
-            #     ops.fused_add_rms_norm_opt(
+                ops.fused_add_rms_norm_opt(
-            #         x,
+                    x,
-            #         residual,
+                    residual,
-            #         self.weight.data,
+                    self.weight.data,
-            #         self.variance_epsilon,
+                    self.variance_epsilon,
-            #     )
+                )
-            # else:
+            else:
                ops.fused_add_rms_norm(
                    x,
                    residual,
@@ -68,14 +68,14 @@ class RMSNorm(CustomOp):
                )
            return x, residual
        out = torch.empty_like(x)
-        # if envs.VLLM_USE_OPT_OP:
+        if envs.VLLM_USE_OPT_OP:
-        #     ops.rms_norm_opt(
+            ops.rms_norm_opt(
-        #         out,
+                out,
-        #         x,
+                x,
-        #         self.weight.data,
+                self.weight.data,
-        #         self.variance_epsilon,
+                self.variance_epsilon,
-        #     )
+            )
-        # else:
+        else:
            ops.rms_norm(
                out,
                x,