From 8e08207c18430fde10e43a593aa1924276934e72 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 30 Dec 2025 00:21:42 +0800 Subject: [PATCH] [diffusion] fix: fix serving with dit-layerwise-offload enabled (#16066) Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> --- .../benchmarks/bench_serving.py | 2 +- .../runtime/entrypoints/openai/image_api.py | 2 - .../pipelines_core/stages/denoising.py | 13 +++- .../runtime/utils/layerwise_offload.py | 7 +- .../test/server/perf_baselines.json | 68 +++++++++++++++++++ .../test/server/test_server_common.py | 3 + .../test/server/test_server_utils.py | 7 +- .../test/server/testcase_configs.py | 17 +++++ 8 files changed, 110 insertions(+), 9 deletions(-) diff --git a/python/sglang/multimodal_gen/benchmarks/bench_serving.py b/python/sglang/multimodal_gen/benchmarks/bench_serving.py index 65106183e..698c43353 100644 --- a/python/sglang/multimodal_gen/benchmarks/bench_serving.py +++ b/python/sglang/multimodal_gen/benchmarks/bench_serving.py @@ -750,7 +750,7 @@ async def benchmark(args): ) ) - print("\n" + "=" * 60) + print("=" * 60) if args.output_file: with open(args.output_file, "w") as f: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 94f23e7b6..074ba14c8 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py @@ -121,8 +121,6 @@ async def generations( server_args=get_global_server_args(), sampling_params=sampling, ) - - # Run synchronously for images and save to disk save_file_path, result = await process_generation_batch( async_scheduler_client, batch ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index d5862a258..4c4dcdc7e 100755 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -12,7 +12,7 @@ import time import weakref from collections.abc import Iterable from functools import lru_cache -from typing import Any +from typing import Any, Optional import torch from einops import rearrange @@ -61,6 +61,9 @@ from sglang.multimodal_gen.runtime.platforms import ( current_platform, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.layerwise_offload import ( + LayerwiseOffloadManager, +) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler @@ -721,6 +724,14 @@ class DenoisingStage(PipelineStage): torch.mps.current_allocated_memory(), ) + # reset offload manager with prefetching first layer for next forward + offload_mgr: Optional[LayerwiseOffloadManager] = None + for transformer in filter(None, [self.transformer, self.transformer_2]): + if ( + offload_mgr := getattr(transformer, "_layerwise_offload_manager", None) + ) is not None: + offload_mgr.prepare_for_next_denoise(non_blocking=True) + def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs): """Shard latents for Sequence Parallelism if applicable.""" if get_sp_world_size() <= 1: diff --git a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py index ffedc3a34..020746b6f 100644 --- a/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/utils/layerwise_offload.py @@ -98,8 +98,11 @@ class LayerwiseOffloadManager: continue self._offload_tensor(name, buf, layer_idx) - self.prefetch_layer(0, non_blocking=False) - if self.copy_stream is not None: + self.prepare_for_next_denoise(non_blocking=False) + + def prepare_for_next_denoise(self, non_blocking=True): + self.prefetch_layer(0, non_blocking=non_blocking) + if not non_blocking and self.copy_stream is not None: torch.cuda.current_stream().wait_stream(self.copy_stream) @torch.compiler.disable diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 45417a45f..34074d832 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -304,6 +304,74 @@ "expected_avg_denoise_ms": 520.09, "expected_median_denoise_ms": 528.0 }, + + "flux_2_image_t2i_layerwise_offload": { + "stages_ms": { + "InputValidationStage": 0.06, + "TextEncodingStage": 513.58, + "ImageVAEEncodingStage": 0.0, + "ConditioningStage": 0.03, + "LatentPreparationStage": 0.46, + "TimestepPreparationStage": 2.38, + "DenoisingStage": 52187.62, + "DecodingStage": 190.31 + }, + "denoise_step_ms": { + "0": 1033.45, + "1": 137.03, + "2": 1046.96, + "3": 1039.28, + "4": 1039.05, + "5": 1043.91, + "6": 1041.75, + "7": 1037.6, + "8": 1043.54, + "9": 1048.63, + "10": 1039.8, + "11": 1042.25, + "12": 1041.54, + "13": 1045.89, + "14": 1038.99, + "15": 1041.82, + "16": 1038.32, + "17": 1045.53, + "18": 1046.54, + "19": 1041.22, + "20": 1044.55, + "21": 1041.31, + "22": 1051.28, + "23": 1043.12, + "24": 1044.65, + "25": 1042.25, + "26": 1046.47, + "27": 1052.9, + "28": 1039.04, + "29": 1042.39, + "30": 1045.33, + "31": 1038.05, + "32": 1037.76, + "33": 1037.93, + "34": 1052.85, + "35": 1045.59, + "36": 1054.32, + "37": 1044.59, + "38": 1043.57, + "39": 1041.93, + "40": 1043.59, + "41": 1046.17, + "42": 1046.92, + "43": 1047.04, + "44": 1046.8, + "45": 1041.86, + "46": 1041.05, + "47": 1044.04, + "48": 1039.77, + "49": 1047.12 + }, + "expected_e2e_ms": 53290.15, + "expected_avg_denoise_ms": 1025.35, + "expected_median_denoise_ms": 1043.33 + }, "flux_2_ti2i": { "stages_ms": { "InputValidationStage": 99.82, diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index 055fa612a..75796486d 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -73,6 +73,9 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: if server_args.ulysses_degree is not None: extra_args += f" --ulysses-degree {server_args.ulysses_degree}" + if server_args.dit_layerwise_offload: + extra_args += f" --dit-layerwise-offload true" + if server_args.ring_degree is not None: extra_args += f" --ring-degree {server_args.ring_degree}" diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index 2cee4371c..46cdd7b7e 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -718,6 +718,7 @@ def get_generate_fn( """Return appropriate generation function for the case.""" # Allow override via environment variable (useful for AMD where large resolutions cause slow VAE) output_size = os.environ.get("SGLANG_TEST_OUTPUT_SIZE", sampling_params.output_size) + n = sampling_params.num_outputs_per_prompt def _create_and_download_video( client, @@ -837,7 +838,7 @@ def get_generate_fn( response = client.images.with_raw_response.generate( model=model_path, prompt=sampling_params.prompt, - n=1, + n=n, size=output_size, response_format="b64_json", ) @@ -906,7 +907,7 @@ def get_generate_fn( model=model_path, image=images, prompt=sampling_params.prompt, - n=1, + n=n, size=output_size, response_format="b64_json", ) @@ -974,7 +975,7 @@ def get_generate_fn( model=model_path, prompt=sampling_params.prompt, image=[], # Only for OpenAI verification - n=1, + n=n, size=sampling_params.output_size, response_format="b64_json", extra_body={"url": image_urls}, diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 0bb2b5a5a..060ee1527 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -152,6 +152,8 @@ class DiffusionServerArgs: # LoRA lora_path: str | None = None # LoRA adapter path (HF repo or local path) + dit_layerwise_offload: bool = False + @dataclass(frozen=True) class DiffusionSamplingParams: @@ -171,6 +173,8 @@ class DiffusionSamplingParams: # URL direct test flag - if True, don't pre-download URL images direct_url_test: bool = False + num_outputs_per_prompt: int = 1 + @dataclass(frozen=True) class DiffusionTestCase: @@ -308,6 +312,19 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [ ), T2I_sampling_params, ), + # TODO: replace with a faster model to test the --dit-layerwise-offload + # TODO: currently, we don't support sending more than one request in test, and setting `num_outputs_per_prompt` to 2 doesn't guarantee the denoising be executed twice, + # so we do one warmup and send one request instead + DiffusionTestCase( + "flux_2_image_t2i_layerwise_offload", + DiffusionServerArgs( + model_path="black-forest-labs/FLUX.2-dev", + modality="image", + dit_layerwise_offload=True, + warmup_text=1, + ), + T2I_sampling_params, + ), DiffusionTestCase( "zimage_image_t2i", DiffusionServerArgs(