[Diffusion][ERNIE] Fuse QKNorm with full-width RoPE (#34620)

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

@@ -1,6 +1,7 @@
 """ERNIE fused norm/scale/shift fast paths must stay bit-exact vs eager."""
 
 import sys
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -10,6 +11,8 @@ from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
 from sglang.multimodal_gen.runtime.models.dits.ernie_image import (
     _ernie_gated_norm_scale_shift,
     _ernie_norm_scale_shift,
+    _ernie_qknorm_rope,
+    _ernie_qknorm_rope_reference,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
 
@@ -53,5 +56,80 @@ def test_fused_norm_scale_shift_is_bit_exact(shape):
     assert not ernie_image._ERNIE_GATED_NORM.disabled
 
 
+def test_fused_qknorm_rope_is_bit_exact():
+    torch.manual_seed(1)
+    ernie_image._ERNIE_QKNORM_ROPE.disabled = False
+    ernie_image._ERNIE_QKNORM_ROPE.verified = False
+    batch, seq, heads, head_dim = 1, 257, 32, 128
+    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)
+    cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
+    sin = torch.randn_like(cos)
+    cache = torch.cat((cos, sin), dim=-1).contiguous()
+    positions = torch.arange(seq, device="cuda", dtype=torch.long)
+
+    q_ref, k_ref = _ernie_qknorm_rope_reference(
+        q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
+    )
+    q_out, k_out = _ernie_qknorm_rope(
+        q,
+        k,
+        q_norm,
+        k_norm,
+        head_dim,
+        cos,
+        sin,
+        cache,
+        positions,
+    )
+
+    assert torch.equal(q_out, q_ref)
+    assert torch.equal(k_out, k_ref)
+    assert ernie_image._ERNIE_QKNORM_ROPE.verified
+    assert not ernie_image._ERNIE_QKNORM_ROPE.disabled
+
+
+def test_qknorm_rope_first_attempt_exception_uses_pristine_inputs():
+    torch.manual_seed(2)
+    ernie_image._ERNIE_QKNORM_ROPE.disabled = False
+    ernie_image._ERNIE_QKNORM_ROPE.verified = False
+    batch, seq, heads, head_dim = 1, 17, 4, 128
+    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)
+    cos = torch.randn(seq, head_dim, device="cuda", dtype=torch.bfloat16)
+    sin = torch.randn_like(cos)
+    cache = torch.cat((cos, sin), dim=-1).contiguous()
+    positions = torch.arange(seq, device="cuda", dtype=torch.long)
+    q_ref, k_ref = _ernie_qknorm_rope_reference(
+        q.clone(), k.clone(), q_norm, k_norm, head_dim, cos, sin
+    )
+
+    def mutate_then_raise(**kwargs):
+        kwargs["q"].zero_()
+        kwargs["k"].zero_()
+        raise RuntimeError("synthetic kernel failure")
+
+    with patch.object(ernie_image, "apply_qk_norm_rope", mutate_then_raise):
+        q_out, k_out = _ernie_qknorm_rope(
+            q,
+            k,
+            q_norm,
+            k_norm,
+            head_dim,
+            cos,
+            sin,
+            cache,
+            positions,
+        )
+
+    assert torch.equal(q_out, q_ref)
+    assert torch.equal(k_out, k_ref)
+    assert ernie_image._ERNIE_QKNORM_ROPE.disabled
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

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

@@ -216,6 +216,45 @@ def test_qknorm_rope_preserves_split_bf16_rounding() -> None:
     assert torch.equal(k_ref, k_fused)
 
 
+def test_qknorm_rope_preserves_full_width_neox_cache() -> None:
+    from sglang.kernels.ops.diffusion.qknorm_rope import fused_inplace_qknorm_rope
+    from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
+
+    num_tokens, num_heads, head_dim = 257, 32, 128
+    q = torch.randn(num_tokens, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
+    k = torch.randn_like(q)
+    q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    positions = torch.arange(num_tokens, device=DEVICE, dtype=torch.int64)
+    cos = torch.randn(num_tokens, head_dim, device=DEVICE, dtype=DTYPE)
+    sin = torch.randn_like(cos)
+    cache = torch.cat((cos, sin), dim=-1).contiguous()
+
+    q_ref, k_ref = q.clone(), k.clone()
+    fused_inplace_qknorm(q_ref, k_ref, q_weight, k_weight, eps=1e-6)
+    half = head_dim // 2
+    q1, q2 = q_ref[..., :half], q_ref[..., half:]
+    k1, k2 = k_ref[..., :half], k_ref[..., half:]
+    q_ref = torch.cat((-q2, q1), dim=-1) * sin[:, None, :] + q_ref * cos[:, None, :]
+    k_ref = torch.cat((-k2, k1), dim=-1) * sin[:, None, :] + k_ref * cos[:, None, :]
+
+    fused_inplace_qknorm_rope(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        cache,
+        positions,
+        is_neox=True,
+        eps=1e-6,
+        round_norm_before_rope=True,
+        cache_has_full_width=True,
+    )
+
+    assert torch.equal(q, q_ref)
+    assert torch.equal(k, k_ref)
+
+
 def test_qknorm_rope_requires_opt_in_for_strided_packed_gqa() -> None:
     from sglang.kernels.ops.diffusion.qknorm_rope import (
         fused_inplace_qknorm_rope,