[Perf]: Fused RMSNorm+3D Rope kernel in Sensenova-U1 (#4669)

Signed-off-by: boatman <109857087+sphinxkkkbc@users.noreply.github.com>

tests/diffusion/models/sensenova_u1/test_sensenova_triton_kernel.py

@@ -0,0 +1,232 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Correctness tests for Sensenova-U1 Triton kernels.
+
+Verifies that fused-RMSNorm+3D Rope Triton kernels
+produce results numerically equivalent to the PyTorch reference.
+All computations follow the kernel's internal float32 promotion rules.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"),
+]
+
+SEED = 42
+B = 1
+H_Q = 32
+H_K = 8
+HEAD_DIM = 128
+T_DIM = 64
+HW_DIM = 64
+HW_ROPE_DIM = 32
+QKV_DIM = (H_Q + H_K + H_K) * HEAD_DIM
+EPS = 1e-6
+DEVICE = torch.device("cuda:0")
+
+
+@dataclass(frozen=True)
+class KernelInput:
+    seq_len: int
+    q: torch.Tensor
+    k: torch.Tensor
+    q_norm_weight: torch.Tensor
+    k_norm_weight: torch.Tensor
+    q_norm_hw_weight: torch.Tensor
+    k_norm_hw_weight: torch.Tensor
+    cos_t: torch.Tensor
+    sin_t: torch.Tensor
+    cos_h: torch.Tensor
+    sin_h: torch.Tensor
+    cos_w: torch.Tensor
+    sin_w: torch.Tensor
+
+    def args(self) -> tuple[torch.Tensor, ...]:
+        return (
+            self.q,
+            self.k,
+            self.q_norm_weight,
+            self.k_norm_weight,
+            self.q_norm_hw_weight,
+            self.k_norm_hw_weight,
+            self.cos_t,
+            self.sin_t,
+            self.cos_h,
+            self.sin_h,
+            self.cos_w,
+            self.sin_w,
+        )
+
+
+try:
+    from vllm_omni.diffusion.models.sensenova_u1.fused_rmsnorm_rope import triton_qk_norm_rope
+
+    _TRITON_AVAILABLE = True
+except ImportError:
+    _TRITON_AVAILABLE = False
+
+triton_available = pytest.mark.skipif(not _TRITON_AVAILABLE, reason="Triton not available on this platform")
+
+
+# ---------------------------------------------------------------------------
+# PyTorch reference implementations (mirror the kernel's internal arithmetic)
+# ---------------------------------------------------------------------------
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rope(
+    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+
+
+def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = EPS) -> torch.Tensor:
+    input_dtype = x.dtype
+    xf = x.float()
+    var = xf.pow(2).mean(dim=-1, keepdim=True)
+    xf = xf * torch.rsqrt(var + eps)
+    return xf.to(input_dtype) * weight
+
+
+def assert_close_with_error_stats(
+    actual: torch.Tensor,
+    expected: torch.Tensor,
+    *,
+    name: str,
+    atol: float,
+    rtol: float,
+) -> None:
+    try:
+        torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)
+    except AssertionError as exc:
+        abs_diff = (actual.float() - expected.float()).abs().flatten()
+        tolerance = atol + rtol * expected.float().abs().flatten()
+        mismatch = abs_diff > tolerance
+        mismatch_count = int(mismatch.sum().item())
+        total = abs_diff.numel()
+        p99 = torch.quantile(abs_diff, 0.99).item() if total else 0.0
+        stats = (
+            f"{name} error stats: "
+            f"max={abs_diff.max().item() if total else 0.0:.6g}, "
+            f"p99={p99:.6g}, "
+            f"mean={abs_diff.mean().item() if total else 0.0:.6g}, "
+            f"mismatch={mismatch_count}/{total} ({mismatch_count / total:.4%}), "
+            f"atol={atol}, rtol={rtol}"
+        )
+        raise AssertionError(f"{stats}\n{exc}") from exc
+
+
+def reference_kernel(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_norm_weight: torch.Tensor,
+    k_norm_weight: torch.Tensor,
+    q_norm_hw_weight: torch.Tensor,
+    k_norm_hw_weight: torch.Tensor,
+    cos_t: torch.Tensor,
+    sin_t: torch.Tensor,
+    cos_h: torch.Tensor,
+    sin_h: torch.Tensor,
+    cos_w: torch.Tensor,
+    sin_w: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """PyTorch reference matching the model's current math."""
+    q_t, q_hw = q.chunk(2, dim=-1)
+    k_t, k_hw = k.chunk(2, dim=-1)
+
+    q_t = rmsnorm(q_t, q_norm_weight).transpose(1, 2)
+    k_t = rmsnorm(k_t, k_norm_weight).transpose(1, 2)
+    q_hw = rmsnorm(q_hw, q_norm_hw_weight).transpose(1, 2)
+    k_hw = rmsnorm(k_hw, k_norm_hw_weight).transpose(1, 2)
+
+    q_h, q_w = q_hw.chunk(2, dim=-1)
+    k_h, k_w = k_hw.chunk(2, dim=-1)
+
+    q_t, k_t = apply_rope(q_t, k_t, cos_t, sin_t)
+    q_h, k_h = apply_rope(q_h, k_h, cos_h, sin_h)
+    q_w, k_w = apply_rope(q_w, k_w, cos_w, sin_w)
+
+    query = torch.cat([q_t, q_h, q_w], dim=-1)
+    key = torch.cat([k_t, k_h, k_w], dim=-1)
+    return query, key
+
+
+def make_input(seq_len: int, dtype: torch.dtype, device: torch.device, seed: int) -> KernelInput:
+    gen = torch.Generator(device=device)
+    gen.manual_seed(seed + seq_len)
+
+    # Match real model layout: q/k/v are views split from fused qkv.
+    qkv = torch.randn(B, seq_len, QKV_DIM, device=device, dtype=dtype, generator=gen)
+    q, k, _v = qkv.split([H_Q * HEAD_DIM, H_K * HEAD_DIM, H_K * HEAD_DIM], dim=-1)
+    q = q.view(B, seq_len, H_Q, HEAD_DIM)
+    k = k.view(B, seq_len, H_K, HEAD_DIM)
+
+    q_norm_weight = torch.randn(T_DIM, device=device, dtype=dtype, generator=gen)
+    k_norm_weight = torch.randn(T_DIM, device=device, dtype=dtype, generator=gen)
+    q_norm_hw_weight = torch.randn(HW_DIM, device=device, dtype=dtype, generator=gen)
+    k_norm_hw_weight = torch.randn(HW_DIM, device=device, dtype=dtype, generator=gen)
+
+    # Match Qwen3RotaryEmbedding: emb = torch.cat((freqs, freqs), dim=-1).
+    def make_rope_pair(dim: int) -> tuple[torch.Tensor, torch.Tensor]:
+        cos_half = torch.randn(B, seq_len, dim // 2, device=device, dtype=dtype, generator=gen)
+        sin_half = torch.randn(B, seq_len, dim // 2, device=device, dtype=dtype, generator=gen)
+        return torch.cat((cos_half, cos_half), dim=-1), torch.cat((sin_half, sin_half), dim=-1)
+
+    cos_t, sin_t = make_rope_pair(T_DIM)
+    cos_h, sin_h = make_rope_pair(HW_ROPE_DIM)
+    cos_w, sin_w = make_rope_pair(HW_ROPE_DIM)
+
+    return KernelInput(
+        seq_len=seq_len,
+        q=q,
+        k=k,
+        q_norm_weight=q_norm_weight,
+        k_norm_weight=k_norm_weight,
+        q_norm_hw_weight=q_norm_hw_weight,
+        k_norm_hw_weight=k_norm_hw_weight,
+        cos_t=cos_t,
+        sin_t=sin_t,
+        cos_h=cos_h,
+        sin_h=sin_h,
+        cos_w=cos_w,
+        sin_w=sin_w,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Fused RMSNorm + 3D Rope
+# ---------------------------------------------------------------------------
+
+
+@triton_available
+@pytest.mark.parametrize("seq_len", [1, 2, 4, 8, 9, 16, 32, 64, 128, 256, 260, 271, 512, 1024, 2048, 4096])
+@pytest.mark.parametrize(
+    ("dtype", "atol", "rtol"),
+    [
+        (torch.float32, 1e-5, 1e-5),
+        (torch.bfloat16, 3e-2, 3e-2),
+    ],
+)
+def test_fused_qk_norm_rope_matches_reference(seq_len: int, dtype: torch.dtype, atol: float, rtol: float):
+    data = make_input(seq_len, dtype=dtype, device=DEVICE, seed=SEED)
+    call_args = data.args()
+    ref_q, ref_k = reference_kernel(*call_args)
+    out_q, out_k = triton_qk_norm_rope(*call_args, EPS)
+
+    assert_close_with_error_stats(out_q, ref_q, name="query", atol=atol, rtol=rtol)
+    assert_close_with_error_stats(out_k, ref_k, name="key", atol=atol, rtol=rtol)