[diffusion] CI: add initial nvfp4 ci test for b200 (#21767)

Co-authored-by: Mick <mickjagger19@icloud.com>

.github/workflows/pr-test-multimodal-gen.yml

@@ -9,6 +9,9 @@ on:
       sgl_kernel:
         required: true
         type: string
+      b200_runner:
+        required: true
+        type: string
       continue_on_error:
         required: false
         type: string
@@ -258,6 +261,52 @@ jobs:
       - uses: ./.github/actions/upload-cuda-coredumps
         if: always()
 
+  multimodal-gen-test-1-b200:
+    if: |
+      (inputs.target_stage == 'multimodal-gen-test-1-b200') ||
+      (
+        !inputs.target_stage &&
+        ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) &&
+        inputs.multimodal_gen == 'true'
+      )
+    runs-on: ${{ inputs.b200_runner }}
+    timeout-minutes: 240
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+
+      - uses: ./.github/actions/check-maintenance
+
+      - name: Download artifacts
+        if: inputs.sgl_kernel == 'true'
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-python3.10-cuda12.9
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
+
+      - name: Run diffusion server tests
+        timeout-minutes: 240
+        env:
+          RUNAI_STREAMER_MEMORY_LIMIT: 0
+          CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
+        run: |
+          cd python
+          python3 sglang/multimodal_gen/test/run_suite.py \
+            --suite 1-gpu-b200 \
+            $CONTINUE_ON_ERROR_FLAG
+
+      - uses: ./.github/actions/upload-cuda-coredumps
+        if: always()
+
   multimodal-gen-unit-test:
     if: |
       (inputs.target_stage == 'multimodal-gen-unit-test') ||

.github/workflows/pr-test.yml

@@ -892,6 +892,7 @@ jobs:
         inputs.target_stage == 'multimodal-gen-test-2-gpu' ||
         inputs.target_stage == 'multimodal-gen-component-accuracy-1-gpu' ||
         inputs.target_stage == 'multimodal-gen-component-accuracy-2-gpu' ||
+        inputs.target_stage == 'multimodal-gen-test-1-b200' ||
         inputs.target_stage == 'multimodal-gen-unit-test' ||
         (
           !inputs.target_stage &&
@@ -903,6 +904,7 @@ jobs:
     with:
       multimodal_gen: ${{ needs.check-changes.outputs.multimodal_gen }}
       sgl_kernel: ${{ needs.check-changes.outputs.sgl_kernel }}
+      b200_runner: ${{ needs.check-changes.outputs.b200_runner }}
       continue_on_error: ${{ needs.check-changes.outputs.continue_on_error }}
       pr_head_sha: ${{ inputs.pr_head_sha || '' }}
       git_ref: ${{ inputs.git_ref || '' }}

python/sglang/multimodal_gen/test/run_suite.py

@@ -52,6 +52,9 @@ SUITES = {
         "test_server_2_gpu_b.py",
         # add new 2-gpu test files here
     ],
+    "1-gpu-b200": [
+        "test_server_c.py",
+    ],
 }
 
 suites_ascend = {
@@ -80,7 +83,7 @@ def parse_args():
         type=str,
         required=True,
         choices=list(SUITES.keys()),
-        help="The test suite to run (e.g., 1-gpu, 2-gpu)",
+        help="The test suite to run (valid names are defined in SUITES)",
     )
     parser.add_argument(
         "--partition-id",

python/sglang/multimodal_gen/test/scripts/gen_diffusion_ci_outputs.py

@@ -8,6 +8,7 @@ ensuring that GT generation uses exactly the same code path as CI tests.
 Usage:
     python gen_diffusion_ci_outputs.py --suite 1-gpu --partition-id 0 --total-partitions 2 --out-dir ./output
     python gen_diffusion_ci_outputs.py --suite 1-gpu --case-ids qwen_image_t2i flux_image_t2i --out-dir ./output
+    python gen_diffusion_ci_outputs.py --suite 1-gpu-b200 --out-dir ./output
 """
 
 import argparse
@@ -27,9 +28,9 @@ def main():
     parser.add_argument(
         "--suite",
         type=str,
-        choices=["1-gpu", "2-gpu"],
+        choices=list(SUITES.keys()),
         required=True,
-        help="Test suite to run (1-gpu or 2-gpu)",
+        help="Test suite to run (choices: " + ", ".join(list(SUITES.keys())) + ")",
     )
     parser.add_argument(
         "--partition-id",

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -2375,220 +2375,71 @@
             "expected_avg_denoise_ms": 0.0,
             "expected_median_denoise_ms": 0.0
         },
-        "helios_base_t2v": {
+        "flux_2_nvfp4_t2i": {
-            "stages_ms": {
-                "InputValidationStage": 0.04,
-                "TextEncodingStage": 1102.45,
-                "LatentPreparationStage": 0.14,
-                "HeliosChunkedDenoisingStage": 116964.69,
-                "HeliosDecodingStage": 664.76,
-                "per_frame_generation": null
-            },
-            "denoise_step_ms": {
-                "0": 1893.3,
-                "1": 1900.93,
-                "2": 1934.08,
-                "3": 1897.65,
-                "4": 1907.59,
-                "5": 1909.1,
-                "6": 1911.51,
-                "7": 1909.25,
-                "8": 1911.69,
-                "9": 1911.77,
-                "10": 1913.35,
-                "11": 1915.44,
-                "12": 1912.11,
-                "13": 1910.08,
-                "14": 1911.77,
-                "15": 1908.22,
-                "16": 1908.83,
-                "17": 1910.11,
-                "18": 1908.19,
-                "19": 1911.99,
-                "20": 1909.96,
-                "21": 1910.32,
-                "22": 1911.76,
-                "23": 1911.87,
-                "24": 1908.91,
-                "25": 1912.41,
-                "26": 1913.15,
-                "27": 1908.34,
-                "28": 1913.21,
-                "29": 1911.98,
-                "30": 1912.16,
-                "31": 1914.17,
-                "32": 1911.45,
-                "33": 1912.5,
-                "34": 1914.48,
-                "35": 1912.64,
-                "36": 1912.24,
-                "37": 1914.48,
-                "38": 1911.06,
-                "39": 1915.45,
-                "40": 1914.0,
-                "41": 1912.99,
-                "42": 1913.68,
-                "43": 1914.09,
-                "44": 1915.83,
-                "45": 1913.36,
-                "46": 1914.84,
-                "47": 1915.31,
-                "48": 1915.58,
-                "49": 1912.63
-            },
-            "expected_e2e_ms": 118821.41,
-            "expected_avg_denoise_ms": 1911.64,
-            "expected_median_denoise_ms": 1912.05
-        },
-        "helios_mid_t2v": {
             "stages_ms": {
                 "InputValidationStage": 0.09,
-                "TextEncodingStage": 1102.28,
+                "TextEncodingStage": 458.68,
-                "LatentPreparationStage": 0.23,
+                "ImageVAEEncodingStage": 0.01,
-                "HeliosChunkedDenoisingStage": 77947.9,
+                "LatentPreparationStage": 0.54,
-                "HeliosDecodingStage": 664.96,
+                "TimestepPreparationStage": 20.88,
-                "per_frame_generation": null
+                "DenoisingStage": 7189.58,
+                "DecodingStage": 13.55
             },
             "denoise_step_ms": {
-                "0": 404.46,
+                "0": 94.51,
-                "1": 404.88,
+                "1": 88.17,
-                "2": 405.35,
+                "2": 132.17,
-                "3": 406.01,
+                "3": 141.26,
-                "4": 404.97,
+                "4": 142.49,
-                "5": 405.07,
+                "5": 141.41,
-                "6": 405.06,
+                "6": 142.62,
-                "7": 404.98,
+                "7": 140.92,
-                "8": 405.39,
+                "8": 141.08,
-                "9": 405.52,
+                "9": 142.68,
-                "10": 405.76,
+                "10": 139.88,
-                "11": 405.53,
+                "11": 144.98,
-                "12": 405.16,
+                "12": 144.37,
-                "13": 405.46,
+                "13": 142.0,
-                "14": 405.75,
+                "14": 142.49,
-                "15": 405.69,
+                "15": 141.24,
-                "16": 405.26,
+                "16": 141.05,
-                "17": 405.23,
+                "17": 140.69,
-                "18": 405.42,
+                "18": 141.48,
-                "19": 405.99,
+                "19": 141.92,
-                "20": 663.39,
+                "20": 146.34,
-                "21": 666.6,
+                "21": 147.32,
-                "22": 665.73,
+                "22": 140.68,
-                "23": 666.37,
+                "23": 141.09,
-                "24": 667.43,
+                "24": 142.51,
-                "25": 668.28,
+                "25": 140.83,
-                "26": 667.96,
+                "26": 145.73,
-                "27": 668.93,
+                "27": 148.47,
-                "28": 667.78,
+                "28": 144.86,
-                "29": 668.15,
+                "29": 140.83,
-                "30": 668.91,
+                "30": 144.76,
-                "31": 667.22,
+                "31": 145.36,
-                "32": 669.31,
+                "32": 140.58,
-                "33": 666.57,
+                "33": 144.49,
-                "34": 669.78,
+                "34": 142.65,
-                "35": 668.38,
+                "35": 141.86,
-                "36": 669.95,
+                "36": 148.19,
-                "37": 668.76,
+                "37": 145.5,
-                "38": 667.82,
+                "38": 145.68,
-                "39": 668.98,
+                "39": 143.64,
-                "40": 1891.05,
+                "40": 143.7,
-                "41": 1893.52,
+                "41": 153.8,
-                "42": 1893.48,
+                "42": 148.57,
-                "43": 1892.79,
+                "43": 143.2,
-                "44": 1892.03,
+                "44": 144.15,
-                "45": 1892.87,
+                "45": 142.11,
-                "46": 1895.55,
+                "46": 146.38,
-                "47": 1892.19,
+                "47": 146.97,
-                "48": 1892.89,
+                "48": 144.62,
-                "49": 1892.32,
+                "49": 146.38
-                "50": 1890.25,
-                "51": 1894.1,
-                "52": 1890.67,
-                "53": 1892.09,
-                "54": 1892.64,
-                "55": 1891.91,
-                "56": 1894.27,
-                "57": 1893.62,
-                "58": 1892.65,
-                "59": 1891.9
             },
-            "expected_e2e_ms": 79824.32,
+            "expected_e2e_ms": 8091.46,
-            "expected_avg_denoise_ms": 988.6,
+            "expected_avg_denoise_ms": 141.37,
-            "expected_median_denoise_ms": 668.05
+            "expected_median_denoise_ms": 142.63
-        },
-        "helios_distilled_t2v": {
-            "stages_ms": {
-                "InputValidationStage": 0.05,
-                "TextEncodingStage": 552.02,
-                "LatentPreparationStage": 0.13,
-                "HeliosChunkedDenoisingStage": 57879.88,
-                "HeliosDecodingStage": 663.31,
-                "per_frame_generation": null
-            },
-            "denoise_step_ms": {
-                "0": 207.03,
-                "1": 204.36,
-                "2": 203.87,
-                "3": 204.51,
-                "4": 206.21,
-                "5": 205.54,
-                "6": 205.06,
-                "7": 205.45,
-                "8": 205.96,
-                "9": 205.95,
-                "10": 205.22,
-                "11": 204.43,
-                "12": 205.14,
-                "13": 205.06,
-                "14": 205.11,
-                "15": 206.09,
-                "16": 205.1,
-                "17": 204.99,
-                "18": 204.55,
-                "19": 205.14,
-                "20": 337.47,
-                "21": 337.06,
-                "22": 337.68,
-                "23": 336.58,
-                "24": 335.98,
-                "25": 335.84,
-                "26": 336.01,
-                "27": 335.61,
-                "28": 335.79,
-                "29": 335.62,
-                "30": 336.69,
-                "31": 335.98,
-                "32": 336.15,
-                "33": 336.55,
-                "34": 336.98,
-                "35": 337.33,
-                "36": 336.34,
-                "37": 335.94,
-                "38": 336.69,
-                "39": 336.14,
-                "40": 954.88,
-                "41": 956.2,
-                "42": 953.9,
-                "43": 953.49,
-                "44": 957.1,
-                "45": 956.95,
-                "46": 955.02,
-                "47": 954.98,
-                "48": 956.0,
-                "49": 956.63,
-                "50": 958.66,
-                "51": 957.26,
-                "52": 956.73,
-                "53": 955.06,
-                "54": 957.04,
-                "55": 958.07,
-                "56": 958.28,
-                "57": 957.99,
-                "58": 957.61,
-                "59": 956.98
-            },
-            "expected_e2e_ms": 59168.9,
-            "expected_avg_denoise_ms": 499.37,
-            "expected_median_denoise_ms": 336.25
         }
     }
 }

python/sglang/multimodal_gen/test/server/test_server_c.py

@@ -0,0 +1,28 @@
+"""
+Config-driven diffusion performance test with pytest parametrization.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.test.server.test_server_common import (  # noqa: F401
+    DiffusionServerBase,
+    diffusion_server,
+)
+from sglang.multimodal_gen.test.server.testcase_configs import (
+    ONE_GPU_CASES_C,
+    DiffusionTestCase,
+)
+
+logger = init_logger(__name__)
+
+
+class TestDiffusionServerOneGpuB200(DiffusionServerBase):
+    """B200-targeted smoke tests for 1-GPU diffusion cases."""
+
+    @pytest.fixture(params=ONE_GPU_CASES_C, ids=lambda c: c.id)
+    def case(self, request) -> DiffusionTestCase:
+        """Provide a DiffusionTestCase for each 1-GPU B200 test."""
+        return request.param

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -798,7 +798,7 @@ if not current_platform.is_hip():
         )
     )
 
-# TODO: enable on 4090/5090/b200
+# TODO: enable on 4090/5090
 ONE_GPU_CASES_C = [
     DiffusionTestCase(
         "flux_2_nvfp4_t2i",

scripts/ci/utils/slash_command_handler.py

@@ -273,6 +273,7 @@ def handle_rerun_stage(
         "multimodal-gen-test-2-gpu",
         "multimodal-gen-component-accuracy-1-gpu",
         "multimodal-gen-component-accuracy-2-gpu",
+        "multimodal-gen-test-1-b200",
     ]
 
     # Valid AMD stage names that support target_stage