XiongZhongwei
40dcfdbb43
Signed-off-by: MarkPoloChina <1228479724@qq.com> Signed-off-by: XiongZhongwei <1228479724@qq.com>
35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.layers.norm import RMSNorm
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
pytestmark = [
|
|
pytest.mark.core_model,
|
|
pytest.mark.diffusion,
|
|
pytest.mark.cuda,
|
|
pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="CUDA platform required"),
|
|
]
|
|
|
|
|
|
def test_rmsnorm_cuda_residual_matches_native_contract() -> None:
|
|
eps = 1e-6
|
|
device = torch.device(current_omni_platform.device_type)
|
|
torch.manual_seed(0)
|
|
norm = RMSNorm(64, eps=eps, dtype=torch.bfloat16).to(device)
|
|
x = torch.randn(2, 4, 64, device=device, dtype=torch.bfloat16)
|
|
residual = torch.randn_like(x)
|
|
x_before = x.clone()
|
|
residual_before = residual.clone()
|
|
|
|
expected_output, expected_residual = norm.forward_native(x_before, residual_before)
|
|
output, updated_residual = norm.forward_cuda(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)
|
|
torch.testing.assert_close(x, x_before, atol=0, rtol=0)
|
|
torch.testing.assert_close(residual, residual_before, atol=0, rtol=0)
|