diff --git a/.github/workflows/_pr-test-check-changes.yml b/.github/workflows/_pr-test-check-changes.yml index dad1acfa8..9591fe97a 100644 --- a/.github/workflows/_pr-test-check-changes.yml +++ b/.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/**" diff --git a/.github/workflows/pr-test-multimodal-gen.yml b/.github/workflows/pr-test-multimodal-gen.yml index cc280aa61..c2c8957c9 100644 --- a/.github/workflows/pr-test-multimodal-gen.yml +++ b/.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: | diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 2bafefc65..104eb3bfa 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/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, diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index a4b98d422..754fafc39 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/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", diff --git a/python/sglang/multimodal_gen/test/single_test_file/test_diffusion_bcg_zimage_turbo.py b/python/sglang/multimodal_gen/test/single_test_file/test_diffusion_bcg_zimage_turbo.py new file mode 100644 index 000000000..e035359a7 --- /dev/null +++ b/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() diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py index 5c52deb57..f4849dd24 100644 --- a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_padding.py +++ b/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]]) diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py index 75344e6ab..df73ed97c 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py +++ b/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 - mempool reclaim per-layer intermediates between segments — storage stays - alive for each segment CUDAGraph's lifetime via its pool use_count. + recurse into tuples/lists; pass-through for non-tensors. Weak-ref'ing + captured args lets the shared mempool reclaim per-layer intermediates + 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, - and dicts of tensors. Returns dst if in-place copy succeeded, - otherwise returns src. + Handles plain tensors, tuples/lists of tensors, dataclass/object with + tensor attributes, and dicts of tensors. Returns dst if in-place copy + 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 - # CUDAGraphs' mempool pin; Python refs don't need to prevent - # pool reuse across layers. + # Weak-ref captured inputs produced by graph segments. Their storage + # is pinned by the segment CUDAGraphs' mempool use-count, so Python + # 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) diff --git a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py index 6ec4e1325..b32cbf1ff 100644 --- a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py +++ b/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."""