minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: 登辉 <yangdenghui.ydh@alibaba-inc.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
iterhui
2026-08-26 19:04:06 +08:00
committed by GitHub
co-authored by 登辉 Xiaoyu Zhang
parent 8005df61d3
commit 170da72c13
1 changed files with 28 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/models/dits/longcat_image.py
+28 -4
View File
@@ -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,