[Diffusion] [NPU] Add CI tests for FLUX (#19001)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>

.github/workflows/pr-test-npu.yml

@@ -281,3 +281,42 @@ jobs:
           export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}"
           cd python
           python3 sglang/multimodal_gen/test/run_suite.py --suite 1-npu
+
+  multimodal-gen-test-2-npu-a3:
+    needs: [check-changes, pr-gate]
+    if: needs.check-changes.outputs.multimodal_gen == 'true'
+    runs-on: linux-aarch64-a3-16
+    container:
+      image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.3.rc2-a3-ubuntu22.04-py3.11
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Install dependencies
+        run: |
+          # speed up by using infra cache services
+          CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local"
+          sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list
+          pip config set global.index-url http://${CACHING_URL}/pypi/simple
+          pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple"
+          pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn"
+
+          bash scripts/ci/npu/npu_ci_install_dependency.sh a3
+          # copy required file from our daily cache
+          cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp
+          # copy download through proxy
+          curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl
+
+      - name: Run test
+        timeout-minutes: 60
+        env:
+          SGLANG_USE_MODELSCOPE: true
+          SGLANG_IS_IN_CI: true
+          HF_ENDPOINT: https://hf-mirror.com
+          TORCH_EXTENSIONS_DIR: /tmp/torch_extensions
+          PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
+          STREAMS_PER_DEVICE: 32
+        run: |
+          export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}"
+          cd python
+          python3 sglang/multimodal_gen/test/run_suite.py --suite 2-npu

python/sglang/multimodal_gen/test/run_suite.py

@@ -56,7 +56,11 @@ suites_ascend = {
     "1-npu": [
         "ascend/test_server_1_npu.py",
         # add new 1-npu test files here
-    ]
+    ],
+    "2-npu": [
+        "ascend/test_server_2_npu.py",
+        # add new 2-npu test files here
+    ],
 }
 
 SUITES.update(suites_ascend)

python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json

@@ -5,6 +5,137 @@
         "description": "Reference numbers captured from the CI diffusion server baseline run"
     },
     "scenarios": {
+        "flux_image_t2i_npu": {
+            "stages_ms": {
+                "InputValidationStage": 0.07,
+                "TextEncodingStage": 154.51,
+                "TimestepPreparationStage": 53.52,
+                "LatentPreparationStage": 0.39,
+                "DenoisingStage": 19423.39,
+                "DecodingStage": 40.14
+            },
+            "denoise_step_ms": {
+                "0": 123.16,
+                "1": 91.7,
+                "2": 265.62,
+                "3": 402.68,
+                "4": 402.86,
+                "5": 402.78,
+                "6": 402.99,
+                "7": 402.77,
+                "8": 402.59,
+                "9": 402.93,
+                "10": 402.05,
+                "11": 402.99,
+                "12": 402.29,
+                "13": 403.07,
+                "14": 402.62,
+                "15": 402.99,
+                "16": 402.68,
+                "17": 403.0,
+                "18": 402.74,
+                "19": 402.85,
+                "20": 402.83,
+                "21": 403.03,
+                "22": 402.56,
+                "23": 402.84,
+                "24": 402.79,
+                "25": 402.95,
+                "26": 402.65,
+                "27": 403.01,
+                "28": 402.66,
+                "29": 402.92,
+                "30": 402.75,
+                "31": 403.0,
+                "32": 402.9,
+                "33": 402.48,
+                "34": 402.85,
+                "35": 402.03,
+                "36": 402.93,
+                "37": 402.3,
+                "38": 403.12,
+                "39": 402.83,
+                "40": 402.84,
+                "41": 402.75,
+                "42": 402.97,
+                "43": 402.62,
+                "44": 402.91,
+                "45": 402.81,
+                "46": 402.97,
+                "47": 402.57,
+                "48": 403.0,
+                "49": 402.75
+            },
+            "expected_e2e_ms": 23819.1,
+            "expected_avg_denoise_ms": 388.22,
+            "expected_median_denoise_ms": 402.82
+        },
+        "flux_2_image_t2i_2npu": {
+            "stages_ms": {
+                "InputValidationStage": 0.06,
+                "TextEncodingStage": 5628.31,
+                "ImageVAEEncodingStage": 0.01,
+                "LatentPreparationStage": 0.75,
+                "TimestepPreparationStage": 30.68,
+                "DenoisingStage": 55002.26,
+                "DecodingStage": 43.73
+            },
+            "denoise_step_ms": {
+                "0": 110.35,
+                "1": 301.82,
+                "2": 1139.81,
+                "3": 1114.17,
+                "4": 1099.34,
+                "5": 1099.12,
+                "6": 1100.16,
+                "7": 1099.67,
+                "8": 1099.09,
+                "9": 1089.81,
+                "10": 1109.73,
+                "11": 1099.97,
+                "12": 1100.26,
+                "13": 1099.67,
+                "14": 1099.79,
+                "15": 1099.6,
+                "16": 1100.16,
+                "17": 1099.87,
+                "18": 1100.02,
+                "19": 1099.34,
+                "20": 1099.6,
+                "21": 1099.45,
+                "22": 1100.2,
+                "23": 1099.29,
+                "24": 1098.86,
+                "25": 1090.38,
+                "26": 1109.19,
+                "27": 1099.67,
+                "28": 1100.06,
+                "29": 1099.22,
+                "30": 1100.08,
+                "31": 1098.86,
+                "32": 1099.73,
+                "33": 1099.11,
+                "34": 1100.13,
+                "35": 1103.97,
+                "36": 1095.26,
+                "37": 1099.38,
+                "38": 1099.34,
+                "39": 1099.17,
+                "40": 1100.08,
+                "41": 1089.89,
+                "42": 1106.69,
+                "43": 1102.57,
+                "44": 1100.17,
+                "45": 1099.21,
+                "46": 1100.42,
+                "47": 1099.38,
+                "48": 1099.59,
+                "49": 1099.47
+            },
+            "expected_e2e_ms": 64195.08,
+            "expected_avg_denoise_ms": 1065.0,
+            "expected_median_denoise_ms": 1099.63
+        },
         "wan2_1_t2v_1.3b_1_npu": {
             "stages_ms": {
                 "InputValidationStage": 0.07,

python/sglang/multimodal_gen/test/server/ascend/test_server_2_npu.py

@@ -0,0 +1,29 @@
+"""
+Config-driven diffusion performance test with pytest parametrization.
+
+
+If the actual run is significantly better than the baseline, the improved cases with their updated baseline will be printed
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.test.server.ascend.testcase_configs_npu import TWO_NPU_CASES
+from sglang.multimodal_gen.test.server.test_server_common import (  # noqa: F401
+    DiffusionServerBase,
+    diffusion_server,
+)
+from sglang.multimodal_gen.test.server.testcase_configs import DiffusionTestCase
+
+logger = init_logger(__name__)
+
+
+class TestDiffusionServerTwoNpu(DiffusionServerBase):
+    """Performance tests for 2-NPU diffusion cases."""
+
+    @pytest.fixture(params=TWO_NPU_CASES, ids=lambda c: c.id)
+    def case(self, request) -> DiffusionTestCase:
+        """Provide a DiffusionTestCase for each 2-NPU test."""
+        return request.param

python/sglang/multimodal_gen/test/server/ascend/testcase_configs_npu.py

@@ -3,9 +3,19 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
     DiffusionSamplingParams,
     DiffusionServerArgs,
     DiffusionTestCase,
+    T2I_sampling_params,
 )
 
 ONE_NPU_CASES: list[DiffusionTestCase] = [
+    # === Text to Image (T2I) ===
+    DiffusionTestCase(
+        "flux_image_t2i_npu",
+        DiffusionServerArgs(
+            model_path="/root/.cache/modelscope/hub/models/black-forest-labs/FLUX.1-dev",
+            modality="image",
+        ),
+        T2I_sampling_params,
+    ),
     # === Text to Video (T2V) ===
     DiffusionTestCase(
         "wan2_1_t2v_1.3b_1_npu",
@@ -19,3 +29,17 @@ ONE_NPU_CASES: list[DiffusionTestCase] = [
         ),
     ),
 ]
+
+TWO_NPU_CASES: list[DiffusionTestCase] = [
+    # === Text to Image (T2I) ===
+    DiffusionTestCase(
+        "flux_2_image_t2i_2npu",
+        DiffusionServerArgs(
+            model_path="/root/.cache/modelscope/hub/models/black-forest-labs/FLUX.2-dev",
+            modality="image",
+            num_gpus=2,
+            tp_size=2,
+        ),
+        T2I_sampling_params,
+    ),
+]