[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

tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py

@@ -54,8 +54,9 @@ def test_real_shape_matches_previous_arithmetic_bf16(seq):
 
 def test_empty_text_stream_on_the_default_backend():
     """The recipe's text-to-image request with text_guidance_scale > 1 runs the
-    context refiner on a zero-token unconditional prompt. Before the guard this
-    raised ``RuntimeError: step must be nonzero`` from the FA varlen fallback."""
+    context refiner on a zero-token unconditional prompt. Before the guards this
+    raised ``RuntimeError: step must be nonzero`` from the FA varlen fallback and
+    left a sticky CUDA error from the fused RMSNorm kernel."""
     torch.manual_seed(0)
     block = (
         TransformerBlock(DIM, HEADS, KV_HEADS, multiple_of=256, ffn_dim_multiplier=1.0, norm_eps=1e-5, modulation=False)
@@ -67,13 +68,10 @@ def test_empty_text_stream_on_the_default_backend():
     mask = torch.ones(1, 0, dtype=torch.bool, device="cuda")
     angles = torch.rand(1, 0, block.head_dim, device="cuda")
     with torch.no_grad():
-        out = block.attn(
-            hidden_states=hidden,
-            encoder_hidden_states=hidden,
-            attention_mask=mask,
-            image_rotary_emb=(angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16)),
-        )
+        out = block(hidden, mask, (angles.cos().to(torch.bfloat16), angles.sin().to(torch.bfloat16)))
     assert out.shape == (1, 0, DIM)
+    # A failed kernel launch is reported asynchronously; this device-to-host copy surfaces it.
+    torch.randn(4, device="cuda").cpu()
 
 
 @pytest.mark.parametrize(

tests/diffusion/models/mammoth_moda2/test_dit_rmsnorm.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MammothModa2 DiT norms run on the shared RMSNorm instead of transformers' Qwen2RMSNorm."""
+
+from types import SimpleNamespace
+
+import pytest
+from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
+
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.layers.norm import RMSNorm
+from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import (
+    SimpleQFormerImageRefiner,
+    Transformer2DModel,
+)
+from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import MammothModa2DiTPipeline
+
+from .test_dit_attention import _SDPA_CONFIG
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_dit_norms_are_shared_rmsnorm_with_their_eps():
+    # norm_eps differs from the fixed 1e-5 of the QK norms and the reinitialized caption
+    # embedder, so a site that lost its eps shows up here.
+    with set_current_diffusion_config(_SDPA_CONFIG):
+        transformer = Transformer2DModel(
+            hidden_size=48,
+            num_layers=1,
+            num_refiner_layers=1,
+            num_attention_heads=6,
+            num_kv_heads=2,
+            multiple_of=8,
+            norm_eps=1e-6,
+            axes_dim_rope=(4, 2, 2),
+            text_feat_dim=16,
+        )
+    refiner = SimpleQFormerImageRefiner(
+        hidden_size=32, output_hidden_size=48, num_queries=4, num_layers=1, norm_eps=1e-6
+    )
+
+    norms = [m for m in (*transformer.modules(), *refiner.modules()) if isinstance(m, (RMSNorm, Qwen2RMSNorm))]
+    assert norms and all(isinstance(m, RMSNorm) for m in norms)
+
+    block = transformer.layers[0]
+    assert block.attn.norm_q.variance_epsilon == block.attn.norm_k.variance_epsilon == 1e-5
+    assert block.norm1.norm.variance_epsilon == block.norm2.variance_epsilon == 1e-6
+    assert transformer.context_refiner[0].norm1.variance_epsilon == 1e-6
+    assert transformer.time_caption_embed.caption_embedder[0].variance_epsilon == 1e-6
+    assert refiner.layers[0]["ln_ffn"].variance_epsilon == 1e-6
+
+    MammothModa2DiTPipeline._reinit_caption_embedder(SimpleNamespace(gen_transformer=transformer), 32)
+    caption_norm = transformer.time_caption_embed.caption_embedder[0]
+    assert isinstance(caption_norm, RMSNorm) and caption_norm.variance_epsilon == 1e-5
+
+
+def test_shared_rmsnorm_loads_qwen2_rmsnorm_weights():
+    RMSNorm(120, eps=1e-5).load_state_dict(Qwen2RMSNorm(120, eps=1e-5).state_dict())

tests/diffusion/models/mammoth_moda2/test_dit_rmsnorm_cuda.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""MammothModa2 DiT norms on the fused shared RMSNorm with real Preview shapes."""
+
+import pytest
+import torch
+from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
+
+from tests.helpers.mark import hardware_marks
+from vllm_omni.diffusion.layers.norm import RMSNorm
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.diffusion,
+    pytest.mark.cuda,
+    *hardware_marks(res={"cuda": "L4"}, num_cards=1),
+]
+
+# MammothModa2-Preview at 1024x1024: 77 text + 4096 image tokens, hidden 2520, 21 heads of 120.
+SEQ = 77 + 4096
+DIM, HEADS = 2520, 21
+EPS = 1e-5
+
+_SHAPES = [(2, SEQ, DIM), (2, SEQ, HEADS, DIM // HEADS)]
+_SHAPE_IDS = ["hidden", "qk"]
+
+
+def _norm_pair(hidden: int) -> tuple[Qwen2RMSNorm, RMSNorm]:
+    """A Qwen2RMSNorm and a shared RMSNorm on CUDA holding the same bf16 weights."""
+    qwen2 = Qwen2RMSNorm(hidden, eps=EPS).to("cuda", torch.bfloat16)
+    with torch.no_grad():
+        qwen2.weight.copy_(torch.rand(hidden) + 0.5)
+    shared = RMSNorm(hidden, eps=EPS, dtype=torch.bfloat16).cuda()
+    shared.load_state_dict(qwen2.state_dict())
+    return qwen2, shared
+
+
+def _float64_reference(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    """RMSNorm evaluated in float64, independent of either implementation."""
+    x64 = x.to(torch.float64)
+    variance = x64.pow(2).mean(-1, keepdim=True)
+    return weight.to(torch.float64) * x64 * torch.rsqrt(variance + EPS)
+
+
+@pytest.mark.parametrize("shape", _SHAPES, ids=_SHAPE_IDS)
+def test_fused_rmsnorm_is_within_one_bf16_step_of_qwen2_rmsnorm(shape):
+    torch.manual_seed(0)
+    qwen2, shared = _norm_pair(shape[-1])
+    x = torch.randn(*shape, device="cuda", dtype=torch.bfloat16)
+
+    with torch.no_grad():
+        got, want = shared(x), qwen2(x)
+        torch.testing.assert_close(got, shared._forward_fused(x), atol=0, rtol=0)
+    # bf16 keeps 7 fraction bits, so one rounding step at value v is at most |v| / 128.
+    assert ((got.double() - want.double()).abs() <= want.double().abs() / 128).all()
+
+
+@pytest.mark.parametrize("shape", _SHAPES, ids=_SHAPE_IDS)
+def test_fused_rmsnorm_is_closer_to_float64_than_qwen2_rmsnorm(shape):
+    """Qwen2RMSNorm rounds to bf16 before the weight multiply and so rounds twice;
+    the fused kernel rounds once."""
+    torch.manual_seed(0)
+    qwen2, shared = _norm_pair(shape[-1])
+    x = torch.randn(*shape, device="cuda", dtype=torch.bfloat16)
+
+    reference = _float64_reference(x, qwen2.weight)
+    scale = reference.abs().mean()
+
+    def relative_error(out: torch.Tensor) -> float:
+        return ((out.to(torch.float64) - reference).abs().mean() / scale).item()
+
+    with torch.no_grad():
+        fused, cast_chain = relative_error(shared._forward_fused(x)), relative_error(qwen2(x))
+    assert fused < cast_chain, f"fused {fused:.3e} not better than Qwen2RMSNorm {cast_chain:.3e}"

vllm_omni/diffusion/layers/norm.py

@@ -93,22 +93,8 @@ class RMSNorm(CustomOp):
         # (returns None) and accesses self.weight.data, which is a DTensor under
         # HSDP. Both patterns confuse inductor's compute_ancestors scheduler.
         # Fall back to forward_native so inductor can fuse the pure-PyTorch ops
-        # itself.
-        if torch.compiler.is_compiling():
-            return self.forward_native(x)
-        try:
-            return self._forward_fused(x)
-        except Exception:
-            return self.forward_native(x)
-
-    def forward_hip(
-        self,
-        x: torch.Tensor,
-        residual: torch.Tensor | None = None,
-    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
-        if residual is not None:
-            return self.forward_native(x, residual)
-        if torch.compiler.is_compiling():
+        # itself. The fused kernel also cannot launch on an empty tensor.
+        if torch.compiler.is_compiling() or x.numel() == 0:
             return self.forward_native(x)
         try:
             return self._forward_fused(x)

vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py

@@ -9,10 +9,10 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_utils import ModelMixin
 from einops import rearrange
 from torch import nn
-from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention as OmniAttention
+from vllm_omni.diffusion.layers.norm import RMSNorm
 
 from .rope_real import RotaryPosEmbedReal, apply_real_rotary_emb
 
@@ -39,7 +39,7 @@ class LuminaRMSNormZero(nn.Module):
             bias=True,
         )
 
-        self.norm = Qwen2RMSNorm(embedding_dim, eps=norm_eps)
+        self.norm = RMSNorm(embedding_dim, eps=norm_eps)
 
     def forward(
         self,
@@ -115,7 +115,7 @@ class LuminaLayerNormContinuous(nn.Module):
         if norm_type == "layer_norm":
             self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
         elif norm_type == "rms_norm":
-            self.norm = Qwen2RMSNorm(embedding_dim, eps=eps)
+            self.norm = RMSNorm(embedding_dim, eps=eps)
         else:
             raise ValueError(f"unknown norm_type {norm_type}")
 
@@ -157,7 +157,7 @@ class Lumina2CombinedTimestepCaptionEmbedding(nn.Module):
         )
 
         self.caption_embedder = nn.Sequential(
-            Qwen2RMSNorm(text_feat_dim, eps=norm_eps),
+            RMSNorm(text_feat_dim, eps=norm_eps),
             nn.Linear(text_feat_dim, hidden_size, bias=True),
         )
 
@@ -194,7 +194,7 @@ class SimpleQFormerImageRefiner(nn.Module):
             num_heads = max(1, hidden_size // 128)
         self.num_heads = self._choose_valid_num_heads(hidden_size, num_heads)
         self.input_proj = nn.Sequential(
-            Qwen2RMSNorm(input_hidden_size, eps=norm_eps),
+            RMSNorm(input_hidden_size, eps=norm_eps),
             nn.Linear(input_hidden_size, hidden_size, bias=True),
         )
 
@@ -208,15 +208,15 @@ class SimpleQFormerImageRefiner(nn.Module):
             self.layers.append(
                 nn.ModuleDict(
                     dict(
-                        ln_q1=Qwen2RMSNorm(hidden_size, eps=norm_eps),
+                        ln_q1=RMSNorm(hidden_size, eps=norm_eps),
                         self_attn=nn.MultiheadAttention(
                             embed_dim=hidden_size, num_heads=self.num_heads, dropout=dropout, batch_first=True
                         ),
-                        ln_q2=Qwen2RMSNorm(hidden_size, eps=norm_eps),
+                        ln_q2=RMSNorm(hidden_size, eps=norm_eps),
                         cross_attn=nn.MultiheadAttention(
                             embed_dim=hidden_size, num_heads=self.num_heads, dropout=dropout, batch_first=True
                         ),
-                        ln_ffn=Qwen2RMSNorm(hidden_size, eps=norm_eps),
+                        ln_ffn=RMSNorm(hidden_size, eps=norm_eps),
                         ffn=LuminaFeedForward(dim=hidden_size, inner_dim=4 * hidden_size),
                     )
                 )
@@ -368,9 +368,9 @@ class TransformerBlock(nn.Module):
             out_bias=False,
             processor=processor,
         )
-        # 显式使用 transformers 的 Qwen2RMSNorm，避免依赖 diffusers 内部创建的 `RMSNorm` 再做递归替换。
-        self.attn.norm_q = Qwen2RMSNorm(self.head_dim, eps=1e-5)
-        self.attn.norm_k = Qwen2RMSNorm(self.head_dim, eps=1e-5)
+        # Set the QK norms here instead of letting diffusers build its own RMSNorm via qk_norm.
+        self.attn.norm_q = RMSNorm(self.head_dim, eps=1e-5)
+        self.attn.norm_k = RMSNorm(self.head_dim, eps=1e-5)
         # The kernel itself runs through the shared Omni attention layer: backend
         # selection, padding-mask handling and native grouped-query heads live
         # there. It owns no parameters, so checkpoint keys are unchanged. The
@@ -393,11 +393,11 @@ class TransformerBlock(nn.Module):
         if modulation:
             self.norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True)
         else:
-            self.norm1 = Qwen2RMSNorm(dim, eps=norm_eps)
+            self.norm1 = RMSNorm(dim, eps=norm_eps)
 
-        self.ffn_norm1 = Qwen2RMSNorm(dim, eps=norm_eps)
-        self.norm2 = Qwen2RMSNorm(dim, eps=norm_eps)
-        self.ffn_norm2 = Qwen2RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm1 = RMSNorm(dim, eps=norm_eps)
+        self.norm2 = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm2 = RMSNorm(dim, eps=norm_eps)
 
     def forward(
         self,

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -11,7 +11,6 @@ from diffusers.image_processor import VaeImageProcessor
 from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
 from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
-from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
 from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 
@@ -22,6 +21,7 @@ from vllm_omni.diffusion.cache.cachedit import (
 )
 from vllm_omni.diffusion.data import DiffusionCacheConfig, DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -465,10 +465,10 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
 
     def _reinit_caption_embedder(self, in_features: int) -> None:
         # Align with upstream Mammothmoda2Model's `reinit_caption_embedder`:
-        # Use Qwen2RMSNorm(in_features) + Linear(in_features -> out_features).
+        # Use RMSNorm(in_features) + Linear(in_features -> out_features).
         out_features = int(getattr(self.gen_transformer, "hidden_size", 0) or self.gen_transformer.config.hidden_size)
         self.gen_transformer.time_caption_embed.caption_embedder = nn.Sequential(
-            Qwen2RMSNorm(in_features, eps=1e-5),
+            RMSNorm(in_features, eps=1e-5),
             nn.Linear(in_features, out_features, bias=True),
         )