[Diffusion] Fuse FLUX.2 token concatenation and NVFP4 quantization (#37141)

test/manual/kernels/bench_flux2_token_cat_nvfp4.py

@@ -0,0 +1,90 @@
+import time
+
+import flashinfer
+import torch
+
+from sglang.kernels.ops.diffusion import try_flux2_token_cat_nvfp4
+
+
+def _benchmark(fn, iterations: int = 100) -> float:
+    for _ in range(10):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iterations):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) * 1000 / iterations
+
+
+def _benchmark_wall(fn, iterations: int = 100) -> float:
+    for _ in range(10):
+        fn()
+    torch.cuda.synchronize()
+    start = time.perf_counter_ns()
+    for _ in range(iterations):
+        fn()
+    torch.cuda.synchronize()
+    return (time.perf_counter_ns() - start) / iterations / 1000
+
+
+def _run_case(token_count: int) -> None:
+    generator = torch.Generator(device="cuda")
+    generator.manual_seed(20260830 + token_count)
+    attention = torch.randn(
+        1,
+        token_count,
+        6144,
+        device="cuda",
+        dtype=torch.bfloat16,
+        generator=generator,
+    )
+    mlp = torch.randn(
+        1,
+        token_count,
+        18432,
+        device="cuda",
+        dtype=torch.bfloat16,
+        generator=generator,
+    )
+    global_scale = torch.tensor(0.625, device="cuda", dtype=torch.float32)
+
+    def baseline():
+        return flashinfer.fp4_quantize(
+            torch.cat([attention, mlp], dim=-1).view(-1, 24576), global_scale
+        )
+
+    def fused():
+        result = try_flux2_token_cat_nvfp4(attention, mlp, global_scale)
+        assert result is not None
+        return result
+
+    expected = baseline()
+    actual = fused()
+    exact = [torch.equal(lhs, rhs) for lhs, rhs in zip(actual, expected)]
+    baseline_us = _benchmark(baseline)
+    fused_us = _benchmark(fused)
+    baseline_wall_us = _benchmark_wall(baseline)
+    fused_wall_us = _benchmark_wall(fused)
+    print(
+        {
+            "tokens": token_count,
+            "baseline_us": baseline_us,
+            "fused_us": fused_us,
+            "speedup": baseline_us / fused_us,
+            "baseline_wall_us": baseline_wall_us,
+            "fused_wall_us": fused_wall_us,
+            "wall_speedup": baseline_wall_us / fused_wall_us,
+            "exact": exact,
+        }
+    )
+
+
+if __name__ == "__main__":
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3):
+        raise RuntimeError("This benchmark requires an NVIDIA Blackwell SM103 GPU")
+    for tokens in (17, 512, 4096, 4608):
+        _run_case(tokens)

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

@@ -51,6 +51,7 @@ from sglang.kernels.ops.diffusion import (
     mount_fused_ln_modulate,
     mount_hunyuan_qknorm,
     mount_ltx2_rms_norm_modulate,
+    try_flux2_token_cat_nvfp4,
     unmount_hunyuan_qknorm,
     unmount_ltx2_rms_norm_modulate,
     wan_rmsnorm_silu,
@@ -371,6 +372,40 @@ class TestFlux2EagerFusions(CustomTestCase):
         self.assertTrue(torch.equal(actual, expected))
         self.assertEqual(len(flux2._FLUX2_SWIGLU_SIGS), 1)
 
+    @pytest.mark.skipif(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3),
+        reason="FLUX.2 token-cat NVFP4 requires Blackwell SM103",
+    )
+    def test_token_cat_nvfp4_matches_flashinfer(self):
+        import flashinfer
+
+        torch.manual_seed(20260830)
+        attention = torch.randn(1, 17, 6144, device="cuda", dtype=torch.bfloat16)
+        mlp = torch.randn(1, 17, 18432, device="cuda", dtype=torch.bfloat16)
+        global_scale = torch.tensor(0.625, device="cuda", dtype=torch.float32)
+
+        expected_fp4, expected_scales = flashinfer.fp4_quantize(
+            torch.cat([attention, mlp], dim=-1).view(-1, 24576), global_scale
+        )
+        actual = try_flux2_token_cat_nvfp4(attention, mlp, global_scale)
+
+        self.assertIsNotNone(actual)
+        actual_fp4, actual_scales = actual
+        self.assertTrue(torch.equal(actual_fp4, expected_fp4))
+        self.assertTrue(
+            torch.equal(
+                actual_scales.view(torch.uint8), expected_scales.view(torch.uint8)
+            )
+        )
+
+    def test_token_cat_nvfp4_falls_back_while_compiling(self):
+        attention = torch.empty(1, 1, 6144, device="cuda", dtype=torch.bfloat16)
+        mlp = torch.empty(1, 1, 18432, device="cuda", dtype=torch.bfloat16)
+        global_scale = torch.ones(1, device="cuda", dtype=torch.float32)
+
+        with patch("torch.compiler.is_compiling", return_value=True):
+            self.assertIsNone(try_flux2_token_cat_nvfp4(attention, mlp, global_scale))
+
 
 # -------------------------------------------------------------------------
 # Qwen-Image -- reuse timestep-only modulation across serial CFG branches

test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py

@@ -10,6 +10,7 @@ from sglang.multimodal_gen.runtime.layers.quantization import (
 from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
     ModelOptFp4Config,
     ModelOptFp4LinearMethod,
+    apply_nvfp4_gemm_prequantized,
 )
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.srt.layers.quantization.modelopt_quant import pad_nvfp4_weight
@@ -363,6 +364,40 @@ def test_flux2_shape_correctness_flashinfer_trtllm(
     assert diff < DEEPGEMM_FP4_MAX_DIFF, f"{m=}, {n=}, {k=}, {diff=:.6f}"
 
 
+@pytest.mark.skipif(
+    not _nvfp4_supported(),
+    reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",
+)
+@pytest.mark.parametrize(
+    "backend", [None, "flashinfer_trtllm"], ids=["default", "flashinfer_trtllm"]
+)
+def test_prequantized_input_matches_regular_apply(
+    monkeypatch: pytest.MonkeyPatch, backend: str | None
+) -> None:
+    _set_diffusion_fp4_backend(monkeypatch, backend)
+    m, n, k = 19, 150, 80
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260831)
+    x = torch.randn((m, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    weight = torch.randn((n, k), device=DEVICE, dtype=DTYPE, generator=generator)
+    input_global_scale = _make_global_scale(x)
+    weight_global_scale = _make_global_scale(weight)
+    weight_fp4, weight_scale_linear = _quantize_weight_for_checkpoint(
+        weight, weight_global_scale
+    )
+    method, layer = _build_layer(
+        weight_fp4, weight_scale_linear, input_global_scale, weight_global_scale
+    )
+
+    expected = method.apply(layer, x)
+    x_fp4, x_scale_interleaved = flashinfer.fp4_quantize(x, input_global_scale)
+    actual = apply_nvfp4_gemm_prequantized(
+        layer, x_fp4, x_scale_interleaved, output_dtype=x.dtype
+    )
+
+    assert torch.equal(actual, expected)
+
+
 @pytest.mark.skipif(
     not _nvfp4_supported(),
     reason="Diffusion NVFP4 scaled mm correctness requires Blackwell GPUs",