[diffusion] fix: enable bcg with tp (#33421)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py

@@ -28,6 +28,7 @@ behavior.
 
 from __future__ import annotations
 
+import contextlib
 import logging
 import os
 from dataclasses import dataclass
@@ -341,6 +342,25 @@ class BaseBreakableCudaGraphRunner:
         """
         return _signature_kwargs(kwargs)
 
+    @contextlib.contextmanager
+    def _tp_graph_capture(self):
+        """Enter the tensor-parallel group's graph-capture context around capture."""
+
+        from sglang.multimodal_gen.runtime.distributed.group_coordinator import (
+            GraphCaptureContext,
+        )
+        from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+            get_tp_group,
+            model_parallel_is_initialized,
+        )
+
+        if not model_parallel_is_initialized() or get_tp_group().world_size <= 1:
+            yield
+            return
+
+        with get_tp_group().graph_capture(GraphCaptureContext(self._capture_stream)):
+            yield
+
     def _empty_cache(self) -> None:
         empty_cache = getattr(self.device_module, "empty_cache", None)
         if callable(empty_cache):
@@ -422,11 +442,12 @@ class BaseBreakableCudaGraphRunner:
         self.device_module.synchronize()
 
         graph = BreakableCUDAGraph()
-        with enable_breakable_cuda_graph():
-            with BreakableCUDAGraphCapture(
-                cuda_graph=graph, pool=self._pool, stream=self._capture_stream
-            ):
-                output = self.transformer(**static_kwargs)
+        with self._tp_graph_capture():
+            with enable_breakable_cuda_graph():
+                with BreakableCUDAGraphCapture(
+                    cuda_graph=graph, pool=self._pool, stream=self._capture_stream
+                ):
+                    output = self.transformer(**static_kwargs)
         self.device_module.synchronize()
 
         logger.info(

python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py

@@ -9,7 +9,7 @@
 # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
 import pickle
 from collections import namedtuple
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -314,7 +314,16 @@ class GroupCoordinator:
             if curr_stream != stream:
                 stream.wait_stream(curr_stream)
 
-            with torch.cuda.stream(stream):
+            # Custom all-reduce has an eager path and a graph path; only inside
+            # capture() does all_reduce pick the graph one, and only capture()
+            # exit registers the graph-pool buffers with the peers. Capturing the
+            # eager path instead faults on replay (unmapped peer IPC address).
+            custom_ar = self.srt_custom_allreduce
+            maybe_ca_context = (
+                nullcontext() if custom_ar is None else custom_ar.capture()
+            )
+
+            with torch.cuda.stream(stream), maybe_ca_context:
                 yield graph_capture_context
         else:
             # For non-CUDA platforms (MPS, CPU), just yield the context without stream management

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

@@ -1157,6 +1157,7 @@ STANDALONE_FILES = {
         "../single_test_file/test_disagg_server.py",
         "../single_test_file/test_ar_models.py",
         "../single_test_file/test_ipc_a2a_2_gpu.py",
+        "../single_test_file/test_diffusion_bcg_tp2_zimage_turbo.py",
         "../single_test_file/test_dp_serving_2_gpu.py",
         "../single_test_file/test_pynccl_a2a_capture_2_gpu.py",
         "../single_test_file/test_usp_replicated_parity_2_gpu.py",
@@ -1194,6 +1195,9 @@ STANDALONE_FILE_EST_TIMES = {
         "../single_test_file/test_ar_models.py": 600.0,
         # no model load; the cost is the one-time JIT build of the sync kernels
         "../single_test_file/test_ipc_a2a_2_gpu.py": 240.0,
+        # ~60 s locally with a warm HF cache (load + one capture + 4 steps);
+        # padded for cold-cache CI.
+        "../single_test_file/test_diffusion_bcg_tp2_zimage_turbo.py": 180.0,
         # zimage server startup dominates; six short requests after warmup
         "../single_test_file/test_dp_serving_2_gpu.py": 900.0,
         # one capture plus three replays on a 32K-element exchange

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

@@ -0,0 +1,160 @@
+"""Breakable CUDA graph must survive tensor parallelism.
+
+With ``--tp-size > 1`` every ``RowParallelLinear`` issues a per-block all-reduce
+that lands *inside* a captured BCG segment (BCG's break points are the attention
+modules, so everything else is captured). Custom all-reduce hands the captured
+kernel a rank-data slot whose peer pointers are only filled in by the
+``register_graph_buffers()`` that runs when ``CustomAllreduce.capture()`` exits,
+and that registration is a host-side IPC exchange -- it cannot happen inside the
+captured region. If capture is not wrapped in the TP group's ``graph_capture()``,
+the slots stay unwritten and replay dereferences them:
+``cudaErrorIllegalAddress``, no image.
+
+The failure is invisible to the single-GPU BCG suite (``--num-gpus 1`` never
+takes the custom all-reduce path at all), which is why this file exists. Capture
+itself succeeds either way -- the log still prints "captured N segment(s)" -- so
+asserting on the capture marker alone is not enough; we also require the
+registration to have happened.
+
+    pytest -v python/sglang/multimodal_gen/test/single_test_file/test_diffusion_bcg_tp2_zimage_turbo.py
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import tempfile
+import unittest
+from pathlib import Path
+
+import torch
+
+from sglang.multimodal_gen.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
+from sglang.test.test_utils import CustomTestCase
+
+IMAGE_SIZE = "512x512"
+NUM_INFERENCE_STEPS = 4
+SEED = 0
+TP_SIZE = 2
+
+BCG_CAPTURE_MARKER = "[Diffusion BCG] captured"
+# Printed by CustomAllreduce.register_graph_buffers(), i.e. only on exit of the
+# capture() context the fix enters. Its absence is the precise signature of the
+# regression: capture ran, registration did not.
+GRAPH_BUFFER_REGISTRATION_MARKER = "cuda graph addresses"
+ILLEGAL_MEMORY_MARKER = "illegal memory access"
+
+# The regression fails during warmup within ~1 min; the timeout only has to be
+# generous enough for weight load plus one capture per bucket.
+GENERATE_TIMEOUT_SECONDS = 900
+
+
+class TestDiffusionBCGTP2ZImageTurbo(CustomTestCase):
+    def test_zimage_turbo_bcg_generates_under_tp2(self):
+        if torch.cuda.device_count() < TP_SIZE:
+            self.skipTest(f"needs {TP_SIZE} GPUs, found {torch.cuda.device_count()}")
+
+        artifact_dir = Path(
+            os.environ.get(
+                "SGLANG_DIFFUSION_ARTIFACT_DIR",
+                tempfile.mkdtemp(prefix="sglang_diffusion_bcg_tp2_"),
+            )
+        )
+        artifact_dir.mkdir(parents=True, exist_ok=True)
+        log_path = artifact_dir / "zimage_turbo_bcg_tp2.log"
+
+        cmd = [
+            "sglang",
+            "generate",
+            "--backend",
+            "sglang",
+            "--model-path",
+            DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
+            "--prompt",
+            "A red fox in fresh snow",
+            "--width",
+            IMAGE_SIZE.split("x")[0],
+            "--height",
+            IMAGE_SIZE.split("x")[1],
+            "--seed",
+            str(SEED),
+            "--num-inference-steps",
+            str(NUM_INFERENCE_STEPS),
+            "--num-gpus",
+            str(TP_SIZE),
+            "--tp-size",
+            str(TP_SIZE),
+            "--warmup-resolutions",
+            IMAGE_SIZE,
+            "--no-save-output",
+            "--guidance-scale",
+            "0.0",
+            "--enable-breakable-cuda-graph",
+            # One bucket keeps the run short; the TP all-reduce path does not
+            # depend on how many prompt buckets get captured.
+            "--bcg-text-buckets",
+            "128",
+            "--enable-torch-compile",
+            "false",
+            "--dit-layerwise-offload",
+            "false",
+            "--dit-cpu-offload",
+            "false",
+        ]
+
+        env = os.environ.copy()
+        env.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
+
+        try:
+            result = subprocess.run(
+                cmd,
+                env=env,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+                text=True,
+                timeout=GENERATE_TIMEOUT_SECONDS,
+            )
+        except subprocess.TimeoutExpired as exc:
+            # An unregistered graph buffer can also hang the worker instead of
+            # returning, so a timeout is a failure of this guard, not flakiness.
+            output = exc.output or b""
+            if isinstance(output, bytes):
+                output = output.decode("utf-8", errors="replace")
+            log_path.write_text(output, encoding="utf-8")
+            self.fail(
+                f"TP={TP_SIZE} BCG generate hung after "
+                f"{GENERATE_TIMEOUT_SECONDS}s. Log: {log_path}\n{output[-4000:]}"
+            )
+
+        log_path.write_text(result.stdout, encoding="utf-8")
+        tail = result.stdout[-4000:]
+
+        self.assertEqual(
+            result.returncode,
+            0,
+            f"TP={TP_SIZE} BCG generate failed. Log: {log_path}\n{tail}",
+        )
+        self.assertNotIn(
+            ILLEGAL_MEMORY_MARKER,
+            result.stdout,
+            f"replay hit an unregistered peer address. Log: {log_path}\n{tail}",
+        )
+        self.assertNotIn("[Diffusion BCG] capture failed", result.stdout)
+        self.assertIn(
+            BCG_CAPTURE_MARKER,
+            result.stdout,
+            f"BCG never captured, so this run did not exercise TP+BCG. "
+            f"Log: {log_path}\n{tail}",
+        )
+        self.assertIn(
+            GRAPH_BUFFER_REGISTRATION_MARKER,
+            result.stdout,
+            f"custom all-reduce graph buffers were never registered, so capture "
+            f"did not go through the TP group's graph_capture(). "
+            f"Log: {log_path}\n{tail}",
+        )
+        self.assertIn("Pixel data generated successfully", result.stdout)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=3)

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

@@ -0,0 +1,207 @@
+"""BCG capture must run inside the TP group's graph-capture context.
+
+Custom all-reduce records every capture-time input into a rank-data slot whose
+peer pointers are written only by the ``register_graph_buffers()`` that runs when
+``CustomAllreduce.capture()`` exits. That registration is a host-side IPC
+exchange, so it cannot happen inside the captured region -- it has to be driven
+by a context that wraps the whole capture. Two things have to hold:
+
+1. ``GroupCoordinator.graph_capture()`` enters ``srt_custom_allreduce.capture()``
+   and leaves it only after the captured body is done.
+2. ``BaseBreakableCudaGraphRunner._capture()`` wraps its capture in the TP
+   group's ``graph_capture()``.
+
+Break either and BCG under ``--tp-size > 1`` faults on replay with
+``cudaErrorIllegalAddress``. The end-to-end guard for that is
+``single_test_file/test_diffusion_bcg_tp2_zimage_turbo.py``, which needs 2 GPUs;
+these tests guard the same wiring on every commit without one.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from sglang.multimodal_gen.runtime.breakable_cuda_graph import runner as runner_mod
+from sglang.multimodal_gen.runtime.breakable_cuda_graph.runner import (
+    BaseBreakableCudaGraphRunner,
+)
+from sglang.multimodal_gen.runtime.distributed.group_coordinator import (
+    GraphCaptureContext,
+    GroupCoordinator,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+def _recording_context(events: list, name: str):
+    """A context manager that appends ``name`` enter/exit to ``events``."""
+
+    @contextlib.contextmanager
+    def _ctx(*args, **kwargs):
+        events.append(f"enter:{name}")
+        try:
+            yield MagicMock()
+        finally:
+            events.append(f"exit:{name}")
+
+    return _ctx
+
+
+class TestBCGTPGraphCapture(CustomTestCase):
+    # --- GroupCoordinator.graph_capture -> CustomAllreduce.capture ---------- #
+
+    def _run_graph_capture(self, custom_ar, events):
+        """Drive the CUDA branch of graph_capture() with a fake custom AR."""
+        group = SimpleNamespace(srt_custom_allreduce=custom_ar)
+        ctx = GraphCaptureContext(MagicMock())
+        with patch(
+            "sglang.multimodal_gen.runtime.distributed.group_coordinator.current_platform.is_cuda_alike",
+            return_value=True,
+        ), patch("torch.cuda.stream"), patch("torch.cuda.current_stream"):
+            with GroupCoordinator.graph_capture(group, ctx) as yielded:
+                events.append("body")
+        return yielded, ctx
+
+    def test_graph_capture_enters_custom_allreduce_capture(self):
+        events = []
+        custom_ar = SimpleNamespace(capture=_recording_context(events, "ca"))
+
+        yielded, ctx = self._run_graph_capture(custom_ar, events)
+
+        # The body must run *inside* capture(), so registration (which happens on
+        # its exit) lands after the captured region is closed.
+        self.assertEqual(events, ["enter:ca", "body", "exit:ca"])
+        self.assertIs(yielded, ctx)
+
+    def test_graph_capture_without_custom_allreduce_is_a_noop(self):
+        events = []
+
+        yielded, ctx = self._run_graph_capture(None, events)
+
+        self.assertEqual(events, ["body"])
+        self.assertIs(yielded, ctx)
+
+    # --- runner._tp_graph_capture -> GroupCoordinator.graph_capture -------- #
+
+    def _tp_graph_capture_events(self, *, world_size: int, initialized: bool = True):
+        events = []
+        capture_stream = MagicMock()
+        tp_group = MagicMock()
+        tp_group.world_size = world_size
+        tp_group.graph_capture = _recording_context(events, "tp")
+        runner = SimpleNamespace(_capture_stream=capture_stream)
+
+        with patch(
+            "sglang.multimodal_gen.runtime.distributed.parallel_state.model_parallel_is_initialized",
+            return_value=initialized,
+        ), patch(
+            "sglang.multimodal_gen.runtime.distributed.parallel_state.get_tp_group",
+            return_value=tp_group,
+        ):
+            with BaseBreakableCudaGraphRunner._tp_graph_capture(runner):
+                events.append("body")
+        return events, capture_stream
+
+    def test_tp_graph_capture_enters_tp_group_context(self):
+        events, _ = self._tp_graph_capture_events(world_size=2)
+
+        self.assertEqual(events, ["enter:tp", "body", "exit:tp"])
+
+    def test_tp_graph_capture_reuses_the_runners_capture_stream(self):
+        # Handing our own stream in keeps graph_capture() from creating a second
+        # one that nothing captures on.
+        events = []
+        capture_stream = MagicMock()
+        tp_group = MagicMock()
+        tp_group.world_size = 2
+        recorded = {}
+
+        @contextlib.contextmanager
+        def _graph_capture(graph_capture_context=None):
+            recorded["ctx"] = graph_capture_context
+            yield graph_capture_context
+
+        tp_group.graph_capture = _graph_capture
+        runner = SimpleNamespace(_capture_stream=capture_stream)
+
+        with patch(
+            "sglang.multimodal_gen.runtime.distributed.parallel_state.model_parallel_is_initialized",
+            return_value=True,
+        ), patch(
+            "sglang.multimodal_gen.runtime.distributed.parallel_state.get_tp_group",
+            return_value=tp_group,
+        ):
+            with BaseBreakableCudaGraphRunner._tp_graph_capture(runner):
+                pass
+
+        self.assertIsNotNone(recorded["ctx"])
+        self.assertIs(recorded["ctx"].stream, capture_stream)
+
+    def test_tp_graph_capture_is_noop_without_tensor_parallelism(self):
+        single_gpu, _ = self._tp_graph_capture_events(world_size=1)
+        self.assertEqual(single_gpu, ["body"])
+
+        uninitialized, _ = self._tp_graph_capture_events(
+            world_size=2, initialized=False
+        )
+        self.assertEqual(uninitialized, ["body"])
+
+    # --- runner._capture wraps the graph capture in the TP context --------- #
+
+    def test_capture_wraps_graph_capture_in_the_tp_context(self):
+        events = []
+        runner = object.__new__(BaseBreakableCudaGraphRunner)
+        runner.transformer = lambda **kwargs: torch.zeros(1)
+        runner.device = "cpu"
+        runner.device_module = MagicMock()
+        runner._pool = (0, 0)
+        runner._capture_stream = MagicMock()
+        runner.entries = {}
+        runner._blocked = set()
+        runner._disabled_reason = None
+        runner.max_entries = 0
+        runner.max_segments = 0
+
+        graph = MagicMock()
+        graph._segments = []
+        kwargs = {"hidden_states": torch.zeros(1)}
+
+        with patch.object(
+            BaseBreakableCudaGraphRunner,
+            "_tp_graph_capture",
+            _recording_context(events, "tp"),
+        ), patch.object(
+            runner_mod, "BreakableCUDAGraph", return_value=graph
+        ), patch.object(
+            runner_mod,
+            "enable_breakable_cuda_graph",
+            _recording_context(events, "bcg_enable"),
+        ), patch.object(
+            runner_mod,
+            "BreakableCUDAGraphCapture",
+            _recording_context(events, "bcg_capture"),
+        ):
+            runner._capture(kwargs, key=runner_mod._signature_kwargs(kwargs))
+
+        # The TP context must be the outermost one: registration on its exit has
+        # to happen after the captured region is closed, and the eager warmup
+        # forwards above it must not run with _IS_CAPTURING set.
+        self.assertEqual(
+            events,
+            [
+                "enter:tp",
+                "enter:bcg_enable",
+                "enter:bcg_capture",
+                "exit:bcg_capture",
+                "exit:bcg_enable",
+                "exit:tp",
+            ],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=3)