Weiming Liaoandwangyu
35a83a0885
Signed-off-by: Weiming Liao <liaowm5@gmail.com> Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com>
41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from tests.helpers.mark import hardware_marks
|
|
from vllm_omni.diffusion.layers.activation import SiluAndMul
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_silu_and_mul_native_matches_packed_reference() -> None:
|
|
packed = torch.randn(257, 256, dtype=torch.float32)
|
|
gate, up = packed.chunk(2, dim=-1)
|
|
|
|
torch.testing.assert_close(SiluAndMul().forward_native(packed), F.silu(gate) * up)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"dtype",
|
|
[
|
|
pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A2"}, num_cards=1), id="A2-bf16"),
|
|
pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A3"}, num_cards=1), id="A3-bf16"),
|
|
pytest.param(torch.bfloat16, marks=hardware_marks(res={"npu": "A5"}, num_cards=1), id="A5-bf16"),
|
|
pytest.param(torch.float16, marks=hardware_marks(res={"npu": "310P"}, num_cards=1), id="310P-fp16"),
|
|
],
|
|
)
|
|
def test_silu_and_mul_npu_matches_packed_reference(dtype: torch.dtype) -> None:
|
|
packed = torch.randn(257, 256, device="npu", dtype=dtype)
|
|
gate, up = packed.chunk(2, dim=-1)
|
|
|
|
torch.testing.assert_close(
|
|
SiluAndMul()(packed),
|
|
F.silu(gate) * up,
|
|
atol=2e-2,
|
|
rtol=2e-2,
|
|
)
|