From ead9d7aa43c6fdfdda46039820cd2ed5e1e0507b Mon Sep 17 00:00:00 2001
From: R0CKSTAR <yeahdongcn@gmail.com>
Date: Wed, 18 Mar 2026 15:44:59 +0800
Subject: [PATCH] [diffusion] fix: fix vae model offload on mps(#20607)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 .../pipelines_core/composed_pipeline_base.py  |  6 +++-
 .../runtime/pipelines_core/stages/decoding.py | 28 ++++++++++++-------
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
index bdbcdcc14..68b5927c0 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
@@ -449,7 +449,11 @@ class ComposedPipelineBase(ABC):
     ) -> "ComposedPipelineBase":
 
         return self.add_stage(
-            DecodingStage(vae=self.get_module(vae_key), pipeline=self),
+            DecodingStage(
+                vae=self.get_module(vae_key),
+                pipeline=self,
+                component_name=vae_key,
+            ),
         )
 
     def add_standard_t2i_stages(
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py
index 39c7c9d41..36ee3c119 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py
+++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py
@@ -56,10 +56,11 @@ class DecodingStage(PipelineStage):
     output format (e.g., pixel values).
     """
 
-    def __init__(self, vae, pipeline=None) -> None:
+    def __init__(self, vae, pipeline=None, component_name: str = "vae") -> None:
         super().__init__()
         self.vae: ParallelTiledVAE = vae
         self.pipeline = weakref.ref(pipeline) if pipeline else None
+        self.component_name = component_name
 
     @property
     def parallelism_type(self) -> StageParallelismType:
@@ -156,14 +157,17 @@ class DecodingStage(PipelineStage):
     def load_model(self):
         # load vae if not already loaded (used for memory constrained devices)
         pipeline = self.pipeline() if self.pipeline else None
-        if not self.server_args.model_loaded["vae"]:
+        if not self.server_args.model_loaded[self.component_name]:
             loader = VAELoader()
-            self.vae = loader.load(
-                self.server_args.model_paths["vae"], self.server_args
+            self.vae, _ = loader.load(
+                self.server_args.model_paths[self.component_name],
+                self.server_args,
+                component_name=self.component_name,
+                transformers_or_diffusers=loader.expected_library,
             )
             if pipeline:
-                pipeline.add_module("vae", self.vae)
-            self.server_args.model_loaded["vae"] = True
+                pipeline.add_module(self.component_name, self.vae)
+            self.server_args.model_loaded[self.component_name] = True
 
     def offload_model(self):
         # Offload models if needed
@@ -173,11 +177,13 @@ class DecodingStage(PipelineStage):
             self.vae.to("cpu", non_blocking=True)
 
         if torch.backends.mps.is_available():
+            # Flush lazy MPS kernels before freeing weights to avoid hangs.
+            torch.mps.synchronize()
             del self.vae
             pipeline = self.pipeline() if self.pipeline else None
-            if pipeline is not None and "vae" in pipeline.modules:
-                del pipeline.modules["vae"]
-            self.server_args.model_loaded["vae"] = False
+            if pipeline is not None and self.component_name in pipeline.modules:
+                del pipeline.modules[self.component_name]
+            self.server_args.model_loaded[self.component_name] = False
 
     @torch.no_grad()
     def forward(
@@ -234,6 +240,8 @@ class DecodingStage(PipelineStage):
             metrics=batch.metrics,
         )
 
-        self.offload_model()
+        # Keep VAE resident during warmup; the real request needs it next.
+        if not getattr(batch, "is_warmup", False):
+            self.offload_model()
 
         return output_batch