[Diffusion][Kernel] Fuse Wan FFN GELU epilogue (#36592)

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md

@@ -383,7 +383,7 @@ Use these as first commands to benchmark, not as universal winners.
 | Z-Image / Z-Image-Turbo | 1024x1024, runtime-default steps/guidance, 1 GPU | `--enable-torch-compile --warmup-mode request` | Keep base Z-Image separate from Turbo: base uses 50-step CFG defaults, Turbo uses 9-step zero-CFG defaults. Mainline has bf16-native Triton RMSNorm scale and tanh-residual fusions. |
 | Wan2.2 A14B T2V/I2V | 1280x720, 81 frames | Nightly: `--num-gpus 4 --enable-cfg-parallel --ulysses-degree 2 --text-encoder-cpu-offload --pin-cpu-memory` | For lowest latency, also benchmark pure Ulysses on the same GPUs. |
 | Wan2.2 TI2V 5B | 1280x720, 81 frames, 1 GPU | `--enable-torch-compile --warmup-mode request` | Keep the input image and motion prompt fixed when comparing sparse attention or Cache-DiT. |
-| Wan2.1 / FastWan / TurboWan variants | 480p or 720p video, family defaults | `--enable-torch-compile --warmup-mode request`; add `--ulysses-degree` / CFG parallel only after measuring | Current registry includes Wan2.1, FastWan2.1, FastWan2.2 TI2V, TurboWan2.1, TurboWan2.2 I2V, and Wan2.1-Fun InP. Use the compatibility matrix and benchmark presets before choosing topology. |
+| Wan2.1 / FastWan / TurboWan variants | 480p or 720p video, family defaults | Compare `--quality lossless` with `--quality high`, then try `--enable-torch-compile --warmup-mode request`; add `--ulysses-degree` / CFG parallel only after measuring | `quality=high` mounts the Wan FFN cublasLt GELU epilogue and the Wan VAE RMSNorm+SiLU fast path when their guards pass; validate video quality against lossless. Current registry includes Wan2.1, FastWan2.1, FastWan2.2 TI2V, TurboWan2.1, TurboWan2.2 I2V, and Wan2.1-Fun InP. Use the compatibility matrix and benchmark presets before choosing topology. |
 | Cosmos3 Nano / Super | T2I: 1024x1024 with `--num-frames 1`; T2V/I2V: 480p/720p video | Start with `--performance-mode auto --warmup-mode request`; use `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` only for benchmark isolation, and compare compile separately | One checkpoint serves T2I/T2V/I2V. Mode is request-driven: `num_frames == 1` means T2I, `--image-path` means I2V. On GPUs with at least 120 GiB available, auto mode keeps the Cosmos3 DiT and VAE resident for every checkpoint in the family; a 1xH200 832x480x9f, 4-step eager ABBA reduced e2e from 1.576 to 0.428 seconds with exact output parity. Cosmos3 runs one DiT per pipeline, so component offload above that threshold only buys a DiT copy out to host memory and back per request -- it cost Cosmos3-Super 720p 81f T2V ~4s of ~115s on 2xH200. |
 | Cosmos3 Edge / distilled Super | Edge T2I: 640x640, 35 steps, 1 GPU; distilled Super T2I: 640x640, fixed 4-step schedule, 4 GPUs | Start eager with `--performance-mode manual`; use `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` only for benchmark isolation | Edge is trained for 256p/480p shapes. Distilled checkpoints own their sigma schedule and force guidance 1.0; do not override steps or flow shift. Do not retry the closed experimental Cosmos BCG path without a new lifecycle design. |
 | Ideogram 4 FP8/NVFP4 | 1024x1024, native preset defaults | `--enable-torch-compile --warmup-mode request` | Do not set `--num-inference-steps` or `--guidance-scale` directly unless you also update the Ideogram preset; sampling params derive them from `preset`. |

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

@@ -12,7 +12,11 @@ import torch.nn as nn
 from sglang.kernels.ops.diffusion import (
     BitExactFusionGate,
     can_use_fused_temb_table_slices,
+    can_use_linear_gelu,
+    fused_gelu_active,
+    fused_linear_gelu_tanh,
     fused_temb_table_slices,
+    mark_fused_gelu_site,
     tensors_equal,
 )
 from sglang.multimodal_gen.configs.models.dits import WanVideoConfig
@@ -76,6 +80,35 @@ if USE_AITER:
     from aiter.ops.rope import rope_cached_2c_fwd_inplace
 
 
+class _WanGELUMLP(MLP):
+    """Wan FFN with a quality-gated cublasLt GELU epilogue."""
+
+    def __init__(
+        self,
+        dim: int,
+        ffn_dim: int,
+        prefix: str,
+        quant_config: QuantizationConfig | None,
+    ):
+        super().__init__(
+            dim,
+            ffn_dim,
+            act_type="gelu_pytorch_tanh",
+            prefix=prefix,
+            quant_config=quant_config,
+        )
+        mark_fused_gelu_site(self, "fc_in")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if fused_gelu_active(self) and can_use_linear_gelu(self.fc_in, x):
+            x = fused_linear_gelu_tanh(x, self.fc_in.weight, self.fc_in.bias)
+        else:
+            x, _ = self.fc_in(x)
+            x = self.act(x)
+        x, _ = self.fc_out(x)
+        return x
+
+
 class WanImageEmbedding(torch.nn.Module):
     def __init__(self, in_features: int, out_features: int):
         super().__init__()
@@ -528,10 +561,9 @@ class WanTransformerBlock(nn.Module):
         )
 
         # 3. Feed-forward
-        self.ffn = MLP(
+        self.ffn = _WanGELUMLP(
             dim,
             ffn_dim,
-            act_type="gelu_pytorch_tanh",
             prefix=add_prefix("ffn", prefix),
             quant_config=quant_config,
         )
@@ -798,10 +830,9 @@ class WanTransformerBlock_VSA(nn.Module):
         )
 
         # 3. Feed-forward
-        self.ffn = MLP(
+        self.ffn = _WanGELUMLP(
             dim,
             ffn_dim,
-            act_type="gelu_pytorch_tanh",
             prefix=add_prefix("ffn", prefix),
             quant_config=quant_config,
         )

python/sglang/multimodal_gen/test/unit/test_wan_gelu_mlp.py

@@ -0,0 +1,46 @@
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    fused_gelu_active,
+    mount_fused_linear_gelu,
+    unmount_fused_linear_gelu,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    maybe_init_distributed_environment_and_model_parallel,
+    model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.models.dits.wanvideo import _WanGELUMLP
+from sglang.multimodal_gen.test.single_test_file.component_accuracy.utils import (
+    ensure_distributed_env_defaults,
+)
+
+
+def _ensure_single_process_parallel_runtime() -> None:
+    if model_parallel_is_initialized():
+        return
+    ensure_distributed_env_defaults()
+    maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=1)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@torch.no_grad()
+def test_wan_gelu_mlp_quality_path_and_lossless_restore():
+    _ensure_single_process_parallel_runtime()
+    torch.manual_seed(0)
+    mlp = _WanGELUMLP(64, 256, prefix="", quant_config=None).to(
+        device="cuda", dtype=torch.bfloat16
+    )
+    for parameter in mlp.parameters():
+        parameter.normal_(mean=0.0, std=0.02)
+
+    x = torch.randn(2, 129, 64, device="cuda", dtype=torch.bfloat16)
+    reference = mlp(x)
+    assert not fused_gelu_active(mlp)
+
+    assert mount_fused_linear_gelu(mlp)
+    torch.testing.assert_close(mlp(x), reference, atol=2e-2, rtol=2e-2)
+
+    unmount_fused_linear_gelu(mlp)
+    assert not fused_gelu_active(mlp)
+    assert torch.equal(mlp(x), reference)