[diffusion] perf: fuse tanh-GELU into the LongCat-Image DiT FFN up-proj (#36322)

Co-authored-by: 登辉 <yangdenghui.ydh@alibaba-inc.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

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

@@ -35,6 +35,10 @@ from diffusers.models.normalization import (
 from sglang.kernels.ops.diffusion import (
     BitExactFusionGate,
     can_use_fused_inplace_qknorm_rope,
+    can_use_linear_gelu,
+    fused_gelu_active,
+    fused_linear_gelu_tanh,
+    mark_fused_gelu_site,
     tensors_equal,
 )
 from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
@@ -217,10 +221,20 @@ class _LongCatFFN(nn.Module):
             ]
         )
         self.act = nn.GELU(approximate="tanh")
+        # quality="high" site: up-proj GEMM + tanh-GELU cublasLt epilogue. Off by
+        # default; the denoising stage mounts it per batch. The ModuleDict holds
+        # `proj` in _modules, so getattr resolves it for the fusion helper.
+        mark_fused_gelu_site(self.net[0], "proj")
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states, _ = self.net[0]["proj"](hidden_states)
+        proj = self.net[0]["proj"]
-        hidden_states = self.act(hidden_states)
+        if fused_gelu_active(self.net[0]) and can_use_linear_gelu(proj, hidden_states):
+            hidden_states = fused_linear_gelu_tanh(
+                hidden_states, proj.weight, proj.bias
+            )
+        else:
+            hidden_states, _ = proj(hidden_states)
+            hidden_states = self.act(hidden_states)
         hidden_states, _ = self.net[2](hidden_states)
         return hidden_states
 
@@ -496,6 +510,9 @@ class _SingleTransformerBlock(nn.Module):
             prefix=f"{prefix}.proj_mlp",
         )
         self.act_mlp = nn.GELU(approximate="tanh")
+        # quality="high" site: proj_mlp GEMM + tanh-GELU cublasLt epilogue,
+        # mounted per batch by the denoising stage; off (bit-exact) by default.
+        mark_fused_gelu_site(self, "proj_mlp")
         # proj_out: RowParallelLinear reduces sharded [attn | mlp] concat via
         # all-reduce, matching Flux2SingleTransformerBlockAttention.to_out.
         self.proj_out = RowParallelLinear(
@@ -560,8 +577,15 @@ class _SingleTransformerBlock(nn.Module):
 
         residual = hidden_states
         norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
-        mlp_hidden_states, _ = self.proj_mlp(norm_hidden_states)
+        if fused_gelu_active(self) and can_use_linear_gelu(
-        mlp_hidden_states = self.act_mlp(mlp_hidden_states)
+            self.proj_mlp, norm_hidden_states
+        ):
+            mlp_hidden_states = fused_linear_gelu_tanh(
+                norm_hidden_states, self.proj_mlp.weight, self.proj_mlp.bias
+            )
+        else:
+            mlp_hidden_states, _ = self.proj_mlp(norm_hidden_states)
+            mlp_hidden_states = self.act_mlp(mlp_hidden_states)
         attn_output = self.attn(
             hidden_states=norm_hidden_states,
             image_rotary_emb=image_rotary_emb,