[diffusion] FLUX.1 bit-exact residual-gate fast path + tanh-GELU epilogue behind quality=high (H200 e2e -1.1% lossless / -4.3% high) (#33819)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

test/registered/kernels/ops/diffusion/test_fused_linear_gelu.py

@@ -37,6 +37,23 @@ def test_fused_matches_reference(dtype):
     torch.testing.assert_close(site(x), ref, atol=atol, rtol=2e-2)
 
 
+def test_flux_gelu_proj_site():
+    """FLUX.1 shared-FF site: gate off is bit-exact, gate on is close."""
+    from sglang.multimodal_gen.runtime.models.dits.flux import FluxFusedGELUProj
+
+    torch.manual_seed(0)
+    proj = nn.Linear(3072, 12288, device="cuda", dtype=torch.bfloat16)
+    site = FluxFusedGELUProj(proj)
+    x = torch.randn(1, 512, 3072, device="cuda", dtype=torch.bfloat16)
+    ref = F.gelu(proj(x), approximate="tanh")
+
+    assert torch.equal(site(x), ref)  # unmounted default: bit-exact reference
+    assert gelu.mount_fused_linear_gelu(site)
+    torch.testing.assert_close(site(x), ref, atol=2e-2, rtol=2e-2)
+    gelu.unmount_fused_linear_gelu(site)
+    assert torch.equal(site(x), ref)
+
+
 def test_mount_guards_and_lossless_path():
     torch.manual_seed(0)
     good, bad = _Site(), _Site(torch.float32)

test/registered/kernels/ops/diffusion/test_residual_gate_add.py

@@ -18,6 +18,10 @@ CASES = [
     ((1, 512, 4096), (1, 512, 4096)),
     ((1, 17, 65), (1, 1, 65)),
     ((1, 17, 65), (1, 17, 65)),
+    # FLUX.1 1024^2 shapes: dual-stream image/text and single-stream joint.
+    ((1, 4096, 3072), (1, 1, 3072)),
+    ((1, 512, 3072), (1, 1, 3072)),
+    ((1, 4608, 3072), (1, 1, 3072)),
 ]