From 3705f9062954ff809585c4b85176c3e32b48d120 Mon Sep 17 00:00:00 2001 From: triple-mu Date: Thu, 22 Jan 2026 23:05:47 +0900 Subject: [PATCH] [diffusion] model: optimize torch.compile (#17472) --- .../pipelines_core/stages/denoising.py | 33 +++++++++---------- 1 file changed, 16 insertions(+), 17 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 47bd3c321..d4bd11387 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -15,6 +15,7 @@ from functools import lru_cache from typing import Any import torch +import torch.nn as nn from einops import rearrange from tqdm.auto import tqdm @@ -92,9 +93,8 @@ class DenoisingStage(PipelineStage): attn_head_size = hidden_size // num_attention_heads # torch compile - if self.server_args.enable_torch_compile: - for transformer in filter(None, [self.transformer, self.transformer_2]): - self.compile_module_with_torch_compile(transformer) + for transformer in filter(None, [self.transformer, self.transformer_2]): + self._maybe_enable_torch_compile(transformer) self.scheduler = scheduler self.vae = vae @@ -116,15 +116,15 @@ class DenoisingStage(PipelineStage): self._cached_num_steps = None self._is_warmed_up = False - def compile_module_with_torch_compile(self, module): + def _maybe_enable_torch_compile(self, module: object) -> None: """ - Compile a module's forward with torch.compile, and enable inductor overlap tweak if available. - No-op if torch compile is disabled or the object has no forward. + 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 self.server_args.enable_torch_compile or module is None: - return module - if not hasattr(module, "forward"): - return module + if not self.server_args.enable_torch_compile or not isinstance( + module, nn.Module + ): + return try: import torch._inductor.config as _inductor_cfg @@ -133,9 +133,8 @@ class DenoisingStage(PipelineStage): pass mode = os.environ.get("SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs") logger.info(f"Compiling transformer with mode: {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_enable_cache_dit(self, num_inference_steps: int, batch: Req) -> None: """Enable cache-dit on the transformers if configured (idempotent). @@ -497,7 +496,7 @@ class DenoisingStage(PipelineStage): ) # enable cache-dit before torch.compile (delayed mounting) self._maybe_enable_cache_dit(cache_dit_num_inference_steps, batch) - self.compile_module_with_torch_compile(self.transformer) + self._maybe_enable_torch_compile(self.transformer) if pipeline: pipeline.add_module("transformer", self.transformer) server_args.model_loaded["transformer"] = True @@ -801,8 +800,8 @@ class DenoisingStage(PipelineStage): def _manage_device_placement( self, - model_to_use: torch.nn.Module, - model_to_offload: torch.nn.Module | None, + model_to_use: nn.Module, + model_to_offload: nn.Module | None, server_args: ServerArgs, ): """ @@ -1218,7 +1217,7 @@ class DenoisingStage(PipelineStage): def _predict_noise_with_cfg( self, - current_model: torch.nn.Module, + current_model: nn.Module, latent_model_input: torch.Tensor, timestep, batch: Req,