[Diffusion] Fuse FLUX.2 ModelOpt FP8 producers and QKV packing (#37162)

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-01 16:14:29 +08:00
parent 379e33d87e
commit 1c3ad92438
16 changed files with 1571 additions and 104 deletions
@@ -0,0 +1,48 @@
+import torch
+
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.ops.diffusion import (
+    fused_layernorm_modulate_fp8_quant_raw,
+    fused_layernorm_modulate_raw,
+)
+from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=8, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+HIDDEN = 6144
+EPS = 1e-6
+
+
+@marker.parametrize("rows", [512, 4096, 4608], [512])
+@marker.benchmark("impl", ["split", "fused"], unit="us")
+def benchmark(rows: int, impl: str):
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260831 + rows)
+    x = torch.randn((1, rows, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator)
+    scale = torch.randn((1, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator)
+    shift = torch.randn((1, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator)
+    input_scale = torch.tensor(0.03125, dtype=torch.float32, device=DEVICE)
+
+    if impl == "split":
+
+        def fn():
+            normalized = fused_layernorm_modulate_raw(x, scale, shift, EPS)
+            return static_quant_fp8(normalized, input_scale)[0]
+
+    else:
+
+        def fn():
+            return fused_layernorm_modulate_fp8_quant_raw(
+                x, scale, shift, input_scale, EPS
+            )
+
+    return marker.do_bench(fn, disable_log_bandwidth=True)
+
+
+if __name__ == "__main__":
+    benchmark.run()
@@ -0,0 +1,46 @@
+import torch
+
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.ops.diffusion import try_flux2_token_cat_fp8
+from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=8, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+
+
+@marker.parametrize("tokens", [512, 4096], [4096])
+@marker.benchmark("impl", ["cat_then_quant", "fused"], unit="us")
+def benchmark(tokens: int, impl: str):
+    generator = torch.Generator(device="cuda")
+    generator.manual_seed(20260831 + tokens)
+    attention = torch.randn(
+        (1, tokens, 6144),
+        dtype=torch.bfloat16,
+        device="cuda",
+        generator=generator,
+    )
+    mlp = torch.randn(
+        (1, tokens, 18432),
+        dtype=torch.bfloat16,
+        device="cuda",
+        generator=generator,
+    )
+    input_scale = torch.tensor([0.013], dtype=torch.float32, device="cuda")
+
+    if impl == "cat_then_quant":
+
+        def fn():
+            return static_quant_fp8(torch.cat([attention, mlp], dim=-1), input_scale)[0]
+
+    else:
+
+        def fn():
+            return try_flux2_token_cat_fp8(attention, mlp, input_scale)
+
+    return marker.do_bench(fn, use_cuda_graph=False, disable_log_bandwidth=True)
+
+
+if __name__ == "__main__":
+    benchmark.run()
@@ -0,0 +1,52 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    fused_layernorm_modulate_fp8_quant_raw,
+    fused_layernorm_modulate_raw,
+)
+from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+HIDDEN = 6144
+EPS = 1e-6
+
+
+def _make_inputs(rows: int):
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260831 + rows)
+    x = torch.randn((1, rows, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator)
+    scale = torch.randn((1, 1, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator)
+    shift = torch.randn((1, 1, HIDDEN), dtype=DTYPE, device=DEVICE, generator=generator)
+    return x, scale, shift
+
+
+@pytest.mark.parametrize("rows", [1, 127, 512])
+@pytest.mark.parametrize(
+    "input_scale_value", [0.005, 0.03125, 0.25, 0.4754464328289032, 1.0]
+)
+def test_flux2_layernorm_modulate_fp8_is_bit_exact(
+    rows: int, input_scale_value: float
+) -> None:
+    x, scale, shift = _make_inputs(rows)
+    input_scale = torch.tensor(input_scale_value, dtype=torch.float32, device=DEVICE)
+
+    normalized = fused_layernorm_modulate_raw(
+        x, scale.squeeze(1), shift.squeeze(1), EPS
+    )
+    expected, _ = static_quant_fp8(normalized, input_scale)
+    actual = fused_layernorm_modulate_fp8_quant_raw(
+        x, scale.squeeze(1), shift.squeeze(1), input_scale, EPS
+    )
+
+    assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))
@@ -0,0 +1,158 @@
+import sys
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    try_fused_flux2_qkv_epilogue,
+)
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    RMSNorm,
+    apply_qk_norm_with_optional_rope,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+DEVICE = "cuda"
+DTYPE = torch.bfloat16
+HEAD_DIM = 128
+
+
+def _packed_qkv(tokens: int, heads: int, generator: torch.Generator):
+    source = [
+        torch.randn(
+            (1, tokens, heads, HEAD_DIM),
+            dtype=DTYPE,
+            device=DEVICE,
+            generator=generator,
+        )
+        for _ in range(3)
+    ]
+    packed = torch.cat([tensor.flatten(2) for tensor in source], dim=-1)
+    views = [
+        tensor.unflatten(-1, (heads, HEAD_DIM)) for tensor in packed.chunk(3, dim=-1)
+    ]
+    assert all(not tensor.is_contiguous() for tensor in views)
+    return views
+
+
+@pytest.mark.parametrize("img_tokens,txt_tokens,heads", [(17, 7, 4), (256, 64, 8)])
+def test_flux2_qkv_epilogue_is_bit_exact(
+    img_tokens: int, txt_tokens: int, heads: int
+) -> None:
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260831 + img_tokens)
+    img_qkv = _packed_qkv(img_tokens, heads, generator)
+    txt_qkv = _packed_qkv(txt_tokens, heads, generator)
+    norms = [
+        RMSNorm(HEAD_DIM, eps=1e-6).to(device=DEVICE, dtype=DTYPE) for _ in range(4)
+    ]
+    for norm in norms:
+        norm.weight.data.normal_(generator=generator)
+
+    angles = torch.randn(
+        (img_tokens + txt_tokens, HEAD_DIM // 2),
+        device=DEVICE,
+        generator=generator,
+    )
+    cache = torch.cat([angles.cos(), angles.sin()], dim=-1).contiguous()
+
+    img_reference = [tensor.contiguous() for tensor in img_qkv]
+    txt_reference = [tensor.contiguous() for tensor in txt_qkv]
+    txt_reference[0], txt_reference[1] = apply_qk_norm_with_optional_rope(
+        txt_reference[0],
+        txt_reference[1],
+        norms[2],
+        norms[3],
+        HEAD_DIM,
+        cache,
+        is_neox=False,
+    )
+    img_reference[0], img_reference[1] = apply_qk_norm_with_optional_rope(
+        img_reference[0],
+        img_reference[1],
+        norms[0],
+        norms[1],
+        HEAD_DIM,
+        cache,
+        is_neox=False,
+        position_offset=txt_tokens,
+    )
+    expected = tuple(
+        torch.cat([txt_reference[index], img_reference[index]], dim=1)
+        for index in range(3)
+    )
+
+    actual = try_fused_flux2_qkv_epilogue(
+        *img_qkv,
+        *txt_qkv,
+        norms[0].weight,
+        norms[1].weight,
+        norms[2].weight,
+        norms[3].weight,
+        cache,
+        1e-6,
+        1e-6,
+    )
+
+    assert actual is not None
+    assert all(
+        torch.equal(result, reference)
+        for result, reference in zip(actual, expected, strict=True)
+    )
+
+
+def test_flux2_qkv_epilogue_rejects_compile() -> None:
+    tensor = torch.empty((1, 1, 1, HEAD_DIM), device=DEVICE, dtype=DTYPE)
+    weight = torch.empty((HEAD_DIM,), device=DEVICE, dtype=DTYPE)
+    cache = torch.empty((2, HEAD_DIM), device=DEVICE, dtype=torch.float32)
+    with patch("torch.compiler.is_compiling", return_value=True):
+        assert (
+            try_fused_flux2_qkv_epilogue(
+                tensor,
+                tensor,
+                tensor,
+                tensor,
+                tensor,
+                tensor,
+                weight,
+                weight,
+                weight,
+                weight,
+                cache,
+                1e-6,
+                1e-6,
+            )
+            is None
+        )
+
+
+def test_flux2_qkv_epilogue_rejects_cuda_graph_capture() -> None:
+    tensor = torch.empty((1, 1, 1, HEAD_DIM), device=DEVICE, dtype=DTYPE)
+    weight = torch.empty((HEAD_DIM,), device=DEVICE, dtype=DTYPE)
+    cache = torch.empty((2, HEAD_DIM), device=DEVICE, dtype=torch.float32)
+    with patch("torch.cuda.is_current_stream_capturing", return_value=True):
+        assert (
+            try_fused_flux2_qkv_epilogue(
+                tensor,
+                tensor,
+                tensor,
+                tensor,
+                tensor,
+                tensor,
+                weight,
+                weight,
+                weight,
+                weight,
+                cache,
+                1e-6,
+                1e-6,
+            )
+            is None
+        )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))
@@ -0,0 +1,56 @@
+import sys
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import try_flux2_token_cat_fp8
+from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+
+@pytest.mark.parametrize("tokens", [1, 127, 4096])
+def test_flux2_token_cat_fp8_is_bit_exact(tokens: int) -> None:
+    generator = torch.Generator(device="cuda")
+    generator.manual_seed(20260831 + tokens)
+    attention = torch.randn(
+        (1, tokens, 6144),
+        dtype=torch.bfloat16,
+        device="cuda",
+        generator=generator,
+    )
+    mlp = torch.randn(
+        (1, tokens, 18432),
+        dtype=torch.bfloat16,
+        device="cuda",
+        generator=generator,
+    )
+    scale = torch.tensor([0.013], dtype=torch.float32, device="cuda")
+
+    expected, _ = static_quant_fp8(torch.cat([attention, mlp], dim=-1), scale)
+    actual = try_flux2_token_cat_fp8(attention, mlp, scale)
+
+    assert actual is not None
+    assert torch.equal(actual, expected)
+
+
+def test_flux2_token_cat_fp8_rejects_compile() -> None:
+    attention = torch.empty((1, 1, 16), device="cuda", dtype=torch.bfloat16)
+    mlp = torch.empty((1, 1, 48), device="cuda", dtype=torch.bfloat16)
+    scale = torch.ones((1,), device="cuda", dtype=torch.float32)
+    with patch("torch.compiler.is_compiling", return_value=True):
+        assert try_flux2_token_cat_fp8(attention, mlp, scale) is None
+
+
+def test_flux2_token_cat_fp8_rejects_cuda_graph_capture() -> None:
+    attention = torch.empty((1, 1, 16), device="cuda", dtype=torch.bfloat16)
+    mlp = torch.empty((1, 1, 48), device="cuda", dtype=torch.bfloat16)
+    scale = torch.ones((1,), device="cuda", dtype=torch.float32)
+    with patch("torch.cuda.is_current_stream_capturing", return_value=True):
+        assert try_flux2_token_cat_fp8(attention, mlp, scale) is None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))