[diffusion] Fuse DiT FFN tanh-GELU into up-proj GEMM (cublasLt epilogue) behind quality=high (Qwen-Image 1024^2 denoise 12.36 -> 12.05 s on H200) (#33536)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-08-04 23:49:43 +08:00
parent 0d0c7d853f
commit 95d0e57e83
6 changed files with 317 additions and 0 deletions
@@ -0,0 +1,59 @@
+"""Core checks for the quality-gated linear + tanh-GELU fusion."""
+
+import sys
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.kernels.ops.diffusion import fused_linear_gelu as gelu
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+class _Site(nn.Module):
+    def __init__(self, dtype=torch.bfloat16, bias=True):
+        super().__init__()
+        self.proj = nn.Linear(64, 256, bias=bias, device="cuda", dtype=dtype)
+        gelu.mark_fused_gelu_site(self, "proj")
+
+    def forward(self, x):
+        if self._sgl_fused_gelu_enabled and gelu.can_fuse_linear_gelu(self.proj, x):
+            return gelu.fused_linear_gelu_tanh(x, self.proj.weight, self.proj.bias)
+        return F.gelu(self.proj(x), approximate="tanh")
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_fused_matches_reference(dtype):
+    torch.manual_seed(0)
+    site = _Site(dtype)
+    x = torch.randn(512, 64, device="cuda", dtype=dtype)
+    ref = site(x)
+    assert gelu.mount_fused_linear_gelu(site)
+    atol = 2e-2 if dtype == torch.bfloat16 else 4e-3
+    torch.testing.assert_close(site(x), ref, atol=atol, rtol=2e-2)
+
+
+def test_mount_guards_and_lossless_path():
+    torch.manual_seed(0)
+    good, bad = _Site(), _Site(torch.float32)
+    model = nn.ModuleList([good, bad])
+    assert not gelu.mount_fused_linear_gelu(model)
+    assert not good._sgl_fused_gelu_enabled
+
+    x = torch.randn(16, 64, device="cuda", dtype=torch.bfloat16)
+    ref = good(x)
+    assert gelu.mount_fused_linear_gelu(good)
+    gelu.unmount_fused_linear_gelu(good)
+    assert torch.equal(good(x), ref)
+
+    no_bias = nn.Linear(8, 8, bias=False, device="cuda", dtype=torch.bfloat16)
+    assert not gelu.can_fuse_linear_gelu_static(no_bias)
+    assert not gelu.can_fuse_linear_gelu(good.proj, x.float())
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))