diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py
index 83897eb15..8d91df8f0 100644
--- a/vllm_omni/diffusion/layers/norm.py
+++ b/vllm_omni/diffusion/layers/norm.py
@@ -65,9 +65,10 @@ class LayerNorm(nn.LayerNorm, CustomOp):
 
 
 class RMSNorm(CustomOp):
-    def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None:
+    def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype | None = None) -> None:
         super().__init__()
-        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
+        weight_dtype = dtype or torch.get_default_dtype()
+        self.weight = nn.Parameter(torch.ones(hidden_size, dtype=weight_dtype))
         self.hidden_size = hidden_size
         self.variance_epsilon = eps