[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/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}"