From e061dd1b470ac7726a0808f31e22518eb903484f Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Thu, 27 Aug 2026 21:06:28 +0800 Subject: [PATCH] [Diffusion][Kernel] Fuse Wan FFN GELU epilogue (#36592) --- .../sglang-diffusion-performance/SKILL.md | 2 +- .../runtime/models/dits/wanvideo.py | 39 ++++++++++++++-- .../test/unit/test_wan_gelu_mlp.py | 46 +++++++++++++++++++ 3 files changed, 82 insertions(+), 5 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_wan_gelu_mlp.py diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md index 9dcaff378..c55b006b7 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md +++ b/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`. | diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 088732303..3333f73fb 100755 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/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, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_wan_gelu_mlp.py b/python/sglang/multimodal_gen/test/unit/test_wan_gelu_mlp.py new file mode 100644 index 000000000..5fd84f2f2 --- /dev/null +++ b/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)