[Bugfix] Fix dtype mismatch in RMSNormGated.forward_native() during torch.compile (#35256)

Signed-off-by: haosdent <haosdent@gmail.com>

vllm/model_executor/layers/layernorm.py

@@ -557,6 +557,11 @@ class RMSNormGated(CustomOp):
             - norm_before_gate=True: out = norm(x) * silu(z)
             - norm_before_gate=False: out = norm(x * silu(z))
         """
+        orig_dtype = x.dtype
+        x = x.float()
+        weight = self.weight.float()
+        z = z.float() if z is not None else None
+
         # Apply gating before normalization if needed
         if z is not None and not self.norm_before_gate:
             x = x * F.silu(z)
@@ -566,7 +571,7 @@ class RMSNormGated(CustomOp):
             # Standard RMS norm across the last dimension
             variance = x.pow(2).mean(dim=-1, keepdim=True)
             x_normed = x * torch.rsqrt(variance + self.eps)
-            out = x_normed * self.weight
+            out = x_normed * weight
         else:
             # Group RMS norm
             from einops import rearrange
@@ -574,13 +579,13 @@ class RMSNormGated(CustomOp):
             x_group = rearrange(x, "... (g d) -> ... g d", d=self.group_size)
             variance = x_group.pow(2).mean(dim=-1, keepdim=True)
             x_normed = x_group * torch.rsqrt(variance + self.eps)
-            out = rearrange(x_normed, "... g d -> ... (g d)") * self.weight
+            out = rearrange(x_normed, "... g d -> ... (g d)") * weight
 
         # Apply gating after normalization if needed
         if z is not None and self.norm_before_gate:
             out = out * F.silu(z)
 
-        return out.to(x.dtype)
+        return out.to(orig_dtype)
 
     def forward_cuda(
         self, x: torch.Tensor, z: torch.Tensor | None = None