[PERF] Wan2.2 support rmsnorm fused op (#2583)
Signed-off-by: fan2956 <zhoufan53@huawei.com> Signed-off-by: gcanlin <canlinguosdu@gmail.com> Co-authored-by: gcanlin <canlinguosdu@gmail.com>
This commit is contained in:
fan2956andgcanlin authored and GitHub committed
1 parent
4d816ff1de
commit
f1cb4ebe4c
4 files changed
+585
-10
No files matched your search
@@ -0,0 +1,453 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for LayerNorm and RMSNorm custom ops in diffusion layers."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
# ── Import tests ──
|
||||
|
||||
|
||||
def test_layernorm_import():
|
||||
"""Verify LayerNorm can be imported from the norm module."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm # noqa: F401
|
||||
|
||||
|
||||
def test_rmsnorm_import():
|
||||
"""Verify RMSNorm can be imported from the norm module."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm # noqa: F401
|
||||
|
||||
|
||||
# ── LayerNorm tests ──
|
||||
|
||||
|
||||
def test_layernorm_forward_shape():
|
||||
"""LayerNorm produces correct output shapes."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
batch = 2
|
||||
seq_len = 4
|
||||
norm = LayerNorm(dim)
|
||||
|
||||
x = torch.randn(batch, seq_len, dim)
|
||||
out = norm(x)
|
||||
|
||||
assert out.shape == (batch, seq_len, dim)
|
||||
|
||||
|
||||
def test_layernorm_forward_shape_2d():
|
||||
"""LayerNorm works with 2D input tensors."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
batch = 2
|
||||
norm = LayerNorm(dim)
|
||||
|
||||
x = torch.randn(batch, dim)
|
||||
out = norm(x)
|
||||
|
||||
assert out.shape == (batch, dim)
|
||||
|
||||
|
||||
def test_layernorm_preserves_dtype_fp32():
|
||||
"""LayerNorm preserves float32 dtype."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
norm = LayerNorm(dim)
|
||||
|
||||
x = torch.randn(2, 4, dim, dtype=torch.float32)
|
||||
out = norm(x)
|
||||
|
||||
assert out.dtype == torch.float32
|
||||
|
||||
|
||||
def test_layernorm_preserves_dtype_fp16():
|
||||
"""LayerNorm preserves float16 dtype."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
norm = LayerNorm(dim)
|
||||
|
||||
x = torch.randn(2, 4, dim, dtype=torch.float16)
|
||||
out = norm(x)
|
||||
|
||||
assert out.dtype == torch.float16
|
||||
|
||||
|
||||
def test_layernorm_preserves_dtype_bf16():
|
||||
"""LayerNorm preserves bfloat16 dtype."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
norm = LayerNorm(dim)
|
||||
|
||||
x = torch.randn(2, 4, dim, dtype=torch.bfloat16)
|
||||
out = norm(x)
|
||||
|
||||
assert out.dtype == torch.bfloat16
|
||||
|
||||
|
||||
def test_layernorm_without_elementwise_affine():
|
||||
"""LayerNorm works without elementwise_affine (no learned parameters)."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
norm = LayerNorm(dim, elementwise_affine=False)
|
||||
|
||||
assert norm.weight is None
|
||||
assert norm.bias is None
|
||||
|
||||
x = torch.randn(2, 4, dim)
|
||||
out = norm(x)
|
||||
|
||||
assert out.shape == (2, 4, dim)
|
||||
|
||||
|
||||
def test_layernorm_custom_eps():
|
||||
"""LayerNorm accepts custom epsilon value."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
eps = 1e-5
|
||||
norm = LayerNorm(dim, eps=eps)
|
||||
|
||||
assert norm.eps == eps
|
||||
|
||||
|
||||
def test_layernorm_has_learnable_parameters():
|
||||
"""LayerNorm has learnable weight and bias by default."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
norm = LayerNorm(dim)
|
||||
|
||||
assert norm.weight is not None
|
||||
assert norm.bias is not None
|
||||
assert norm.weight.shape == (dim,)
|
||||
assert norm.bias.shape == (dim,)
|
||||
|
||||
|
||||
def test_layernorm_matches_fp32_reference():
|
||||
"""Verify LayerNorm produces identical output to FP32 nn.LayerNorm."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
eps = 1e-6
|
||||
torch.manual_seed(42)
|
||||
|
||||
ours = LayerNorm(dim, eps=eps)
|
||||
ref = torch.nn.LayerNorm(dim, eps=eps)
|
||||
|
||||
# Copy weights
|
||||
ref.weight.data.copy_(ours.weight.data)
|
||||
ref.bias.data.copy_(ours.bias.data)
|
||||
|
||||
x = torch.randn(2, 4, dim)
|
||||
|
||||
out_ours = ours(x)
|
||||
out_ref = ref(x.float()).to(x.dtype)
|
||||
|
||||
torch.testing.assert_close(out_ours, out_ref, atol=1e-5, rtol=1e-5)
|
||||
|
||||
|
||||
def test_layernorm_matches_diffusers_fp32layernorm():
|
||||
"""Verify LayerNorm produces identical output to diffusers FP32LayerNorm."""
|
||||
from diffusers.models.normalization import FP32LayerNorm
|
||||
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
eps = 1e-6
|
||||
torch.manual_seed(42)
|
||||
|
||||
ours = LayerNorm(dim, eps=eps)
|
||||
ref = FP32LayerNorm(dim, eps=eps)
|
||||
|
||||
# Copy weights
|
||||
ref.weight.data.copy_(ours.weight.data)
|
||||
ref.bias.data.copy_(ours.bias.data)
|
||||
|
||||
# Test with fp16 input to verify FP32 computation
|
||||
x = torch.randn(2, 4, dim, dtype=torch.float16)
|
||||
|
||||
out_ours = ours(x)
|
||||
out_ref = ref(x)
|
||||
|
||||
torch.testing.assert_close(out_ours, out_ref, atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
# ── RMSNorm tests ──
|
||||
|
||||
|
||||
def test_rmsnorm_forward_shape():
|
||||
"""RMSNorm produces correct output shapes."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
hidden_size = 64
|
||||
batch = 2
|
||||
seq_len = 4
|
||||
norm = RMSNorm(hidden_size)
|
||||
|
||||
x = torch.randn(batch, seq_len, hidden_size)
|
||||
out = norm(x)
|
||||
|
||||
assert out.shape == (batch, seq_len, hidden_size)
|
||||
|
||||
|
||||
def test_rmsnorm_forward_shape_2d():
|
||||
"""RMSNorm works with 2D input tensors."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
hidden_size = 64
|
||||
batch = 2
|
||||
norm = RMSNorm(hidden_size)
|
||||
|
||||
x = torch.randn(batch, hidden_size)
|
||||
out = norm(x)
|
||||
|
||||
assert out.shape == (batch, hidden_size)
|
||||
|
||||
|
||||
def test_rmsnorm_preserves_dtype_fp32():
|
||||
"""RMSNorm preserves float32 dtype."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
hidden_size = 64
|
||||
norm = RMSNorm(hidden_size)
|
||||
|
||||
x = torch.randn(2, 4, hidden_size, dtype=torch.float32)
|
||||
out = norm(x)
|
||||
|
||||
assert out.dtype == torch.float32
|
||||
|
||||
|
||||
def test_rmsnorm_preserves_dtype_fp16():
|
||||
"""RMSNorm preserves float16 dtype."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
hidden_size = 64
|
||||
norm = RMSNorm(hidden_size)
|
||||
|
||||
x = torch.randn(2, 4, hidden_size, dtype=torch.float16)
|
||||
out = norm(x)
|
||||
|
||||
assert out.dtype == torch.float16
|
||||
|
||||
|
||||
def test_rmsnorm_preserves_dtype_bf16():
|
||||
"""RMSNorm preserves bfloat16 dtype."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
hidden_size = 64
|
||||
norm = RMSNorm(hidden_size)
|
||||
|
||||
x = torch.randn(2, 4, hidden_size, dtype=torch.bfloat16)
|
||||
out = norm(x)
|
||||
|
||||
assert out.dtype == torch.bfloat16
|
||||
|
||||
|
||||
def test_rmsnorm_custom_eps():
|
||||
"""RMSNorm accepts custom epsilon value."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
hidden_size = 64
|
||||
eps = 1e-5
|
||||
norm = RMSNorm(hidden_size, eps=eps)
|
||||
|
||||
assert norm.variance_epsilon == eps
|
||||
|
||||
|
||||
def test_rmsnorm_has_weight_parameter():
|
||||
"""RMSNorm has learnable weight parameter initialized to ones."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
hidden_size = 64
|
||||
norm = RMSNorm(hidden_size)
|
||||
|
||||
assert norm.weight is not None
|
||||
assert norm.weight.shape == (hidden_size,)
|
||||
torch.testing.assert_close(norm.weight, torch.ones(hidden_size))
|
||||
|
||||
|
||||
def test_rmsnorm_numerical_correctness():
|
||||
"""Verify RMSNorm produces numerically correct output."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
hidden_size = 64
|
||||
eps = 1e-6
|
||||
torch.manual_seed(42)
|
||||
|
||||
norm = RMSNorm(hidden_size, eps=eps)
|
||||
x = torch.randn(2, 4, hidden_size)
|
||||
|
||||
# Compute expected output manually
|
||||
x_fp32 = x.to(torch.float32)
|
||||
variance = x_fp32.pow(2).mean(-1, keepdim=True)
|
||||
expected = x_fp32 * torch.rsqrt(variance + eps)
|
||||
expected = norm.weight.to(torch.float32) * expected
|
||||
expected = expected.to(x.dtype)
|
||||
|
||||
out = norm(x)
|
||||
|
||||
torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)
|
||||
|
||||
|
||||
def test_rmsnorm_matches_reference_implementation():
|
||||
"""Verify RMSNorm matches a reference implementation."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
def reference_rmsnorm(x, weight, eps):
|
||||
"""Reference RMSNorm implementation."""
|
||||
input_dtype = x.dtype
|
||||
x = x.to(torch.float32)
|
||||
variance = x.pow(2).mean(-1, keepdim=True)
|
||||
out = x * torch.rsqrt(variance + eps)
|
||||
out = weight.to(torch.float32) * out
|
||||
return out.to(input_dtype)
|
||||
|
||||
hidden_size = 128
|
||||
eps = 1e-6
|
||||
torch.manual_seed(123)
|
||||
|
||||
norm = RMSNorm(hidden_size, eps=eps)
|
||||
|
||||
# Test with various dtypes
|
||||
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
||||
x = torch.randn(4, 8, hidden_size, dtype=dtype)
|
||||
expected = reference_rmsnorm(x, norm.weight, eps)
|
||||
out = norm(x)
|
||||
torch.testing.assert_close(out, expected, atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
# ── CustomOp dispatch tests ──
|
||||
|
||||
|
||||
def test_layernorm_inherits_from_customop():
|
||||
"""LayerNorm inherits from CustomOp for platform dispatch."""
|
||||
from vllm_omni.diffusion.layers.custom_op import CustomOp
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
norm = LayerNorm(64)
|
||||
assert isinstance(norm, CustomOp)
|
||||
|
||||
|
||||
def test_rmsnorm_inherits_from_customop():
|
||||
"""RMSNorm inherits from CustomOp for platform dispatch."""
|
||||
from vllm_omni.diffusion.layers.custom_op import CustomOp
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
norm = RMSNorm(64)
|
||||
assert isinstance(norm, CustomOp)
|
||||
|
||||
|
||||
def test_layernorm_has_platform_methods():
|
||||
"""LayerNorm has forward methods for each platform."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
norm = LayerNorm(64)
|
||||
|
||||
assert hasattr(norm, "forward_cuda")
|
||||
assert hasattr(norm, "forward_hip")
|
||||
assert hasattr(norm, "forward_xpu")
|
||||
assert hasattr(norm, "forward_npu")
|
||||
assert hasattr(norm, "forward_native")
|
||||
|
||||
|
||||
def test_rmsnorm_has_platform_methods():
|
||||
"""RMSNorm has forward methods for each platform."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
norm = RMSNorm(64)
|
||||
|
||||
assert hasattr(norm, "forward_cuda")
|
||||
assert hasattr(norm, "forward_hip")
|
||||
assert hasattr(norm, "forward_xpu")
|
||||
assert hasattr(norm, "forward_npu")
|
||||
assert hasattr(norm, "forward_native")
|
||||
|
||||
|
||||
def test_layernorm_forward_native_directly():
|
||||
"""LayerNorm.forward_native can be called directly."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
norm = LayerNorm(dim)
|
||||
x = torch.randn(2, 4, dim)
|
||||
|
||||
out = norm.forward_native(x)
|
||||
|
||||
assert out.shape == (2, 4, dim)
|
||||
|
||||
|
||||
def test_rmsnorm_forward_native_directly():
|
||||
"""RMSNorm.forward_native can be called directly."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
hidden_size = 64
|
||||
norm = RMSNorm(hidden_size)
|
||||
x = torch.randn(2, 4, hidden_size)
|
||||
|
||||
out = norm.forward_native(x)
|
||||
|
||||
assert out.shape == (2, 4, hidden_size)
|
||||
|
||||
|
||||
# ── Edge case tests ──
|
||||
|
||||
|
||||
def test_layernorm_with_large_dim():
|
||||
"""LayerNorm works with large hidden dimensions."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 4096
|
||||
norm = LayerNorm(dim)
|
||||
x = torch.randn(1, 16, dim)
|
||||
|
||||
out = norm(x)
|
||||
|
||||
assert out.shape == (1, 16, dim)
|
||||
|
||||
|
||||
def test_rmsnorm_with_large_dim():
|
||||
"""RMSNorm works with large hidden dimensions."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
hidden_size = 4096
|
||||
norm = RMSNorm(hidden_size)
|
||||
x = torch.randn(1, 16, hidden_size)
|
||||
|
||||
out = norm(x)
|
||||
|
||||
assert out.shape == (1, 16, hidden_size)
|
||||
|
||||
|
||||
def test_layernorm_with_single_element_batch():
|
||||
"""LayerNorm works with batch size of 1."""
|
||||
from vllm_omni.diffusion.layers.norm import LayerNorm
|
||||
|
||||
dim = 64
|
||||
norm = LayerNorm(dim)
|
||||
x = torch.randn(1, 1, dim)
|
||||
|
||||
out = norm(x)
|
||||
|
||||
assert out.shape == (1, 1, dim)
|
||||
|
||||
|
||||
def test_rmsnorm_with_single_element_batch():
|
||||
"""RMSNorm works with batch size of 1."""
|
||||
from vllm_omni.diffusion.layers.norm import RMSNorm
|
||||
|
||||
hidden_size = 64
|
||||
norm = RMSNorm(hidden_size)
|
||||
x = torch.randn(1, 1, hidden_size)
|
||||
|
||||
out = norm(x)
|
||||
|
||||
assert out.shape == (1, 1, hidden_size)
|
||||
Reference in new issue
Block a user