From 170da72c13b75a3778a5ca9452667922a00ce3ee Mon Sep 17 00:00:00 2001 From: iterhui <46839689+ITerydh@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:04:06 +0800 Subject: [PATCH] [diffusion] perf: fuse tanh-GELU into the LongCat-Image DiT FFN up-proj (#36322) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 登辉 Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- .../runtime/models/dits/longcat_image.py | 32 ++++++++++++++++--- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py b/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py index f7c2dbc08..5739d1dba 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py +++ b/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) - hidden_states = self.act(hidden_states) + proj = self.net[0]["proj"] + 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) - mlp_hidden_states = self.act_mlp(mlp_hidden_states) + if fused_gelu_active(self) and can_use_linear_gelu( + 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,