chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[diffusion] improve: improve torch.compile for MOVA (#18914)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
triple-mu
2026-02-18 00:47:38 +08:00
committed by GitHub
parent bf52388354
commit 504b2c58cf
1 changed files with 13 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py
View File

@@ -17,6 +17,7 @@ import os
from collections.abc import Iterable
import torch
import torch.nn as nn
from diffusers.utils.torch_utils import randn_tensor
from tqdm.auto import tqdm
@@ -200,11 +201,13 @@ class MOVADenoisingStage(PipelineStage):
partial = (1 - guidance_scale) * neg
return cfg_model_parallel_all_reduce(partial)
def compile_module_with_torch_compile(self, module, server_args: ServerArgs):
if not server_args.enable_torch_compile or module is None:
return module
if not hasattr(module, "forward"):
return module
def _maybe_enable_torch_compile(self, module: nn.Module, server_args: ServerArgs):
"""
Compile a module with torch.compile, and enable inductor overlap tweak if available.
No-op if torch compile is disabled or the object is not a nn.Module.
"""
if not server_args.enable_torch_compile or not isinstance(module, nn.Module):
return
try:
import torch._inductor.config as _inductor_cfg
@@ -213,15 +216,14 @@ class MOVADenoisingStage(PipelineStage):
pass
mode = os.environ.get("SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs")
logger.info("Compiling %s with mode: %s", module.__class__.__name__, mode)
compiled_forward = torch.compile(getattr(module, "forward"), mode=mode)
setattr(module, "forward", compiled_forward)
return module
# TODO(triple-mu): support customized fullgraph and dynamic in the future
module.compile(mode=mode, fullgraph=False, dynamic=None)
def _maybe_compile_dits(self, server_args: ServerArgs):
if self._torch_compiled or not server_args.enable_torch_compile:
return
for module in filter(None, [self.video_dit, self.video_dit_2, self.audio_dit]):
self.compile_module_with_torch_compile(module, server_args)
self._maybe_enable_torch_compile(module, server_args)
self._torch_compiled = True
def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
@@ -310,8 +312,8 @@ class MOVADenoisingStage(PipelineStage):
def _manage_device_placement(
self,
model_to_use: torch.nn.Module | None,
model_to_offload: torch.nn.Module | None,
model_to_use: nn.Module | None,
model_to_offload: nn.Module | None,
server_args: ServerArgs,
):
if not server_args.dit_cpu_offload: