vllm-project--vllm-omni
wehub-resource-sync
613 行
17 KiB
Python
613 行
17 KiB
Python
# 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
|
|
from pytest_mock import MockerFixture
|
|
|
|
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)
|
|
|
|
|
|
# ── RMSNorm compile-path regression tests ──
|
|
|
|
|
|
def test_rmsnorm_forward_cuda_does_not_call_fused_during_compile(mocker: MockerFixture) -> None:
|
|
"""Regression: _forward_fused must not be called during torch.compile tracing.
|
|
|
|
Under HSDP, RMSNorm.weight is a DTensor. Accessing .data on a DTensor inside
|
|
_forward_fused during tracing produces an orphan all-gather node outside the
|
|
compile boundary, causing inductor's compute_ancestors to raise KeyError.
|
|
The is_compiling() guard in forward_cuda/forward_hip prevents this by routing
|
|
to forward_native during tracing.
|
|
|
|
If someone removes the guard, this test will catch the regression by asserting
|
|
that _forward_fused was not called while is_compiling() returns True.
|
|
"""
|
|
from vllm_omni.diffusion.layers.norm import RMSNorm
|
|
|
|
norm = RMSNorm(hidden_size=64)
|
|
x = torch.randn(2, 4, 64)
|
|
|
|
mock_fused = mocker.patch.object(norm, "_forward_fused", wraps=norm._forward_fused)
|
|
mocker.patch("torch.compiler.is_compiling", return_value=True)
|
|
|
|
out = norm.forward_cuda(x)
|
|
|
|
mock_fused.assert_not_called()
|
|
assert out.shape == x.shape
|
|
|
|
|
|
def test_rmsnorm_forward_hip_does_not_call_fused_during_compile(mocker: MockerFixture) -> None:
|
|
"""Regression: same guard must be present in forward_hip.
|
|
|
|
forward_hip is the entry point on ROCm (AMD GPU). It must behave identically
|
|
to forward_cuda with respect to the is_compiling() guard.
|
|
"""
|
|
from vllm_omni.diffusion.layers.norm import RMSNorm
|
|
|
|
norm = RMSNorm(hidden_size=64)
|
|
x = torch.randn(2, 4, 64)
|
|
|
|
mock_fused = mocker.patch.object(norm, "_forward_fused", wraps=norm._forward_fused)
|
|
mocker.patch("torch.compiler.is_compiling", return_value=True)
|
|
|
|
out = norm.forward_hip(x)
|
|
|
|
mock_fused.assert_not_called()
|
|
assert out.shape == x.shape
|
|
|
|
|
|
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)
|
|
|
|
|
|
# ── Fused vs Native correctness tests ──
|
|
|
|
|
|
def test_rmsnorm_fused_native_parity_cuda():
|
|
"""Verify fused and native RMSNorm produce identical results on CUDA.
|
|
|
|
The fused kernel uses vllm._custom_ops.rms_norm while native uses pure PyTorch.
|
|
Both should produce numerically equivalent results.
|
|
"""
|
|
from vllm_omni.diffusion.layers.norm import RMSNorm
|
|
|
|
hidden_size = 128
|
|
eps = 1e-6
|
|
torch.manual_seed(42)
|
|
|
|
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
|
torch.manual_seed(42)
|
|
norm = RMSNorm(hidden_size, eps=eps).cuda()
|
|
norm.weight.data = norm.weight.data.to(dtype)
|
|
|
|
x = torch.randn(4, 8, hidden_size, dtype=dtype, device="cuda")
|
|
|
|
out_fused = norm._forward_fused(x)
|
|
out_native = norm.forward_native(x)
|
|
|
|
torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_rmsnorm_fused_native_parity_large_tensor():
|
|
"""Verify fused and native RMSNorm parity with large tensors.
|
|
|
|
Tests realistic tensor sizes used in diffusion models.
|
|
"""
|
|
from vllm_omni.diffusion.layers.norm import RMSNorm
|
|
|
|
hidden_size = 4096
|
|
batch_size = 16
|
|
seq_len = 1024
|
|
eps = 1e-6
|
|
|
|
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
|
torch.manual_seed(123)
|
|
norm = RMSNorm(hidden_size, eps=eps).cuda()
|
|
norm.weight.data = norm.weight.data.to(dtype)
|
|
|
|
x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda")
|
|
|
|
out_fused = norm._forward_fused(x)
|
|
out_native = norm.forward_native(x)
|
|
|
|
torch.testing.assert_close(out_fused, out_native, atol=1e-2, rtol=1e-2)
|
|
|
|
|
|
def test_rmsnorm_fused_native_parity_2d_input():
|
|
"""Verify fused and native RMSNorm parity with 2D input tensors."""
|
|
from vllm_omni.diffusion.layers.norm import RMSNorm
|
|
|
|
hidden_size = 256
|
|
batch_size = 32
|
|
eps = 1e-6
|
|
|
|
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
|
torch.manual_seed(456)
|
|
norm = RMSNorm(hidden_size, eps=eps).cuda()
|
|
norm.weight.data = norm.weight.data.to(dtype)
|
|
|
|
x = torch.randn(batch_size, hidden_size, dtype=dtype, device="cuda")
|
|
|
|
out_fused = norm._forward_fused(x)
|
|
out_native = norm.forward_native(x)
|
|
|
|
torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_rmsnorm_fused_native_parity_custom_weight():
|
|
"""Verify fused and native RMSNorm parity with non-default weights.
|
|
|
|
Note: bfloat16 requires wider tolerance because the fused CUDA kernel
|
|
(vllm._custom_ops.rms_norm) computes in bf16 precision internally,
|
|
while forward_native uses float32 intermediates. With random custom
|
|
weights (not the default ones-weight), this precision gap exceeds 1e-3.
|
|
|
|
The gap surfaced after upstream vLLM routed RMSNorm through the vLLM IR
|
|
op (ir.ops.rms_norm) — see upstream commit 40bb17502
|
|
"[vLLM IR] 1/N Implement IR skeleton and rms_norm op (#33825)" — which
|
|
changed which fused kernel/precision path the layer dispatches to.
|
|
"""
|
|
from vllm_omni.diffusion.layers.norm import RMSNorm
|
|
|
|
hidden_size = 64
|
|
eps = 1e-6
|
|
|
|
# Per-dtype tolerances: bf16 fused kernel has inherent precision limit
|
|
atol_map = {torch.float32: 1e-3, torch.float16: 1e-3, torch.bfloat16: 2e-2}
|
|
rtol_map = {torch.float32: 1e-3, torch.float16: 1e-3, torch.bfloat16: 1e-2}
|
|
|
|
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
|
torch.manual_seed(789)
|
|
norm = RMSNorm(hidden_size, eps=eps).cuda()
|
|
norm.weight.data = torch.randn(hidden_size, dtype=dtype, device="cuda")
|
|
|
|
x = torch.randn(2, 4, hidden_size, dtype=dtype, device="cuda")
|
|
|
|
out_fused = norm._forward_fused(x)
|
|
out_native = norm.forward_native(x)
|
|
|
|
torch.testing.assert_close(out_fused, out_native, atol=atol_map[dtype], rtol=rtol_map[dtype])
|