[Bugfix][XPU] Add MoTRMSNorm.forward_xpu delegating to forward_native (#5568)

Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>
Signed-off-by: Joshna-Medisetty <joshna.medisetty@intel.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

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),
+    )

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)

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,

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.

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,

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,

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)