From acb64db9e27e8c6969fb1c69748e514d1375c68c Mon Sep 17 00:00:00 2001 From: Yihao Wang <42559837+AgainstEntropy@users.noreply.github.com> Date: Fri, 7 Aug 2026 04:01:05 -0700 Subject: [PATCH] [diffusion] fix: enable bcg with tp (#33421) Co-authored-by: Claude Opus 5 (1M context) --- .../runtime/breakable_cuda_graph/runner.py | 31 ++- .../runtime/distributed/group_coordinator.py | 13 +- .../multimodal_gen/test/server/gpu_cases.py | 4 + .../test_diffusion_bcg_tp2_zimage_turbo.py | 160 ++++++++++++++ .../test_diffusion_bcg_tp_graph_capture.py | 207 ++++++++++++++++++ 5 files changed, 408 insertions(+), 7 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/single_test_file/test_diffusion_bcg_tp2_zimage_turbo.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py diff --git a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py index 46f8af05b..44db42df4 100644 --- a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py b/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py index 4045e4a67..bc564f82e 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/group_coordinator.py +++ b/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 diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index eba4aa175..0cc8af740 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/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 diff --git a/python/sglang/multimodal_gen/test/single_test_file/test_diffusion_bcg_tp2_zimage_turbo.py b/python/sglang/multimodal_gen/test/single_test_file/test_diffusion_bcg_tp2_zimage_turbo.py new file mode 100644 index 000000000..36dc3d153 --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py new file mode 100644 index 000000000..7cf37e570 --- /dev/null +++ b/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)