diff --git a/tests/diffusion/layers/test_custom_op.py b/tests/diffusion/layers/test_custom_op.py new file mode 100644 index 000000000..9b4913730 --- /dev/null +++ b/tests/diffusion/layers/test_custom_op.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM 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.layers.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), + ) diff --git a/vllm_omni/diffusion/layers/activation.py b/vllm_omni/diffusion/layers/activation.py index 97a5a60fc..761a681fb 100644 --- a/vllm_omni/diffusion/layers/activation.py +++ b/vllm_omni/diffusion/layers/activation.py @@ -35,6 +35,3 @@ class SiluAndMul(CustomOp): import torch_npu return torch_npu.npu_swiglu(x, dim=-1) - - def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: - return self.forward_native(x) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index d147bdcfe..5052807b5 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -70,14 +70,6 @@ class AdaLayerNorm(CustomOp): return output - def forward_xpu( - self, - x: torch.Tensor, - scale: torch.Tensor, - shift: torch.Tensor, - ) -> torch.Tensor: - return self.forward_native(x, scale, shift) - def forward_native( self, x: torch.Tensor, diff --git a/vllm_omni/diffusion/layers/custom_op.py b/vllm_omni/diffusion/layers/custom_op.py index 27e3bce1f..38a6a2891 100644 --- a/vllm_omni/diffusion/layers/custom_op.py +++ b/vllm_omni/diffusion/layers/custom_op.py @@ -48,7 +48,9 @@ class CustomOp(nn.Module): raise NotImplementedError def forward_xpu(self, *args, **kwargs): - raise NotImplementedError + # By default, we assume that XPU ops are compatible with the + # PyTorch-native implementation. + return self.forward_native(*args, **kwargs) def forward_hip(self, *args, **kwargs): # By default, we assume that HIP ops are compatible with CUDA ops. diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py index 3618388e6..d9e116ec7 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/vllm_omni/diffusion/layers/norm.py @@ -36,9 +36,6 @@ class LayerNorm(nn.LayerNorm, CustomOp): def forward_hip(self, x: torch.Tensor) -> torch.Tensor: return self.forward_native(x) - def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: - return self.forward_native(x) - def forward_npu(self, x: torch.Tensor) -> torch.Tensor: if _HAS_MINDIESD: try: @@ -151,15 +148,6 @@ class RMSNorm(CustomOp): return output - def forward_xpu( - self, - x: torch.Tensor, - residual: torch.Tensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - if residual is None: - return self.forward_native(x) - return self.forward_native(x, residual) - def forward_native( self, x: torch.Tensor, @@ -231,12 +219,6 @@ class RMSNormVAE(CustomOp): out = out + self.bias return out - def forward_xpu( - self, - x: torch.Tensor, - ) -> torch.Tensor: - return self.forward_native(x) - def forward_native( self, x: torch.Tensor, diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index e78574b04..59d2c9a62 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -179,14 +179,6 @@ class RotaryEmbedding(CustomOp): else: return self.forward_native(x, cos, sin) - def forward_xpu( - self, - x: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - ) -> torch.Tensor: - return self.forward_native(x, cos, sin) - def forward_musa( self, x: torch.Tensor, diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 3a47e2912..0308cda7b 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -72,9 +72,6 @@ class _RMSNorm(CustomOp): def forward_cuda(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.forward_native(hidden_states) - def forward_xpu(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.forward_native(hidden_states) - def forward_native(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32)