[Kernel] Fuse Q/K RMSNorm and RoPE (#5990)

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: Mengjie Zhao <zmj0129@gmail.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

tests/diffusion/layers/test_fused_qk_norm_rope.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import torch.nn.functional as F
+from vllm.triton_utils import HAS_TRITON
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda, pytest.mark.diffusion]
+
+_HEAD_DIM = 128
+_ROTARY_DIM = 96
+_EPS = 1e-5
+
+
+def _reference(q, k, q_weight, k_weight, rope_table):
+    q = F.rms_norm(q, (_HEAD_DIM,), q_weight, _EPS)
+    k = F.rms_norm(k, (_HEAD_DIM,), k_weight, _EPS)
+    half = _ROTARY_DIM // 2
+    cos = rope_table[..., :half].unsqueeze(1)
+    sin = rope_table[..., half:].unsqueeze(1)
+
+    def apply(x):
+        first = x[..., :half]
+        second = x[..., half:_ROTARY_DIM]
+        return torch.cat(
+            (
+                first * cos - second * sin,
+                second * cos + first * sin,
+                x[..., _ROTARY_DIM:],
+            ),
+            dim=-1,
+        )
+
+    return apply(q), apply(k)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not HAS_TRITON, reason="Triton required")
+@pytest.mark.parametrize("seq_len", [1, 257, 1024])
+def test_fused_qk_norm_rope_matches_bf16_reference(seq_len):
+    from vllm_omni.diffusion.layers.fused_qk_norm_rope import (
+        fused_qk_norm_rope,
+    )
+
+    torch.manual_seed(17)
+    heads = 14
+    qkv = torch.randn(
+        seq_len,
+        heads * _HEAD_DIM * 3,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    q = qkv[:, : heads * _HEAD_DIM].view(seq_len, heads, _HEAD_DIM)
+    k = qkv[:, heads * _HEAD_DIM : 2 * heads * _HEAD_DIM].view(
+        seq_len,
+        heads,
+        _HEAD_DIM,
+    )
+    q_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    k_weight = torch.randn(_HEAD_DIM, device="cuda", dtype=torch.bfloat16)
+    freqs = torch.randn(seq_len, _ROTARY_DIM // 2, device="cuda")
+    rope_table = torch.cat((torch.cos(freqs), torch.sin(freqs)), dim=-1).to(torch.bfloat16)
+
+    expected_q, expected_k = _reference(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        rope_table,
+    )
+    actual_q, actual_k = fused_qk_norm_rope(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        rope_table,
+        _EPS,
+    )
+
+    torch.testing.assert_close(actual_q, expected_q, atol=0.0625, rtol=0.02)
+    torch.testing.assert_close(actual_k, expected_k, atol=0.0625, rtol=0.02)