AliciaandHongsheng Liu
d02681bf01
Signed-off-by: congw729 <115451386+congw729@users.noreply.github.com> Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Unit tests for the CustomOp platform-dispatch base class."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.layers.custom_op import CustomOp
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
|
|
|
|
|
class _NativeOnlyOp(CustomOp):
|
|
"""An op that implements only the PyTorch-native path."""
|
|
|
|
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
|
|
return x * 2
|
|
|
|
|
|
def test_forward_xpu_falls_back_to_native():
|
|
"""XPU has no bespoke kernels, so the base class must reuse forward_native.
|
|
|
|
Without this fallback every CustomOp subclass has to define an identical
|
|
forward_xpu, and any that forgets crashes with NotImplementedError on XPU.
|
|
"""
|
|
op = _NativeOnlyOp()
|
|
x = torch.randn(4, 8)
|
|
|
|
assert torch.equal(op.forward_xpu(x), op.forward_native(x))
|
|
|
|
|
|
def test_forward_cuda_still_requires_an_implementation():
|
|
"""The native fallback must not mask a missing CUDA kernel."""
|
|
op = _NativeOnlyOp()
|
|
|
|
with pytest.raises(NotImplementedError):
|
|
op.forward_cuda(torch.randn(4, 8))
|
|
|
|
|
|
def test_mot_rmsnorm_dispatches_on_xpu_without_an_override():
|
|
"""MoTRMSNorm relies on the inherited fallback (regression test)."""
|
|
from vllm_omni.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
|
|
|
|
assert "forward_xpu" not in MoTRMSNorm.__dict__
|
|
|
|
norm = MoTRMSNorm(16)
|
|
x = torch.randn(8, 16)
|
|
text_indices = torch.arange(0, 4)
|
|
vae_indices = torch.arange(4, 8)
|
|
|
|
assert torch.equal(norm.forward_xpu(x), norm.forward_native(x))
|
|
assert torch.equal(
|
|
norm.forward_xpu(x, text_indices, vae_indices),
|
|
norm.forward_native(x, text_indices, vae_indices),
|
|
)
|