From 575a49dc0e999f7f66e1aeefb67c448ae2a07e78 Mon Sep 17 00:00:00 2001 From: ryang <38470282+ryang-max@users.noreply.github.com> Date: Mon, 22 Dec 2025 17:01:35 +0800 Subject: [PATCH] [diffusion] chore: remove default post-denoising dit offload in local mode (#15573) --- .../runtime/pipelines_core/stages/denoising.py | 14 -------------- 1 file changed, 14 deletions(-) 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 211171c42..b2e1274ba 100755 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -754,20 +754,6 @@ class DenoisingStage(PipelineStage): torch.mps.current_allocated_memory(), ) - # In offline local mode (`sglang generate`), offload transformer weights to CPU - # after denoising to reduce peak VRAM during VAE decoding. - if current_platform.is_cuda_alike() and server_args.is_local_mode: - for model in (self.transformer, self.transformer_2): - if model is not None: - model.to("cpu") - logger.info( - "Offloaded denoiser transformer weights to CPU after denoising to reduce peak VRAM during VAE decoding." - ) - try: - torch.cuda.empty_cache() - except Exception: - pass - def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs): """Shard latents for Sequence Parallelism if applicable.""" if get_sp_world_size() <= 1: