From 533851fbcb29dadafbb636dc9e9c4e482f93afc5 Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Wed, 17 Dec 2025 20:44:21 +0800
Subject: [PATCH] [diffusion] ci: add flux2 tp2 test into ci to avoid breaking
 tensor parallel (#15237)

---
 .../test/server/perf_baselines.json           | 67 +++++++++++++++++++
 .../test/server/test_server_common.py         |  3 +
 .../test/server/testcase_configs.py           | 14 ++++
 3 files changed, 84 insertions(+)

diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json
index 29951e088..705a57fe2 100644
--- a/python/sglang/multimodal_gen/test/server/perf_baselines.json
+++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json
@@ -1265,6 +1265,73 @@
             "expected_e2e_ms": 204257.12,
             "expected_avg_denoise_ms": 3855.55,
             "expected_median_denoise_ms": 3774.03
+        },
+        "flux_2_image_t2i_2_gpus": {
+            "stages_ms": {
+                "InputValidationStage": 0.05,
+                "TextEncodingStage": 518.88,
+                "ImageVAEEncodingStage": 0.0,
+                "ConditioningStage": 0.01,
+                "LatentPreparationStage": 0.45,
+                "TimestepPreparationStage": 3.41,
+                "DenoisingStage": 26377.63,
+                "DecodingStage": 321.94
+            },
+            "denoise_step_ms": {
+                "0": 66.94,
+                "1": 511.02,
+                "2": 537.24,
+                "3": 516.92,
+                "4": 537.18,
+                "5": 520.79,
+                "6": 531.95,
+                "7": 522.69,
+                "8": 532.77,
+                "9": 522.36,
+                "10": 530.91,
+                "11": 526.67,
+                "12": 530.23,
+                "13": 526.16,
+                "14": 528.55,
+                "15": 524.96,
+                "16": 526.36,
+                "17": 527.71,
+                "18": 528.84,
+                "19": 528.73,
+                "20": 527.95,
+                "21": 525.1,
+                "22": 528.21,
+                "23": 528.06,
+                "24": 527.11,
+                "25": 530.93,
+                "26": 526.27,
+                "27": 528.66,
+                "28": 526.33,
+                "29": 530.99,
+                "30": 528.51,
+                "31": 528.53,
+                "32": 528.73,
+                "33": 527.02,
+                "34": 527.49,
+                "35": 528.25,
+                "36": 529.96,
+                "37": 528.43,
+                "38": 528.5,
+                "39": 525.72,
+                "40": 529.65,
+                "41": 525.62,
+                "42": 526.24,
+                "43": 528.77,
+                "44": 526.68,
+                "45": 528.06,
+                "46": 524.92,
+                "47": 529.5,
+                "48": 524.22,
+                "49": 528.21
+            },
+            "expected_e2e_ms": 27624.8,
+            "expected_avg_denoise_ms": 518.23,
+            "expected_median_denoise_ms": 528.06
         }
     }
 }
diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py
index 8720108bd..02c80380b 100644
--- a/python/sglang/multimodal_gen/test/server/test_server_common.py
+++ b/python/sglang/multimodal_gen/test/server/test_server_common.py
@@ -54,6 +54,9 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
     extra_args = os.environ.get("SGLANG_TEST_SERVE_ARGS", "")
     extra_args += f" --num-gpus {server_args.num_gpus}"
 
+    if server_args.tp_size is not None:
+        extra_args += f" --tp-size {server_args.tp_size}"
+
     if server_args.ulysses_degree is not None:
         extra_args += f" --ulysses-degree {server_args.ulysses_degree}"
 
diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py
index 4c8845a73..679e60148 100644
--- a/python/sglang/multimodal_gen/test/server/testcase_configs.py
+++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py
@@ -119,6 +119,7 @@ class DiffusionServerArgs:
     custom_validator: str | None = None  # optional custom validator name
     # resources
     num_gpus: int = 1
+    tp_size: int | None = None
     ulysses_degree: int | None = None
     ring_degree: int | None = None
     # LoRA
@@ -518,6 +519,19 @@ TWO_GPU_CASES_B = [
             modality="image",
             warmup_text=1,
             warmup_edit=0,
+            num_gpus=2,
+        ),
+        T2I_sampling_params,
+    ),
+    DiffusionTestCase(
+        "flux_2_image_t2i_2_gpus",
+        DiffusionServerArgs(
+            model_path="black-forest-labs/FLUX.2-dev",
+            modality="image",
+            warmup_text=1,
+            warmup_edit=0,
+            num_gpus=2,
+            tp_size=2,
         ),
         T2I_sampling_params,
     ),