[Perf] Minimax-H3 support fused RMSNorm and RoPE opt (#5801)

Signed-off-by: Fan <fan@FandeMacBook-Air.local>
Co-authored-by: Fan <fan@FandeMacBook-Air.local>

vllm_omni/diffusion/layers/norm.py

@@ -65,9 +65,9 @@ class LayerNorm(nn.LayerNorm, CustomOp):
 
 
 class RMSNorm(CustomOp):
-    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+    def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None:
         super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
         self.variance_epsilon = eps
 
     def _forward_fused(self, x: torch.Tensor) -> torch.Tensor: