From 267ff1b5f9ed63db8768d43a88fa45c936e27b42 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Mon, 6 Jul 2026 19:15:21 -0700 Subject: [PATCH] Fix LTX2 RoPE JIT kernel CI (#30278) Co-authored-by: Brayden Zhong --- .../csrc/diffusion/ltx2_qknorm_split_rope.cuh | 5 ++--- .../diffusion/test_ltx2_qknorm_split_rope.py | 18 +++++++----------- 2 files changed, 9 insertions(+), 14 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh b/python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh index 794007197..3cf399055 100644 --- a/python/sglang/jit_kernel/csrc/diffusion/ltx2_qknorm_split_rope.cuh +++ b/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: -// torch.nn.RMSNorm(input) returns fp32 under bf16 autocast, then split RoPE -// runs in fp32 and rounds once to bf16 at the final attention input. +// This mirrors the LTX2 eager oracle: RMSNorm and split RoPE both run in +// fp32, rounding to bf16 only once at the final attention input. #pragma once diff --git a/test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py b/test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py index a52991148..dbf70bb02 100644 --- a/test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py +++ b/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)