[Kernel][MammothModa2] Use shared RMSNorm in the DiT (#7489)

Signed-off-by: Jerry Xu <xjr2423@gmail.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

tests/diffusion/layers/test_norm.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for LayerNorm and RMSNorm custom ops in diffusion layers."""
 
 import sys
@@ -480,6 +480,19 @@ def test_rmsnorm_forward_hip_does_not_call_fused_during_compile(mocker: MockerFi
     assert out.shape == x.shape
 
 
+def test_rmsnorm_skips_fused_kernel_for_empty_input(mocker: MockerFixture) -> None:
+    """Regression: the fused kernel fails to launch on a zero-token tensor and leaves a sticky CUDA error."""
+    from vllm_omni.diffusion.layers.norm import RMSNorm
+
+    norm = RMSNorm(hidden_size=64)
+    x = torch.randn(1, 0, 64)
+    mock_fused = mocker.patch.object(norm, "_forward_fused", wraps=norm._forward_fused)
+
+    assert norm.forward_cuda(x).shape == x.shape
+    assert norm.forward_hip(x).shape == x.shape
+    mock_fused.assert_not_called()
+
+
 def test_rmsnorm_matches_reference_implementation():
     """Verify RMSNorm matches a reference implementation."""
     from vllm_omni.diffusion.layers.norm import RMSNorm