From 4901693110c354e66c60d84c80f4678cb19faea4 Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Tue, 16 Dec 2025 01:18:47 +0800
Subject: [PATCH] [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>
---
 .../configs/models/dits/zimage.py             |  4 +++
 .../runtime/layers/activation.py              |  9 ++++--
 .../runtime/models/dits/zimage.py             | 32 +++++++++++++------
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/python/sglang/multimodal_gen/configs/models/dits/zimage.py b/python/sglang/multimodal_gen/configs/models/dits/zimage.py
index 1f4f2bb22..fc0b17830 100644
--- a/python/sglang/multimodal_gen/configs/models/dits/zimage.py
+++ b/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),
         }
     )
 
diff --git a/python/sglang/multimodal_gen/runtime/layers/activation.py b/python/sglang/multimodal_gen/runtime/layers/activation.py
index 4eff9ba1c..afbad4ed3 100644
--- a/python/sglang/multimodal_gen/runtime/layers/activation.py
+++ b/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()."""
diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py
index ad1e9059b..c81a6f054 100644
--- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py
+++ b/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,