Fix diffusion BCG lifetime and add Z-Image-Turbo CI (#30584)

.github/workflows/_pr-test-check-changes.yml

@@ -91,6 +91,7 @@ jobs:
               - ".github/workflows/pr-test-multimodal-gen.yml"
               - "python/pyproject.toml"
               - "python/sglang/multimodal_gen/**/!(*.md|*.ipynb)"
+              - "python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/**"
               - "python/sglang/srt/observability/**"
               - "python/sglang/jit_kernel/**"
               - "test/registered/jit/diffusion/**"

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

@@ -207,6 +207,61 @@ jobs:
         with:
           artifact-suffix: 1-gpu-5090
 
+  bcg-diffusion:
+    if: |
+      ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) &&
+      inputs.multimodal_gen == 'true'
+    runs-on: 1-gpu-h100
+    timeout-minutes: 90
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-pr-test-health
+
+      - 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-cuda*
+
+      - 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 BCG diffusion tests
+        timeout-minutes: 60
+        env:
+          RUNAI_STREAMER_MEMORY_LIMIT: 0
+          CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
+          SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-bcg-artifacts
+        run: |
+          cd python
+          python3 sglang/multimodal_gen/test/run_suite.py \
+            --suite bcg-diffusion \
+            $CONTINUE_ON_ERROR_FLAG
+
+      - name: Upload BCG diffusion artifacts
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: diffusion-bcg-artifacts-${{ github.run_attempt }}
+          path: diffusion-bcg-artifacts/
+          if-no-files-found: ignore
+          retention-days: 7
+
+      - uses: ./.github/actions/upload-cuda-coredumps
+        if: failure()
+        with:
+          artifact-suffix: bcg-diffusion
+
   multimodal-gen-test-2-gpu:
     needs: compute-diffusion-partitions
     if: |

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -74,6 +74,13 @@ from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import
 
 logger = init_logger(__name__)  # pylint: disable=invalid-name
 
+
+def _attn_mask_meta_local_pad(attn_mask_meta) -> int:
+    if attn_mask_meta is None or isinstance(attn_mask_meta, DynamicVarlenMaskMeta):
+        return 0
+    return attn_mask_meta.get("local_pad", 0)
+
+
 try:
     from nunchaku.models.attention import NunchakuFeedForward  # type: ignore[import]
 except Exception:
@@ -694,7 +701,7 @@ class QwenImageCrossAttention(nn.Module):
         # fully sequence-parallel, so no leading tokens are replicated.
         sp_text_sharded = cross_attention_kwargs.get("sp_text_sharded", False)
         # Rows of tail padding inside THIS rank's text chunk (sp_shard meta).
-        sp_txt_pad = (attn_mask_meta or {}).get("local_pad", 0)
+        sp_txt_pad = _attn_mask_meta_local_pad(attn_mask_meta)
 
         (
             img_query,

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

@@ -947,9 +947,13 @@ PARAMETRIZED_CASE_GROUPS = {
     "2-gpu": [
         ("test_server_2_gpu.py", TWO_GPU_CASES),
     ],
+    "bcg-diffusion": [],
 }
 
 STANDALONE_FILES = {
+    "bcg-diffusion": [
+        "../single_test_file/test_diffusion_bcg_zimage_turbo.py",
+    ],
     "1-gpu": [
         "../single_test_file/test_generate_zimage_turbo_cli.py",
         "../single_test_file/test_update_weights_from_disk.py",
@@ -964,6 +968,9 @@ STANDALONE_FILES = {
 # CI will use a fallback estimate for sharding, run the test, then print a
 # 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,
+    },
     "1-gpu": {
         "../single_test_file/test_update_weights_from_disk.py": 1200.0,
     },
@@ -985,7 +992,7 @@ SUITES = {
     },
 }
 
-STRICT_SUITES = {"unit"}
+STRICT_SUITES = {"unit", "bcg-diffusion"}
 COMPONENT_ACCURACY_SUITES = {
     "component-accuracy",
     "component-accuracy-1-gpu",

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

@@ -0,0 +1,99 @@
+import json
+import os
+import subprocess
+import tempfile
+import unittest
+from pathlib import Path
+
+from sglang.multimodal_gen.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+from sglang.test.test_utils import CustomTestCase
+
+
+class TestDiffusionBCGZImageTurbo(CustomTestCase):
+    def test_zimage_turbo_true_bcg_generate(self):
+        artifact_dir = Path(
+            os.environ.get(
+                "SGLANG_DIFFUSION_ARTIFACT_DIR",
+                tempfile.mkdtemp(prefix="sglang_diffusion_bcg_"),
+            )
+        )
+        artifact_dir.mkdir(parents=True, exist_ok=True)
+        log_path = artifact_dir / "zimage_turbo_bcg.log"
+        perf_path = artifact_dir / "zimage_turbo_bcg_perf.json"
+
+        cmd = [
+            "sglang",
+            "generate",
+            "--backend",
+            "sglang",
+            "--model-path",
+            DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            "--prompt",
+            (
+                "A detailed cinematic scene of a glass observatory above a quiet "
+                "lake at sunrise, with soft mist, warm reflections, and crisp "
+                "architectural detail"
+            ),
+            "--width",
+            "512",
+            "--height",
+            "512",
+            "--seed",
+            "42",
+            "--num-inference-steps",
+            "9",
+            "--warmup-resolutions",
+            "512x512",
+            "--no-save-output",
+            "--guidance-scale",
+            "0.0",
+            "--enable-breakable-cuda-graph",
+            "--bcg-text-buckets",
+            "128",
+            "--enable-torch-compile",
+            "false",
+            "--dit-layerwise-offload",
+            "false",
+            "--dit-cpu-offload",
+            "false",
+            "--perf-dump-path",
+            str(perf_path),
+        ]
+
+        env = os.environ.copy()
+        env.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
+
+        result = subprocess.run(
+            cmd,
+            env=env,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            timeout=300,
+        )
+        log_path.write_text(result.stdout, encoding="utf-8")
+
+        self.assertEqual(
+            result.returncode,
+            0,
+            f"Z-Image-Turbo BCG command failed. Log: {log_path}\n"
+            f"{result.stdout[-4000:]}",
+        )
+        self.assertNotIn("Falling back to diffusers backend", result.stdout)
+        self.assertNotIn("Using diffusers backend", result.stdout)
+        self.assertNotIn("Loaded diffusers pipeline", result.stdout)
+        self.assertNotIn("[Diffusion BCG] capture failed", result.stdout)
+        self.assertIn("[Diffusion BCG] captured", result.stdout)
+        self.assertIn("Pixel data generated successfully", result.stdout)
+
+        self.assertTrue(perf_path.exists(), f"perf dump not found: {perf_path}")
+        perf = json.loads(perf_path.read_text(encoding="utf-8"))
+        stage_names = {
+            step.get("name") for step in perf.get("steps", []) if isinstance(step, dict)
+        }
+        self.assertIn("DenoisingStage", stage_names)
+        self.assertGreater(len(perf.get("denoise_steps_ms", [])), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py

@@ -13,6 +13,9 @@ from sglang.multimodal_gen.runtime.layers.attention import (
     DynamicVarlenMaskMeta,
     build_varlen_mask_meta,
 )
+from sglang.multimodal_gen.runtime.models.dits.qwen_image import (
+    _attn_mask_meta_local_pad,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import (
     DenoisingStage,
 )
@@ -109,6 +112,11 @@ class TestDiffusionBCGPadding(unittest.TestCase):
         )
         self.assertEqual(_signature_kwargs(first), _signature_kwargs(second))
 
+    def test_qwen_dynamic_varlen_meta_is_not_tail_pad_meta(self):
+        self.assertEqual(_attn_mask_meta_local_pad(None), 0)
+        self.assertEqual(_attn_mask_meta_local_pad({"local_pad": 7}), 7)
+        self.assertEqual(_attn_mask_meta_local_pad(DynamicVarlenMaskMeta()), 0)
+
     def test_qwen_default_bucket_preserves_mask(self):
         def kwargs(valid_len: int):
             mask = torch.zeros(1, 64, dtype=torch.bool)
@@ -248,6 +256,23 @@ class TestDiffusionBCGPadding(unittest.TestCase):
             BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS,
         )
 
+    def test_image_generation_models_are_registered_as_bcg_supported(self):
+        for model_id in (
+            "qwen/qwen-image",
+            "qwen/qwen-image-2512",
+            "tongyi-mai/z-image",
+            "tongyi-mai/z-image-turbo",
+            "zai-org/glm-image",
+        ):
+            self.assertIn(model_id, BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS)
+
+        for config_name in (
+            "GlmImagePipelineConfig",
+            "QwenImagePipelineConfig",
+            "ZImagePipelineConfig",
+        ):
+            self.assertIn(config_name, BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS)
+
     def test_dynamic_varlen_mask_meta_rebuilds_once_per_replay_token(self):
         builder = DynamicVarlenMaskMeta()
         mask = torch.tensor([[True, True, False, False]])

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py

@@ -155,9 +155,10 @@ def _uninstall_wait_stream_hook():
 
 def _weak_ref_if_tensor(x):
     """Return a weak-ref tensor view (shared storage, no refcount) for tensors;
-    pass-through for non-tensors. Weak-ref'ing captured args lets the shared
+    recurse into tuples/lists; pass-through for non-tensors. Weak-ref'ing
-    mempool reclaim per-layer intermediates between segments — storage stays
+    captured args lets the shared mempool reclaim per-layer intermediates
-    alive for each segment CUDAGraph's lifetime via its pool use_count.
+    between segments — storage stays alive for each segment CUDAGraph's
+    lifetime via its pool use_count.
 
     weak_ref_tensors is imported lazily because it hard-raises on
     platforms without a CUDA/HIP/NPU backend; we only reach this code during
@@ -166,20 +167,32 @@ def _weak_ref_if_tensor(x):
         from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
 
         return weak_ref_tensors(x)
+    if isinstance(x, tuple):
+        return tuple(_weak_ref_if_tensor(e) for e in x)
+    if isinstance(x, list):
+        return [_weak_ref_if_tensor(e) for e in x]
     return x
 
 
 def _copy_output(dst: Any, src: Any) -> Any:
     """Copy src output into dst in-place where possible.
 
-    Handles plain tensors, dataclass/object with tensor attributes,
+    Handles plain tensors, tuples/lists of tensors, dataclass/object with
-    and dicts of tensors. Returns dst if in-place copy succeeded,
+    tensor attributes, and dicts of tensors. Returns dst if in-place copy
-    otherwise returns src.
+    succeeded, otherwise returns src.
     """
     if torch.is_tensor(dst) and torch.is_tensor(src):
         dst.copy_(src)
         return dst
 
+    if (
+        isinstance(dst, (tuple, list))
+        and isinstance(src, (tuple, list))
+        and len(dst) == len(src)
+    ):
+        copied = [_copy_output(d, s) for d, s in zip(dst, src)]
+        return tuple(copied) if isinstance(dst, tuple) else copied
+
     if hasattr(dst, "__dict__") and hasattr(src, "__dict__"):
         for key, src_val in src.__dict__.items():
             dst_val = getattr(dst, key, None)
@@ -220,13 +233,17 @@ def eager_on_graph(enable: bool):
             # writes real data into them.
             output = inner(*args, **kwargs)
 
-            # Weak-ref the closure state. Storage lives with the segment
+            # Weak-ref captured inputs produced by graph segments. Their storage
-            # CUDAGraphs' mempool pin; Python refs don't need to prevent
+            # is pinned by the segment CUDAGraphs' mempool use-count, so Python
-            # pool reuse across layers.
+            # refs do not need to keep every intermediate alive.
             captured_inner = inner
             captured_args = tuple(_weak_ref_if_tensor(a) for a in args)
             captured_kwargs = {k: _weak_ref_if_tensor(v) for k, v in kwargs.items()}
-            captured_output = _weak_ref_if_tensor(output)
+            # The eager break output is different: it is allocated between graph
+            # captures and is the static input address consumed by the next
+            # captured segment. Keep a strong reference so replay can safely
+            # copy fresh eager output into that bridge buffer.
+            captured_output = output
 
             def replay_fn():
                 new_out = captured_inner(*captured_args, **captured_kwargs)

test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py

@@ -168,6 +168,28 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         torch.cuda.synchronize()
         self.assertTrue(torch.allclose(y, torch.full((4,), 33.0, device=self.device)))
 
+    def test_eager_output_is_held_strongly_for_replay_bridge(self):
+        """The replay closure must keep the eager output bridge buffer alive."""
+        x = torch.zeros(4, device=self.device)
+        y = torch.zeros(4, device=self.device)
+
+        @self.eager_on_graph(enable=True)
+        def scale(src):
+            return src * 3.0
+
+        graph = self.BreakableCUDAGraph()
+        stream = torch.cuda.Stream(self.device)
+        with self.BreakableCUDAGraphCapture(graph, stream=stream):
+            t = x + 1.0
+            broken = scale(t)
+            y.copy_(broken)
+
+        replay_closure = graph._break_fns[0].__closure__ or ()
+        self.assertTrue(
+            any(cell.cell_contents is broken for cell in replay_closure),
+            "eager output bridge buffer must be strongly captured",
+        )
+
 
 class TestCopyOutput(CustomTestCase):
     """Test the _copy_output helper for structured output writeback."""