[diffusion] perf: support FFN pack gate and up proj for Z-Image(#15201)

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

python/sglang/multimodal_gen/configs/models/dits/zimage.py

@@ -32,6 +32,8 @@ class ZImageArchConfig(DiTArchConfig):
             (".to_qkv", ".to_q", "q"),
             (".to_qkv", ".to_k", "k"),
             (".to_qkv", ".to_v", "v"),
+            (".feed_forward.w13", ".feed_forward.w1", "gate"),
+            (".feed_forward.w13", ".feed_forward.w3", "up"),
         ]
     )
 
@@ -40,6 +42,8 @@ class ZImageArchConfig(DiTArchConfig):
             r"(.*)\.to_q\.weight$": (r"\1.to_qkv.weight", 0, 3),
             r"(.*)\.to_k\.weight$": (r"\1.to_qkv.weight", 1, 3),
             r"(.*)\.to_v\.weight$": (r"\1.to_qkv.weight", 2, 3),
+            r"(.*)\.feed_forward\.w1\.weight$": (r"\1.feed_forward.w13.weight", 0, 2),
+            r"(.*)\.feed_forward\.w3\.weight$": (r"\1.feed_forward.w13.weight", 1, 2),
         }
     )
 

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

@@ -9,6 +9,7 @@ from typing import Any
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from sgl_kernel import silu_and_mul
 
 # TODO (will): remove this dependency
 from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
@@ -28,8 +29,12 @@ class SiluAndMul(CustomOp):
     def __init__(self) -> None:
         super().__init__()
 
-    def forward_cuda(self, *args, **kwargs) -> Any:
-        return self.forward_native(*args, **kwargs)
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        d = x.shape[-1] // 2
+        output_shape = x.shape[:-1] + (d,)
+        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
+        silu_and_mul(x, out)
+        return out
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         """PyTorch-native implementation equivalent to forward()."""

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

@@ -3,14 +3,16 @@ from typing import Any, List, Optional, Tuple
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from sglang.multimodal_gen.configs.models.dits.zimage import ZImageDitConfig
+from sglang.multimodal_gen.runtime.layers.activation import SiluAndMul
 from sglang.multimodal_gen.runtime.layers.attention import USPAttention
 from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
 from sglang.multimodal_gen.runtime.layers.linear import (
+    MergedColumnParallelLinear,
     QKVParallelLinear,
     ReplicatedLinear,
+    RowParallelLinear,
 )
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import _apply_rotary_emb
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
@@ -77,17 +79,22 @@ class TimestepEmbedder(nn.Module):
 class FeedForward(nn.Module):
     def __init__(self, dim: int, hidden_dim: int):
         super().__init__()
-        self.w1 = ReplicatedLinear(dim, hidden_dim, bias=False)
-        self.w2 = ReplicatedLinear(hidden_dim, dim, bias=False)
-        self.w3 = ReplicatedLinear(dim, hidden_dim, bias=False)
-
-    def _forward_silu_gating(self, x1, x3):
-        return F.silu(x1) * x3
+        self.w13 = MergedColumnParallelLinear(
+            input_size=dim,
+            output_sizes=[hidden_dim] * 2,
+            bias=False,
+        )
+        self.w2 = RowParallelLinear(
+            input_size=hidden_dim,
+            output_size=dim,
+            bias=False,
+        )
+        self.act = SiluAndMul()
 
     def forward(self, x):
-        x1, _ = self.w1(x)
-        x3, _ = self.w3(x)
-        out, _ = self.w2(self._forward_silu_gating(x1, x3))
+        x13, _ = self.w13(x)
+        x = self.act(x13)
+        out, _ = self.w2(x)
         return out
 
 
@@ -353,6 +360,11 @@ class ZImageTransformer2DModel(CachableDiT):
     _no_split_modules = ["ZImageTransformerBlock"]
     param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping
 
+    param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping
+    reverse_param_names_mapping = (
+        ZImageDitConfig().arch_config.reverse_param_names_mapping
+    )
+
     def __init__(
         self,
         config: ZImageDitConfig,