Fix LTX2 RoPE JIT kernel CI (#30278)

Co-authored-by: Brayden Zhong <brayden@radixark.ai>

test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py

@@ -2,6 +2,7 @@ import sys
 
 import pytest
 import torch
+import torch.nn.functional as F
 
 from sglang.jit_kernel.diffusion.ltx2_qknorm_split_rope import (
     can_use_ltx2_qknorm_split_rope_cuda,
@@ -72,17 +73,12 @@ def _reference(
     k_weight: torch.Tensor,
     eps: float,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    q_norm = torch.nn.RMSNorm(q.shape[-1], eps=eps, device="cuda").to(
-        dtype=torch.bfloat16
-    )
-    k_norm = torch.nn.RMSNorm(k.shape[-1], eps=eps, device="cuda").to(
-        dtype=torch.bfloat16
-    )
-    q_norm.weight.data.copy_(q_weight)
-    k_norm.weight.data.copy_(k_weight)
-    with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
-        q_ref = _apply_split_rotary_ref(q_norm(q), q_cos, q_sin)
-        k_ref = _apply_split_rotary_ref(k_norm(k), k_cos, k_sin)
+    # rms_norm isn't autocast fp32-preserving, so feed fp32 inputs directly
+    # to keep the normalized value unrounded until the final RoPE output.
+    q_norm = F.rms_norm(q.float(), (q.shape[-1],), q_weight.float(), eps)
+    k_norm = F.rms_norm(k.float(), (k.shape[-1],), k_weight.float(), eps)
+    q_ref = _apply_split_rotary_ref(q_norm, q_cos, q_sin)
+    k_ref = _apply_split_rotary_ref(k_norm, k_cos, k_sin)
     return q_ref.to(dtype=torch.bfloat16), k_ref.to(dtype=torch.bfloat16)