# 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 sys
import types

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)


def test_rmsnorm_dtype_matches_parameter_and_fp32_accumulation_contract():
    """dtype controls gamma storage; reduced-precision inputs still use fp32 accumulation."""
    from vllm_omni.diffusion.layers.norm import RMSNorm

    hidden_size = 64
    eps = 1e-6
    norm = RMSNorm(hidden_size, eps=eps, dtype=torch.bfloat16)
    x = torch.randn(2, 4, hidden_size, dtype=torch.bfloat16)

    assert norm.weight.dtype == torch.bfloat16

    x_fp32 = x.float()
    expected = (x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + eps) * norm.weight.float()).to(x.dtype)

    torch.testing.assert_close(norm.forward_native(x), expected, atol=0, rtol=0)


def test_rmsnorm_npu_receives_gamma_with_configured_dtype(monkeypatch: pytest.MonkeyPatch):
    """The NPU fused op receives the bf16 checkpoint gamma used by H3."""
    from vllm_omni.diffusion.layers.norm import RMSNorm

    captured: dict[str, torch.Tensor] = {}

    def npu_rms_norm(x: torch.Tensor, gamma: torch.Tensor, epsilon: float):
        captured["gamma"] = gamma
        captured["epsilon"] = torch.tensor(epsilon)
        return (x,)

    monkeypatch.setitem(sys.modules, "torch_npu", types.SimpleNamespace(npu_rms_norm=npu_rms_norm))
    norm = RMSNorm(64, eps=1e-5, dtype=torch.bfloat16)
    x = torch.randn(2, 4, 64, dtype=torch.bfloat16)

    assert norm.forward_npu(x) is x
    assert captured["gamma"] is norm.weight
    assert captured["gamma"].dtype == torch.bfloat16
    assert captured["epsilon"].item() == pytest.approx(1e-5)


# ── RMSNorm residual tests ──


def test_rmsnorm_matches_native_residual_contract():
    """The non-NPU path preserves separate residual-add and RMSNorm math."""
    from vllm_omni.diffusion.layers.norm import RMSNorm

    eps = 1e-6
    norm = RMSNorm(4, eps=eps, dtype=torch.bfloat16)
    norm.weight.data.copy_(torch.tensor([1.0, 0.5, 1.5, 2.0], dtype=torch.bfloat16))
    x = torch.tensor([[1.0, -2.0, 3.0, -4.0]], dtype=torch.bfloat16)
    residual = torch.tensor([[0.00390625, 0.0078125, -0.0078125, 0.015625]], dtype=torch.bfloat16)

    expected_residual = residual + x
    expected_fp32 = expected_residual.float()
    expected_output = expected_fp32 * torch.rsqrt(expected_fp32.square().mean(-1, keepdim=True) + eps)
    expected_output = (expected_output * norm.weight.float()).to(x.dtype)
    output, updated_residual = norm.forward_native(x, residual)

    torch.testing.assert_close(updated_residual, expected_residual, atol=0, rtol=0)
    torch.testing.assert_close(output, expected_output, atol=0, rtol=0)


@pytest.mark.parametrize(
    "forward_name",
    ["forward_cuda", "forward_hip", "forward_musa", "forward_xpu"],
)
def test_rmsnorm_non_npu_residual_delegates_to_native(
    monkeypatch: pytest.MonkeyPatch,
    forward_name: str,
):
    """Platform wrappers keep residual arithmetic in forward_native."""
    from vllm_omni.diffusion.layers.norm import RMSNorm

    norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
    x = torch.randn(2, 4, dtype=torch.bfloat16)
    residual = torch.randn_like(x)
    expected_output = torch.randn_like(x)
    expected_residual = torch.randn_like(x)
    captured: dict[str, object] = {}

    def forward_native(native_x: torch.Tensor, native_residual: torch.Tensor | None = None):
        captured.update(x=native_x, residual=native_residual)
        return expected_output, expected_residual

    monkeypatch.setattr(norm, "forward_native", forward_native)

    output, updated_residual = getattr(norm, forward_name)(x, residual)

    assert output is expected_output
    assert updated_residual is expected_residual
    assert captured["x"] is x
    assert captured["residual"] is residual


def test_rmsnorm_npu_residual_uses_torch_npu_fused_op(monkeypatch: pytest.MonkeyPatch):
    """NPU keeps the direct fused AddRMSNorm path instead of IR fallback."""
    from vllm_omni.diffusion.layers.norm import RMSNorm

    captured: dict[str, object] = {}
    fused_output = torch.randn(2, 4, dtype=torch.bfloat16)
    fused_residual = torch.randn(2, 4, dtype=torch.bfloat16)

    def npu_add_rms_norm(
        x: torch.Tensor,
        residual: torch.Tensor,
        gamma: torch.Tensor,
        epsilon: float,
    ):
        captured.update(x=x, residual=residual, gamma=gamma, epsilon=epsilon)
        return fused_output, torch.empty(0), fused_residual

    monkeypatch.setitem(
        sys.modules,
        "torch_npu",
        types.SimpleNamespace(npu_add_rms_norm=npu_add_rms_norm),
    )
    norm = RMSNorm(4, eps=1e-5, dtype=torch.bfloat16)
    x = torch.randn(2, 4, dtype=torch.bfloat16)
    residual = torch.randn(2, 4, dtype=torch.bfloat16)

    output, updated_residual = norm.forward_npu(x, residual)

    assert output is fused_output
    assert updated_residual is fused_residual
    assert captured["x"] is x
    assert captured["residual"] is residual
    assert captured["gamma"] is norm.weight
    assert captured["epsilon"] == 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])