diff --git a/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py b/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py index a0c58b4ec..324d6b5e1 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py +++ b/python/sglang/multimodal_gen/configs/models/dits/mova_audio.py @@ -18,6 +18,7 @@ class MOVAAudioArchConfig(DiTArchConfig): default_factory=lambda: { r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_norm.\2", r"^text_embedding\.0\.(.*)$": r"text_embedding.fc_in.\1", r"^text_embedding\.2\.(.*)$": r"text_embedding.fc_out.\1", r"^time_embedding\.0\.(.*)$": r"time_embedding.fc_in.\1", diff --git a/python/sglang/multimodal_gen/configs/models/dits/mova_video.py b/python/sglang/multimodal_gen/configs/models/dits/mova_video.py index 0ec03932a..a7983a45a 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/mova_video.py +++ b/python/sglang/multimodal_gen/configs/models/dits/mova_video.py @@ -18,6 +18,7 @@ class MOVAVideoArchConfig(DiTArchConfig): default_factory=lambda: { r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_norm.\2", r"^text_embedding\.0\.(.*)$": r"text_embedding.fc_in.\1", r"^text_embedding\.2\.(.*)$": r"text_embedding.fc_out.\1", r"^time_embedding\.0\.(.*)$": r"time_embedding.fc_in.\1", diff --git a/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py b/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py index e947728de..31d43a458 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py @@ -19,7 +19,9 @@ from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAt # Reuse SGLang's optimized RMSNorm instead of torch.nn.RMSNorm or custom SlowRMSNorm from sglang.multimodal_gen.runtime.layers.layernorm import ( + LayerNormScaleShift, RMSNorm, + ScaleResidualLayerNormScaleShift, tensor_parallel_rms_norm, ) from sglang.multimodal_gen.runtime.layers.linear import ( @@ -247,14 +249,12 @@ class CrossAttention(nn.Module): return x -class GateModule(nn.Module): - def __init__( - self, - ): +class MulAdd(nn.Module): + def __init__(self): super().__init__() def forward(self, x, gate, residual): - return x + gate * residual + return residual + gate * x class DiTBlock(nn.Module): @@ -272,12 +272,18 @@ class DiTBlock(nn.Module): self.self_attn = SelfAttention(dim, num_heads, eps) self.cross_attn = CrossAttention(dim, num_heads, eps) - self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) - self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) - self.norm3 = nn.LayerNorm(dim, eps=eps) + self.norm1 = LayerNormScaleShift( + dim, eps=eps, elementwise_affine=False, dtype=torch.float32 + ) + self.self_attn_norm = nn.LayerNorm(dim, eps=eps) + # Fused: residual + 1 * cross_attn_out → layernorm + scale/shift + # Replaces the old norm2 (LayerNormScaleShift) + residual add for cross-attention + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, eps=eps, elementwise_affine=False, dtype=torch.float32 + ) self.ffn = MLP(dim, ffn_dim, output_dim=dim, act_type="gelu_pytorch_tanh") self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) - self.gate = GateModule() + self.mlp_residual = MulAdd() def forward(self, x, context, t_mod, freqs): has_seq = len(t_mod.shape) == 4 @@ -295,11 +301,23 @@ class DiTBlock(nn.Module): scale_mlp.squeeze(2), gate_mlp.squeeze(2), ) - input_x = modulate(self.norm1(x), shift_msa, scale_msa) - x = self.gate(x, gate_msa, self.self_attn(input_x, freqs)) - x = x + self.cross_attn(self.norm3(x), context) - input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) - x = self.gate(x, gate_mlp, self.ffn(input_x)) + orig_dtype = x.dtype + # 1. Self-attention, fuse: + # - layernorm(x) * (1 + scale_msa) + shift_msa + input_x = self.norm1(x, shift_msa, scale_msa) + # 2. torch.compile may fuse mlp_residual and self_attn_norm + x = self.mlp_residual(self.self_attn(input_x, freqs), gate_msa, x) + norm_x = self.self_attn_norm(x) + # 3. Cross-attention, fuse: + # - x = x + 1 * cross_output + # - input_x = layernorm(x) * (1 + scale_mlp) + shift_mlp + cross_output = self.cross_attn(norm_x, context) + input_x, x = self.cross_attn_residual_norm( + x, cross_output, 1, shift_mlp, scale_mlp + ) + # 4. Feed-forward + x = self.mlp_residual(self.ffn(input_x), gate_mlp, x) + x = x.to(orig_dtype) return x @@ -310,7 +328,9 @@ class Head(nn.Module): super().__init__() self.dim = dim self.patch_size = patch_size - self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm = LayerNormScaleShift( + dim, eps=eps, elementwise_affine=False, dtype=torch.float32 + ) # Output dim is small for MOVA; replicate to avoid TP shape coupling. self.head = ReplicatedLinear(dim, out_dim * math.prod(patch_size)) self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) @@ -321,12 +341,12 @@ class Head(nn.Module): self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(2) ).chunk(2, dim=2) - x, _ = self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2)) + x, _ = self.head(self.norm(x, shift.squeeze(2), scale.squeeze(2))) else: shift, scale = ( self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod ).chunk(2, dim=1) - x, _ = self.head(self.norm(x) * (1 + scale) + shift) + x, _ = self.head(self.norm(x, shift, scale)) return x