[diffusion] FLUX.1 fused adaLN modulate (bit-exact) + RoPE cache hoist, LN-affine folding behind quality=high (H200 e2e -3.5% lossless / -6.9% high) (#34004)

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

@@ -0,0 +1,56 @@
+import pytest
+import torch
+import torch.nn as nn
+
+from sglang.kernels.ops.diffusion.fused_ln_modulate import (
+    can_fuse_ln_modulate,
+    fused_ln_modulate,
+    fused_ln_modulate_active,
+    mark_fused_ln_modulate_site,
+    mount_fused_ln_modulate,
+    unmount_fused_ln_modulate,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+@pytest.fixture(autouse=True)
+def cuda_setup():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    torch.cuda.manual_seed(0)
+
+
+@pytest.mark.parametrize("seq_len", [4096, 512])
+def test_fused_ln_modulate_matches_reference(seq_len):
+    x = torch.randn((1, seq_len, 3072), device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn((1, 3072), device="cuda", dtype=torch.bfloat16)
+    shift = torch.randn_like(scale)
+    assert can_fuse_ln_modulate(x, scale, shift)
+    out = fused_ln_modulate(x, scale, shift, eps=1e-6)
+    norm = nn.LayerNorm(3072, eps=1e-6, elementwise_affine=False).cuda()
+    ref = norm(x) * (1 + scale[:, None]) + shift[:, None]
+    # Contract: bf16 rounding-order-level difference only, not bit-exact.
+    torch.testing.assert_close(out, ref, atol=0.0625, rtol=0.05)
+
+
+def test_fused_ln_modulate_guards_and_mount_protocol():
+    x = torch.randn((2, 64, 3072), device="cuda", dtype=torch.bfloat16)
+    row = torch.randn((2, 3072), device="cuda", dtype=torch.bfloat16)
+    assert not can_fuse_ln_modulate(x, row, row)  # folded affine needs B == 1
+    root = nn.Module()
+    root.child = nn.Module()
+    mark_fused_ln_modulate_site(root.child)
+    assert not fused_ln_modulate_active(root.child)
+    assert mount_fused_ln_modulate(root)
+    assert fused_ln_modulate_active(root.child)
+    unmount_fused_ln_modulate(root)
+    assert not fused_ln_modulate_active(root.child)
+    assert not mount_fused_ln_modulate(nn.Module())  # no marked sites
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__]))

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

@@ -0,0 +1,56 @@
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion.modulate_scale_shift import (
+    can_use_modulate_scale_shift_cuda,
+    modulate_scale_shift_cuda,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+# FLUX.1 1024^2 adaLN shapes (D=3072) plus batched and odd-length coverage.
+CASES = [(1, 4096, 3072), (1, 512, 3072), (1, 4608, 3072), (2, 1024, 3072), (1, 17, 64)]
+
+
+@pytest.fixture(autouse=True)
+def cuda_setup():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    torch.cuda.manual_seed(0)
+
+
+def _eager(x, scale, shift):
+    return x * (1 + scale[:, None]) + shift[:, None]
+
+
+@pytest.mark.parametrize("shape", CASES)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+def test_modulate_scale_shift_matches_eager(shape, dtype):
+    x = torch.randn(shape, device="cuda", dtype=dtype)
+    scale = torch.randn((shape[0], shape[-1]), device="cuda", dtype=dtype)
+    shift = torch.randn_like(scale)
+    out = modulate_scale_shift_cuda(x, scale, shift)
+    assert torch.equal(out, _eager(x, scale, shift))  # bitwise contract
+
+
+def test_modulate_scale_shift_adaln_chunk_views():
+    x = torch.randn((1, 4096, 3072), device="cuda", dtype=torch.bfloat16)
+    emb = torch.randn((1, 6 * 3072), device="cuda", dtype=torch.bfloat16)
+    shift, scale = emb.chunk(6, dim=1)[:2]
+    assert can_use_modulate_scale_shift_cuda(x, scale, shift)
+    out = modulate_scale_shift_cuda(x, scale, shift)
+    assert torch.equal(out, _eager(x, scale, shift))
+
+
+def test_modulate_scale_shift_guards_reject_fp32():
+    x = torch.randn((1, 64, 64), device="cuda", dtype=torch.float32)
+    row = torch.randn((1, 64), device="cuda", dtype=torch.float32)
+    assert not can_use_modulate_scale_shift_cuda(x, row, row)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__]))