Fix LTX2 RoPE JIT kernel CI (#30278)

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

python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh

@@ -3,9 +3,8 @@
 // Developed with MIT HAN Lab Kernel Design Agents:
 // https://github.com/mit-han-lab/kernel-design-agents
 //
-// This mirrors the LTX2 eager oracle:
+// This mirrors the LTX2 eager oracle: RMSNorm and split RoPE both run in
-//   torch.nn.RMSNorm(input) returns fp32 under bf16 autocast, then split RoPE
+// fp32, rounding to bf16 only once at the final attention input.
-//   runs in fp32 and rounds once to bf16 at the final attention input.
 
 #pragma once
 

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(
+    # rms_norm isn't autocast fp32-preserving, so feed fp32 inputs directly
-        dtype=torch.bfloat16
+    # 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 = torch.nn.RMSNorm(k.shape[-1], eps=eps, device="cuda").to(
+    k_norm = F.rms_norm(k.float(), (k.shape[-1],), k_weight.float(), eps)
-        dtype=torch.bfloat16
+    q_ref = _apply_split_rotary_ref(q_norm, q_cos, q_sin)
-    )
+    k_ref = _apply_split_rotary_ref(k_norm, k_cos, k_sin)
-    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)
     return q_ref.to(dtype=torch.bfloat16), k_ref.to(dtype=torch.bfloat16)