233 lines
7.4 KiB
Python
233 lines
7.4 KiB
Python
# 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)
|