[diffusion] Fuse LongCat-Image QKNorm and interleaved RoPE (#35995)

test/registered/kernels/ops/diffusion/test_model_fast_paths.py

@@ -34,6 +34,7 @@ import sglang.multimodal_gen.runtime.models.dits.ernie_image as ernie_image
 import sglang.multimodal_gen.runtime.models.dits.flux as flux
 import sglang.multimodal_gen.runtime.models.dits.flux_2 as flux2
 import sglang.multimodal_gen.runtime.models.dits.glm_image as glm_image
+import sglang.multimodal_gen.runtime.models.dits.longcat_image as longcat_image
 import sglang.multimodal_gen.runtime.models.dits.ltx_2 as ltx2_module
 import sglang.multimodal_gen.runtime.models.dits.sana as sana
 from sglang.kernels.ops.diffusion import (
@@ -56,7 +57,11 @@ from sglang.kernels.ops.diffusion.common.platform import is_cuda
 from sglang.multimodal_gen.configs.models.vaes.stablediffusion3 import (
     StableDiffusion3VAEConfig,
 )
-from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm, RMSNormNoWeight
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    RMSNorm,
+    RMSNormNoWeight,
+    apply_qk_norm,
+)
 from sglang.multimodal_gen.runtime.layers.rotary_embedding.utils import (
     _apply_rotary_emb,
 )
@@ -85,6 +90,9 @@ from sglang.multimodal_gen.runtime.models.dits.hunyuanvideo import (
     _hunyuan_pack_qkv,
     _hunyuan_qknorm,
 )
+from sglang.multimodal_gen.runtime.models.dits.longcat_image import (
+    _apply_longcat_qknorm_rope,
+)
 from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_rms_norm_modulate
 from sglang.multimodal_gen.runtime.models.dits.sana import (
     _eager_ln_modulate as _sana_eager_ln_modulate,
@@ -490,6 +498,54 @@ def test_ernie_qknorm_rope_first_attempt_exception_uses_pristine_inputs():
     assert ernie_image._ERNIE_QKNORM_ROPE.disabled
 
 
+# -------------------------------------------------------------------------
+# LongCat-Image -- full-width interleaved QKNorm + RoPE
+# -------------------------------------------------------------------------
+
+
+@requires_inline_ptx
+def test_longcat_qknorm_rope_is_bit_exact():
+    torch.manual_seed(3)
+    batch, seq, heads, head_dim = 2, 17, 24, 128
+    offset = 11
+    q = torch.randn(batch, seq, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn_like(q)
+    q_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    k_norm = RMSNorm(head_dim, eps=1e-6).to(device="cuda", dtype=torch.bfloat16)
+    with torch.no_grad():
+        q_norm.weight.copy_(torch.randn_like(q_norm.weight))
+        k_norm.weight.copy_(torch.randn_like(k_norm.weight))
+
+    cos = torch.randn(offset + seq, head_dim, device="cuda")
+    sin = torch.randn_like(cos)
+    image_rotary_emb = (cos[offset:], sin[offset:])
+    cache = torch.cat((cos, sin), dim=-1).contiguous()
+    positions = torch.arange(offset, offset + seq, device="cuda", dtype=torch.int64)
+
+    q_ref, k_ref = apply_qk_norm(q.clone(), k.clone(), q_norm, k_norm, head_dim)
+    q_ref = longcat_image.apply_rotary_emb(q_ref, image_rotary_emb, sequence_dim=1)
+    k_ref = longcat_image.apply_rotary_emb(k_ref, image_rotary_emb, sequence_dim=1)
+
+    q_fused, k_fused = q.clone(), k.clone()
+    q_out, k_out = _apply_longcat_qknorm_rope(
+        q_fused,
+        k_fused,
+        q_norm,
+        k_norm,
+        head_dim,
+        image_rotary_emb,
+        cache,
+        positions,
+    )
+
+    assert q_out.data_ptr() == q_fused.data_ptr()
+    assert k_out.data_ptr() == k_fused.data_ptr()
+    assert torch.equal(q_out, q_ref)
+    assert torch.equal(k_out, k_ref)
+    assert longcat_image._LONGCAT_QKNORM_ROPE.verified
+    assert not longcat_image._LONGCAT_QKNORM_ROPE.disabled
+
+
 # -------------------------------------------------------------------------
 # LTX-2 -- weightless RMSNorm + modulate (quality-gated)
 # -------------------------------------------------------------------------