Make RMSNormGated a CustomOp

Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>

Author: vadiklyutiy

vllm/model_executor/layers/layernorm.py

@@ -14,6 +14,7 @@
 )
 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import direct_register_custom_op
+from vllm.model_executor.layers.fla.ops.layernorm_guard import rmsnorm_fn
 
 
 def is_rocm_aiter_rmsnorm_enabled() -> bool:
@@ -368,8 +369,8 @@ def forward_cuda(
             self._is_compiled = True
         return self.forward_native(x, residual)
 
-
-class RMSNormGated(nn.Module):
+@CustomOp.register("rms_norm_gated")
+class RMSNormGated(CustomOp):
     """RMS Normalization with optional gating.
 
     This is a native PyTorch implementation that supports:
@@ -413,7 +414,7 @@ def __init__(
     def reset_parameters(self):
         torch.nn.init.ones_(self.weight)
 
-    def forward(self, x: torch.Tensor, z: torch.Tensor | None = None) -> torch.Tensor:
+    def forward_native(self, x: torch.Tensor, z: torch.Tensor | None = None) -> torch.Tensor:
         """
         Native PyTorch implementation of RMS normalization with gating.
 
@@ -453,6 +454,17 @@ def forward(self, x: torch.Tensor, z: torch.Tensor | None = None) -> torch.Tenso
 
         return out
 
+    def forward_cuda(self, x: torch.Tensor, z: torch.Tensor | None = None) -> torch.Tensor:
+        return rmsnorm_fn(
+            x,
+            self.weight,
+            self.bias,
+            z=z,
+            eps=self.eps,
+            group_size=self.group_size,
+            norm_before_gate=self.norm_before_gate,
+        )
+
 
 class LayerNorm(nn.Module):
     """