From 965ad68f1dce29e095bb8576aad5c5da6aac5960 Mon Sep 17 00:00:00 2001 From: "Jierui (Jerry) Xu" <132116599+Jerry2423@users.noreply.github.com> Date: Mon, 28 Sep 2026 09:12:16 -0700 Subject: [PATCH] [Kernel][MammothModa2] Use shared RMSNorm in the DiT (#7489) Signed-off-by: Jerry Xu Co-authored-by: Claude Opus 5 --- tests/diffusion/layers/test_norm.py | 15 +++- .../mammoth_moda2/test_dit_attention_cuda.py | 14 ++-- .../models/mammoth_moda2/test_dit_rmsnorm.py | 58 +++++++++++++++ .../mammoth_moda2/test_dit_rmsnorm_cuda.py | 74 +++++++++++++++++++ vllm_omni/diffusion/layers/norm.py | 18 +---- .../mammoth_moda2/mammothmoda2_dit_model.py | 30 ++++---- .../pipeline_mammothmoda2_dit.py | 6 +- 7 files changed, 172 insertions(+), 43 deletions(-) create mode 100644 tests/diffusion/models/mammoth_moda2/test_dit_rmsnorm.py create mode 100644 tests/diffusion/models/mammoth_moda2/test_dit_rmsnorm_cuda.py diff --git a/tests/diffusion/layers/test_norm.py b/tests/diffusion/layers/test_norm.py index 4794ca1c2..529cfa134 100644 --- a/tests/diffusion/layers/test_norm.py +++ b/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 diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py b/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py index 62415953a..050d404bd 100644 --- a/tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py +++ b/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( diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_rmsnorm.py b/tests/diffusion/models/mammoth_moda2/test_dit_rmsnorm.py new file mode 100644 index 000000000..261d4c05d --- /dev/null +++ b/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()) diff --git a/tests/diffusion/models/mammoth_moda2/test_dit_rmsnorm_cuda.py b/tests/diffusion/models/mammoth_moda2/test_dit_rmsnorm_cuda.py new file mode 100644 index 000000000..71cf8181d --- /dev/null +++ b/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}" diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py index 478fc0d09..1c98c6f01 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/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) diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py index 415a17ed8..3a0291398 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py +++ b/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, diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py index c7d0e7310..fd716aa9a 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/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), )