minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] fix: fix vae model offload on mps(#20607)

Browse Source 
Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
R0CKSTAR
2026-03-18 15:44:59 +08:00
committed by GitHub
co-authored by gemini-code-assist[bot]
parent 532470bcca
commit ead9d7aa43
2 changed files with 23 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
+5 -1
View File
@@ -449,7 +449,11 @@ class ComposedPipelineBase(ABC):
) -> "ComposedPipelineBase": ) -> "ComposedPipelineBase":
return self.add_stage( 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( def add_standard_t2i_stages(
python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding.py
+18 -10
View File
@@ -56,10 +56,11 @@ class DecodingStage(PipelineStage):
output format (e.g., pixel values). 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__() super().__init__()
self.vae: ParallelTiledVAE = vae self.vae: ParallelTiledVAE = vae
self.pipeline = weakref.ref(pipeline) if pipeline else None self.pipeline = weakref.ref(pipeline) if pipeline else None
self.component_name = component_name
@property @property
def parallelism_type(self) -> StageParallelismType: def parallelism_type(self) -> StageParallelismType:
@@ -156,14 +157,17 @@ class DecodingStage(PipelineStage):
def load_model(self): def load_model(self):
# load vae if not already loaded (used for memory constrained devices) # load vae if not already loaded (used for memory constrained devices)
pipeline = self.pipeline() if self.pipeline else None 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() loader = VAELoader()
self.vae = loader.load( self.vae, _ = loader.load(
self.server_args.model_paths["vae"], self.server_args self.server_args.model_paths[self.component_name],
self.server_args,
component_name=self.component_name,
transformers_or_diffusers=loader.expected_library,
) )
if pipeline: if pipeline:
pipeline.add_module("vae", self.vae) pipeline.add_module(self.component_name, self.vae)
self.server_args.model_loaded["vae"] = True self.server_args.model_loaded[self.component_name] = True
def offload_model(self): def offload_model(self):
# Offload models if needed # Offload models if needed
@@ -173,11 +177,13 @@ class DecodingStage(PipelineStage):
self.vae.to("cpu", non_blocking=True) self.vae.to("cpu", non_blocking=True)
if torch.backends.mps.is_available(): if torch.backends.mps.is_available():
# Flush lazy MPS kernels before freeing weights to avoid hangs.
torch.mps.synchronize()
del self.vae del self.vae
pipeline = self.pipeline() if self.pipeline else None pipeline = self.pipeline() if self.pipeline else None
if pipeline is not None and "vae" in pipeline.modules: if pipeline is not None and self.component_name in pipeline.modules:
del pipeline.modules["vae"] del pipeline.modules[self.component_name]
self.server_args.model_loaded["vae"] = False self.server_args.model_loaded[self.component_name] = False
@torch.no_grad() @torch.no_grad()
def forward( def forward(
@@ -234,6 +240,8 @@ class DecodingStage(PipelineStage):
metrics=batch.metrics, 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 return output_batch