Add diffusion BCG prompt conditioning guard (#30782)

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

@@ -980,7 +980,7 @@ STANDALONE_FILES = {
 # measured value that must be copied into STANDALONE_FILE_EST_TIMES.
 STANDALONE_FILE_EST_TIMES = {
     "bcg-diffusion": {
-        "../single_test_file/test_diffusion_bcg_zimage_turbo.py": 300.0,
+        "../single_test_file/test_diffusion_bcg_zimage_turbo.py": 420.0,
     },
     "1-gpu": {
         "../single_test_file/test_update_weights_from_disk.py": 1200.0,

python/sglang/multimodal_gen/test/single_test_file/test_diffusion_bcg_zimage_turbo.py

@@ -5,9 +5,20 @@ import tempfile
 import unittest
 from pathlib import Path
 
-from sglang.multimodal_gen.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+from openai import OpenAI
+
+from sglang.multimodal_gen.test.server.test_server_utils import ServerManager
+from sglang.multimodal_gen.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+    get_dynamic_server_port,
+)
 from sglang.test.test_utils import CustomTestCase
 
+IMAGE_SIZE = "512x512"
+PROMPT_SWITCH_NUM_INFERENCE_STEPS = 4
+PROMPT_SWITCH_SEED = 0
+BCG_CAPTURE_MARKER = "[Diffusion BCG] captured"
+
 
 class TestDiffusionBCGZImageTurbo(CustomTestCase):
     def test_zimage_turbo_true_bcg_generate(self):
@@ -94,6 +105,68 @@ class TestDiffusionBCGZImageTurbo(CustomTestCase):
         self.assertIn("DenoisingStage", stage_names)
         self.assertGreater(len(perf.get("denoise_steps_ms", [])), 0)
 
+    def test_zimage_turbo_bcg_prompt_switch_reuses_captured_graph(self):
+        port = get_dynamic_server_port()
+        extra_args = (
+            "--model-type diffusion "
+            "--num-gpus 1 "
+            "--strict-ports "
+            "--enable-breakable-cuda-graph "
+            f"--warmup-resolutions {IMAGE_SIZE} "
+            "--bcg-text-buckets 256"
+        )
+        manager = ServerManager(
+            model=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            port=port,
+            wait_deadline=float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")),
+            extra_args=extra_args,
+        )
+        ctx = manager.start()
+        try:
+            client = OpenAI(
+                api_key="sglang-anything",
+                base_url=f"http://localhost:{ctx.port}/v1",
+                timeout=float(
+                    os.environ.get("SGLANG_TEST_OPENAI_REQUEST_TIMEOUT_SECS", "600")
+                ),
+                max_retries=0,
+            )
+            self._generate_prompt_switch_image(
+                client,
+                "a red cube on a white table, centered product photo",
+            )
+            self._generate_prompt_switch_image(
+                client,
+                "a blue butterfly flying above green grass, watercolor illustration",
+            )
+        finally:
+            server_log = ctx.stdout_file.read_text(encoding="utf-8", errors="ignore")
+            ctx.cleanup()
+
+        capture_count = server_log.count(BCG_CAPTURE_MARKER)
+        self.assertEqual(
+            capture_count,
+            1,
+            "BCG should be captured during warmup and then reused by different "
+            f"prompt requests. Observed {capture_count} capture log line(s).\n"
+            f"Server log tail:\n{ctx.log_tail(lines=400)}",
+        )
+
+    def _generate_prompt_switch_image(self, client: OpenAI, prompt: str):
+        response = client.images.with_raw_response.generate(
+            model=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            prompt=prompt,
+            n=1,
+            size=IMAGE_SIZE,
+            response_format="b64_json",
+            extra_body={
+                "num_inference_steps": PROMPT_SWITCH_NUM_INFERENCE_STEPS,
+                "seed": PROMPT_SWITCH_SEED,
+            },
+        )
+        parsed = response.parse()
+        self.assertTrue(parsed.data[0].b64_json)
+
 
 if __name__ == "__main__":
     unittest.main()