[diffusion] perf: apply mul add fusion for Qwen-Image (#16299)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

python/sglang/multimodal_gen/runtime/layers/elementwise.py

@@ -0,0 +1,35 @@
+import torch
+
+from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
+from sglang.multimodal_gen.runtime.layers.triton_ops import fuse_scale_shift_kernel
+
+
+class MulAdd(CustomOp):
+    """
+    Fuse elementwise mul and add
+    Input: a, b, c, OptionalInt[k]
+    Output: a * (k + b) + c
+    """
+
+    def __init__(self, prefix: str = ""):
+        super().__init__()
+
+    def forward_native(
+        self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
+    ) -> torch.Tensor:
+        # a.shape: [batch_size, seq_len, inner_dim]
+        if b.dim() == 4:
+            # b.shape: [batch_size, num_frames, 1, inner_dim]
+            num_frames = b.shape[1]
+            frame_seqlen = a.shape[1] // num_frames
+            return c + (
+                a.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * b
+            ).flatten(1, 2)
+        else:
+            # b.shape: [batch_size, 1, inner_dim]
+            return c + a * (k + b)
+
+    def forward_cuda(
+        self, a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, k: int = 0
+    ):
+        return fuse_scale_shift_kernel(a, b, c, scale_constant=k)

python/sglang/multimodal_gen/runtime/layers/layernorm.py

@@ -238,31 +238,6 @@ class LayerNorm(CustomOp):
         return s
 
 
-class ScaleResidual(nn.Module):
-    """
-    Applies gated residual connection.
-    """
-
-    def __init__(self, prefix: str = ""):
-        super().__init__()
-
-    def forward(
-        self, residual: torch.Tensor, x: torch.Tensor, gate: torch.Tensor
-    ) -> torch.Tensor:
-        """Apply gated residual connection."""
-        # x.shape: [batch_size, seq_len, inner_dim]
-        if gate.dim() == 4:
-            # gate.shape: [batch_size, num_frames, 1, inner_dim]
-            num_frames = gate.shape[1]
-            frame_seqlen = x.shape[1] // num_frames
-            return residual + (
-                x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate
-            ).flatten(1, 2)
-        else:
-            # gate.shape: [batch_size, 1, inner_dim]
-            return residual + x * gate
-
-
 # adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py
 # NOTE(will): Needed to match behavior of diffusers and wan2.1 even while using
 # FSDP's MixedPrecisionPolicy

python/sglang/multimodal_gen/runtime/layers/triton_ops.py

@@ -25,6 +25,7 @@ def _fused_scale_shift_4d_kernel(
     normalized_ptr,
     scale_ptr,
     shift_ptr,
+    scale_constant: tl.constexpr,  # scale_constant is either 0 or 1.
     rows,
     inner_dim,
     seq_len,
@@ -56,8 +57,8 @@ def _fused_scale_shift_4d_kernel(
     scale = tl.load(scale_ptrs, mask=mask, other=0.0)
     shift = tl.load(shift_ptrs, mask=mask, other=0.0)
 
-    one = tl.full([BLOCK_N], 1.0, dtype=scale.dtype)
-    output = normalized * (one + scale) + shift
+    scale_const_tensor = tl.full([BLOCK_N], scale_constant, dtype=scale.dtype)
+    output = normalized * (scale_const_tensor + scale) + shift
 
     tl.store(out_ptrs, output, mask=mask)
 
@@ -67,6 +68,7 @@ def fuse_scale_shift_kernel_blc_opt(
     x_ptr,
     shift_ptr,
     scale_ptr,
+    scale_constant: tl.constexpr,  # scale_constant is either 0 or 1.,
     y_ptr,
     B,
     L,
@@ -125,7 +127,7 @@ def fuse_scale_shift_kernel_blc_opt(
         )
         scale = tl.load(scale_ptr + sc_off, mask=mask, other=0)
 
-    y = x * (1 + scale) + shift
+    y = x * (scale_constant + scale) + shift
     tl.store(y_ptr + x_off, y, mask=mask)
 
 
@@ -221,6 +223,7 @@ def fuse_scale_shift_kernel(
     x: torch.Tensor,
     scale: torch.Tensor,
     shift: torch.Tensor,
+    scale_constant: float = 1.0,
     block_l: int = 128,
     block_c: int = 128,
 ):
@@ -251,6 +254,7 @@ def fuse_scale_shift_kernel(
             x_2d,
             scale_reshaped,
             shift_reshaped,
+            scale_constant,
             rows,
             C,
             L,
@@ -306,6 +310,7 @@ def fuse_scale_shift_kernel(
             x,
             shift_blc if need_shift_scalar else shift_exp,
             scale_blc if need_scale_scalar else scale_exp,
+            scale_constant,
             output,
             B,
             L,

python/sglang/multimodal_gen/runtime/models/dits/causal_wanvideo.py

@@ -26,11 +26,11 @@ import torch.distributed as dist
 from sglang.multimodal_gen.configs.models.dits import WanVideoConfig
 from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size
 from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
+from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd
 from sglang.multimodal_gen.runtime.layers.layernorm import (
     FP32LayerNorm,
     LayerNormScaleShift,
     RMSNorm,
-    ScaleResidual,
     ScaleResidualLayerNormScaleShift,
 )
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
@@ -318,7 +318,7 @@ class CausalWanTransformerBlock(nn.Module):
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
-        self.mlp_residual = ScaleResidual()
+        self.mlp_residual = MulAdd()
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
 
@@ -417,7 +417,7 @@ class CausalWanTransformerBlock(nn.Module):
 
         # 3. Feed-forward
         ff_output = self.ffn(norm_hidden_states)
-        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
+        hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states)
         hidden_states = hidden_states.to(orig_dtype)
 
         return hidden_states

python/sglang/multimodal_gen/runtime/models/dits/hunyuanvideo.py

@@ -15,10 +15,10 @@ from sglang.multimodal_gen.runtime.layers.attention import (
     LocalAttention,
     UlyssesAttention,
 )
+from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd
 from sglang.multimodal_gen.runtime.layers.layernorm import (
     LayerNormScaleShift,
     RMSNorm,
-    ScaleResidual,
     ScaleResidualLayerNormScaleShift,
 )
 from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
@@ -81,7 +81,7 @@ class MMDoubleStreamBlock(nn.Module):
         self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
             hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype
         )
-        self.img_mlp_residual = ScaleResidual()
+        self.img_mlp_residual = MulAdd()
 
         # Image attention components
         self.img_attn_qkv = ReplicatedLinear(
@@ -127,7 +127,7 @@ class MMDoubleStreamBlock(nn.Module):
         self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
             hidden_size, norm_type="layer", elementwise_affine=False, dtype=dtype
         )
-        self.txt_mlp_residual = ScaleResidual()
+        self.txt_mlp_residual = MulAdd()
 
         # Text attention components
         self.txt_attn_qkv = ReplicatedLinear(
@@ -231,7 +231,7 @@ class MMDoubleStreamBlock(nn.Module):
 
         # Process image MLP
         img_mlp_out = self.img_mlp(img_mlp_input)
-        img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate)
+        img = self.img_mlp_residual(img_mlp_out, img_mlp_gate, img_residual)
 
         # Process text attention output
         txt_attn_out, _ = self.txt_attn_proj(
@@ -245,7 +245,7 @@ class MMDoubleStreamBlock(nn.Module):
 
         # Process text MLP
         txt_mlp_out = self.txt_mlp(txt_mlp_input)
-        txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate)
+        txt = self.txt_mlp_residual(txt_mlp_out, txt_mlp_gate, txt_residual)
 
         return img, txt
 
@@ -304,7 +304,7 @@ class MMSingleStreamBlock(nn.Module):
             elementwise_affine=False,
             dtype=dtype,
         )
-        self.output_residual = ScaleResidual()
+        self.output_residual = MulAdd()
 
         # Activation function
         self.mlp_act = nn.GELU(approximate="tanh")
@@ -384,7 +384,7 @@ class MMSingleStreamBlock(nn.Module):
         output, _ = self.linear2(combined)
 
         # Apply residual connection with gating using fused operation
-        return self.output_residual(x, output, mod_gate)
+        return self.output_residual(output, mod_gate, x)
 
 
 class HunyuanVideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -17,6 +17,7 @@ from diffusers.models.normalization import AdaLayerNormContinuous
 from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
+from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd
 from sglang.multimodal_gen.runtime.layers.layernorm import (
     LayerNorm,
     RMSNorm,
@@ -28,7 +29,6 @@ from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
 )
 from sglang.multimodal_gen.runtime.layers.triton_ops import (
     fuse_scale_shift_gate_select01_kernel,
-    fuse_scale_shift_kernel,
 )
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.platforms import (
@@ -350,7 +350,7 @@ class QwenEmbedLayer3DRope(nn.Module):
             if idx != layer_num:
                 video_freq = self._compute_video_freqs(frame, height, width, idx)
             else:
-                ### For the condition image, we set the layer index to -1
+                # For the condition image, we set the layer index to -1
                 video_freq = self._compute_condition_freqs(frame, height, width)
             video_freq = video_freq.to(device)
             vid_freqs.append(video_freq)
@@ -673,6 +673,8 @@ class QwenImageTransformerBlock(nn.Module):
         self.txt_mlp = FeedForward(
             dim=dim, dim_out=dim, activation_fn="gelu-approximate"
         )
+        # Utils
+        self.fuse_mul_add = MulAdd()
 
     def _modulate(self, x, mod_params, index=None):
         shift, scale, gate = mod_params.chunk(3, dim=-1)
@@ -714,14 +716,14 @@ class QwenImageTransformerBlock(nn.Module):
                 )
                 gate_result = torch.where(mask, gate0.unsqueeze(1), gate1.unsqueeze(1))
                 return (
-                    fuse_scale_shift_kernel(x, scale_result, shift_result),
+                    self.fuse_mul_add(x, scale_result, shift_result, k=1.0),
                     gate_result,
                 )
         else:
             shift_result = shift.unsqueeze(1)
             scale_result = scale.unsqueeze(1)
             gate_result = gate.unsqueeze(1)
-            return fuse_scale_shift_kernel(x, scale_result, shift_result), gate_result
+            return self.fuse_mul_add(x, scale_result, shift_result, k=1.0), gate_result
 
     def forward(
         self,
@@ -759,8 +761,10 @@ class QwenImageTransformerBlock(nn.Module):
         # 4. Splits results back to separate streams
         joint_attention_kwargs = joint_attention_kwargs or {}
         attn_output = self.attn(
-            hidden_states=img_modulated,  # Image stream (will be processed as "sample")
-            encoder_hidden_states=txt_modulated,  # Text stream (will be processed as "context")
+            # Image stream (will be processed as "sample")
+            hidden_states=img_modulated,
+            # Text stream (will be processed as "context")
+            encoder_hidden_states=txt_modulated,
             encoder_hidden_states_mask=encoder_hidden_states_mask,
             image_rotary_emb=image_rotary_emb,
             **joint_attention_kwargs,
@@ -780,13 +784,15 @@ class QwenImageTransformerBlock(nn.Module):
             img_normed2, img_mod2, modulate_index
         )
         img_mlp_output = self.img_mlp(img_modulated2)
-        hidden_states = hidden_states + img_gate2 * img_mlp_output
+        hidden_states = self.fuse_mul_add(img_mlp_output, img_gate2, hidden_states)
 
         # Process text stream - norm2 + MLP
         txt_normed2 = self.txt_norm2(encoder_hidden_states)
         txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2)
         txt_mlp_output = self.txt_mlp(txt_modulated2)
-        encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
+        encoder_hidden_states = self.fuse_mul_add(
+            txt_mlp_output, txt_gate2, encoder_hidden_states
+        )
 
         # Clip to prevent overflow for fp16
         if encoder_hidden_states.dtype == torch.float16:

python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py

@@ -20,11 +20,11 @@ from sglang.multimodal_gen.runtime.layers.attention import (
     UlyssesAttention_VSA,
     USPAttention,
 )
+from sglang.multimodal_gen.runtime.layers.elementwise import MulAdd
 from sglang.multimodal_gen.runtime.layers.layernorm import (
     FP32LayerNorm,
     LayerNormScaleShift,
     RMSNorm,
-    ScaleResidual,
     ScaleResidualLayerNormScaleShift,
     tensor_parallel_rms_norm,
 )
@@ -382,7 +382,7 @@ class WanTransformerBlock(nn.Module):
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
-        self.mlp_residual = ScaleResidual()
+        self.mlp_residual = MulAdd()
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
 
@@ -488,7 +488,7 @@ class WanTransformerBlock(nn.Module):
 
         # 3. Feed-forward
         ff_output = self.ffn(norm_hidden_states)
-        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
+        hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states)
         hidden_states = hidden_states.to(orig_dtype)
 
         return hidden_states
@@ -582,7 +582,7 @@ class WanTransformerBlock_VSA(nn.Module):
 
         # 3. Feed-forward
         self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
-        self.mlp_residual = ScaleResidual()
+        self.mlp_residual = MulAdd()
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
 
@@ -669,7 +669,7 @@ class WanTransformerBlock_VSA(nn.Module):
 
         # 3. Feed-forward
         ff_output = self.ffn(norm_hidden_states)
-        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
+        hidden_states = self.mlp_residual(ff_output, c_gate_msa, hidden_states)
         hidden_states = hidden_states.to(orig_dtype)
 
         return hidden_states