[diffusion] chore: consolidate pipeline core hygiene (#33843)

python/sglang/multimodal_gen/runtime/models/dits/zimage.py

@@ -746,8 +746,6 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
     _supports_gradient_checkpointing = True
     _no_split_modules = ["ZImageTransformerBlock"]
     _fsdp_shard_conditions = ZImageDitConfig().arch_config._fsdp_shard_conditions
-    param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping
-
     param_names_mapping = ZImageDitConfig().arch_config.param_names_mapping
     reverse_param_names_mapping = (
         ZImageDitConfig().arch_config.reverse_param_names_mapping

python/sglang/multimodal_gen/runtime/pipelines/flux.py

@@ -6,29 +6,14 @@ from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
+    calculate_linear_shift,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.flux import (
     FluxProgressiveDenoisingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-# TODO(will): move PRECISION_TO_TYPE to better place
-
-logger = init_logger(__name__)
-
-
-def calculate_shift(
-    image_seq_len,
-    base_seq_len: int = 256,
-    max_seq_len: int = 4096,
-    base_shift: float = 0.5,
-    max_shift: float = 1.15,
-):
-    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
-    b = base_shift - m * base_seq_len
-    mu = image_seq_len * m + b
-    return mu
 
 
 def prepare_mu(batch: Req, server_args: ServerArgs):
@@ -41,15 +26,7 @@ def prepare_mu(batch: Req, server_args: ServerArgs):
         int(width) // (vae_scale_factor * 2)
     )
 
-    mu = calculate_shift(
+    return "mu", calculate_linear_shift(image_seq_len)
-        image_seq_len,
-        # hard code, since scheduler_config is not in PipelineConfig now
-        256,
-        4096,
-        0.5,
-        1.15,
-    )
-    return "mu", mu
 
 
 class FluxPipeline(LoRAPipeline, ComposedPipelineBase):

python/sglang/multimodal_gen/runtime/pipelines/flux_2.py

@@ -11,9 +11,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.
     Flux2ProgressiveDenoisingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
 
 
 def compute_empirical_mu(batch: Req, server_args: ServerArgs):

python/sglang/multimodal_gen/runtime/pipelines/glm_image.py

@@ -8,9 +8,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.g
     GlmImageBeforeDenoisingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
 
 
 class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):

python/sglang/multimodal_gen/runtime/pipelines/helios_pipeline.py

@@ -20,9 +20,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.h
     HeliosChunkedDenoisingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
 
 
 class HeliosPipeline(LoRAPipeline, ComposedPipelineBase):

python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py

@@ -16,11 +16,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages import (
     TextEncodingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-# TODO(will): move PRECISION_TO_TYPE to better place
-
-logger = init_logger(__name__)
 
 
 class HunyuanVideoPipeline(ComposedPipelineBase):

python/sglang/multimodal_gen/runtime/pipelines/krea2.py

@@ -25,9 +25,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.k
     Krea2BeforeDenoisingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
 
 _TEXT_MAX_LENGTH = 512
 

python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py

@@ -8,6 +8,9 @@ from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
+    calculate_linear_shift,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.qwen_image_layered import (
     QwenImageLayeredBeforeDenoisingStage,
@@ -16,26 +19,8 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.
     QwenImageProgressiveDenoisingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 
-# TODO(will): move PRECISION_TO_TYPE to better place
-
-logger = init_logger(__name__)
-
-
-def calculate_shift(
-    image_seq_len,
-    base_seq_len: int = 256,
-    max_seq_len: int = 4096,
-    base_shift: float = 0.5,
-    max_shift: float = 1.15,
-):
-    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
-    b = base_shift - m * base_seq_len
-    mu = image_seq_len * m + b
-    return mu
-
 
 def prepare_mu(batch: Req, server_args: ServerArgs):
     height = batch.height
@@ -44,15 +29,11 @@ def prepare_mu(batch: Req, server_args: ServerArgs):
     image_seq_len = (int(height) // vae_scale_factor // 2) * (
         int(width) // vae_scale_factor // 2
     )
-    mu = calculate_shift(
+    return "mu", calculate_linear_shift(
         image_seq_len,
-        # hard code, since scheduler_config is not in PipelineConfig now
+        max_seq_len=8192,
-        256,
+        max_shift=0.9,
-        8192,
-        0.5,
-        0.9,
     )
-    return "mu", mu
 
 
 class QwenImagePipeline(LoRAPipeline, ComposedPipelineBase):

python/sglang/multimodal_gen/runtime/pipelines/sana.py

@@ -20,9 +20,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages import (
     TextEncodingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
 
 
 class SanaPipeline(LoRAPipeline, ComposedPipelineBase):

python/sglang/multimodal_gen/runtime/pipelines/stable_diffusion_3.py

@@ -16,9 +16,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages import (
     TextEncodingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
 
 
 class SD3ConditioningStage(PipelineStage):

python/sglang/multimodal_gen/runtime/pipelines/wan_causal_dmd_pipeline.py

@@ -18,12 +18,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages import (
     InputValidationStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 # isort: on
 
-logger = init_logger(__name__)
-
 
 class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
     pipeline_name = "WanCausalDMDPipeline"

python/sglang/multimodal_gen/runtime/pipelines/wan_dmd_pipeline.py

@@ -16,7 +16,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import
 )
 from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 # isort: off
 from sglang.multimodal_gen.runtime.pipelines_core.stages import (
@@ -26,8 +25,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages import (
 
 # isort: on
 
-logger = init_logger(__name__)
-
 
 class WanDMDPipeline(LoRAPipeline, ComposedPipelineBase):
     """

python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py

@@ -17,9 +17,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import
 from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
 from sglang.multimodal_gen.runtime.pipelines_core.stages import DmdDenoisingStage
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
 
 
 class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase):

python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_pipeline.py

@@ -16,9 +16,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import
 )
 from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
 
 
 class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):

python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py

@@ -22,9 +22,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.
     WanProgressiveDenoisingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
 
 
 class WanPipeline(LoRAPipeline, ComposedPipelineBase):

python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py

@@ -5,26 +5,13 @@ from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline, Req
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.diffusion_scheduler_utils import (
+    calculate_linear_shift,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.progressive_resolution.zimage import (
     ZImageProgressiveDenoisingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
-
-
-def calculate_shift(
-    image_seq_len,
-    base_seq_len: int = 256,
-    max_seq_len: int = 4096,
-    base_shift: float = 0.5,
-    max_shift: float = 1.15,
-):
-    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
-    b = base_shift - m * base_seq_len
-    mu = image_seq_len * m + b
-    return mu
 
 
 def prepare_mu(batch: Req, server_args: ServerArgs):
@@ -34,15 +21,7 @@ def prepare_mu(batch: Req, server_args: ServerArgs):
     image_seq_len = ((int(height) // vae_scale_factor) // 2) * (
         (int(width) // vae_scale_factor) // 2
     )
-    mu = calculate_shift(
+    return "mu", calculate_linear_shift(image_seq_len)
-        image_seq_len,
-        # hard code, since scheduler_config is not in PipelineConfig now
-        256,
-        4096,
-        0.5,
-        1.15,
-    )
-    return "mu", mu
 
 
 class ZImagePipeline(LoRAPipeline, ComposedPipelineBase):

python/sglang/multimodal_gen/runtime/pipelines_core/diffusion_scheduler_utils.py

@@ -11,6 +11,19 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.platforms import current_platform
 
 
+def calculate_linear_shift(
+    image_seq_len: int,
+    *,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+) -> float:
+    """return the affine dynamic shift used by native flow schedulers"""
+    slope = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    return image_seq_len * slope + base_shift - slope * base_seq_len
+
+
 def clone_scheduler_runtime(scheduler: Any) -> Any:
     """Create an isolated scheduler runtime from a scheduler template or runtime."""
     return deepcopy(scheduler)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py

@@ -69,10 +69,10 @@ class DmdDenoisingStage(DenoisingStage):
         """
         prepared_vars = self._prepare_denoising_loop(batch, server_args)
 
-        target_dtype = prepared_vars["target_dtype"]
+        target_dtype = prepared_vars.target_dtype
-        autocast_enabled = prepared_vars["autocast_enabled"]
+        autocast_enabled = prepared_vars.autocast_enabled
-        num_warmup_steps = prepared_vars["num_warmup_steps"]
+        num_warmup_steps = prepared_vars.num_warmup_steps
-        latents = prepared_vars["latents"]
+        latents = prepared_vars.latents
         video_raw_latent_shape = latents.shape
         scheduler = self.scheduler
 
@@ -95,7 +95,7 @@ class DmdDenoisingStage(DenoisingStage):
             },
         )
 
-        pos_cond_kwargs = prepared_vars["pos_cond_kwargs"]
+        pos_cond_kwargs = prepared_vars.pos_cond_kwargs
 
         denoising_loop_start_time = time.time()
         with self.progress_bar(total=len(timesteps), batch=batch) as progress_bar:
@@ -113,15 +113,13 @@ class DmdDenoisingStage(DenoisingStage):
                 ):
                     t_int = int(t.item())
                     if self.transformer_2 is not None:
-                        current_model, current_guidance_scale = (
+                        current_model, _ = self._select_and_manage_model(
-                            self._select_and_manage_model(
+                            t_int=t_int,
-                                t_int=t_int,
+                            boundary_timestep=self._handle_boundary_ratio(
-                                boundary_timestep=self._handle_boundary_ratio(
+                                server_args, batch, scheduler
-                                    server_args, batch, scheduler
+                            ),
-                                ),
+                            server_args=server_args,
-                                server_args=server_args,
+                            batch=batch,
-                                batch=batch,
-                            )
                         )
                     else:
                         current_model = self.transformer
@@ -228,54 +226,3 @@ class DmdDenoisingStage(DenoisingStage):
         )
 
         return batch
-
-    def _select_and_manage_model(
-        self,
-        t_int: int,
-        boundary_timestep: float | None,
-        server_args: ServerArgs,
-        batch: Req,
-    ):
-        if boundary_timestep is None or t_int >= boundary_timestep:
-            # High-noise stage
-            current_model = self.transformer
-            current_guidance_scale = batch.guidance_scale
-            current_phase = "transformer"
-        else:
-            # Low-noise stage
-            current_model = self.transformer_2
-            current_guidance_scale = batch.guidance_scale_2
-            current_phase = "transformer_2"
-
-        self._manage_dit_use_site(current_model, current_phase, batch)
-
-        assert current_model is not None, "The model for the current step is not set."
-        return current_model, current_guidance_scale
-
-    def _handle_boundary_ratio(
-        self,
-        server_args,
-        batch,
-        scheduler,
-    ):
-        """
-        (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
-        """
-        boundary_ratio = server_args.pipeline_config.dit_config.boundary_ratio
-        if batch.boundary_ratio is not None:
-            logger.info(
-                "Overriding boundary ratio from %s to %s",
-                boundary_ratio,
-                batch.boundary_ratio,
-            )
-            boundary_ratio = batch.boundary_ratio
-
-        if boundary_ratio is not None:
-            num_train_timesteps = getattr(scheduler, "num_train_timesteps", None)
-            if num_train_timesteps is None:
-                num_train_timesteps = scheduler.config.num_train_timesteps
-            boundary_timestep = boundary_ratio * num_train_timesteps
-        else:
-            boundary_timestep = None
-
-        return boundary_timestep

python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py

@@ -40,7 +40,7 @@ class LatentPreparationFingerprint:
 
 @dataclass(frozen=True)
 class LatentPreparationSpec:
-    """ "dataclass for controlling the LatentPreparationStage runtime semantics"""
+    """Runtime configuration for LatentPreparationStage."""
 
     shape: tuple[int, ...]
     dtype: torch.dtype

python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py

@@ -82,8 +82,6 @@ class TimestepPreparationStage(PipelineStage):
         """
         Prepare timesteps for the diffusion process.
 
-
-
         Returns:
             The batch with prepared timesteps.
         """

python/sglang/multimodal_gen/runtime/server_args/server_args.py

@@ -262,7 +262,8 @@ class ServerArgs(DisaggServerArgsMixin):
     # filename logic.
     component_transformer_weights_paths: dict[str, str] = field(default_factory=dict)
 
-    # Quantization method for online quantization
+    # Explicit quantization method override (e.g. "mxfp8", "fp8", "modelslim").
+    # When set, the transformer loader uses it instead of auto-detection.
     quantization: str | None = None
     # Layer name patterns to skip during online quantization
     quantization_ignored_layers: list[str] | None = None
@@ -334,10 +335,6 @@ class ServerArgs(DisaggServerArgsMixin):
 
     disable_autocast: bool | None = None
 
-    # Explicit quantization method override (e.g. "mxfp8", "fp8", "modelslim").
-    # When set, the transformer loader will use this instead of auto-detection.
-    quantization: str | None = None
-
     # Quantization / Nunchaku SVDQuant configuration
     nunchaku_config: NunchakuSVDQuantArgs | NunchakuConfig | None = field(
         default_factory=NunchakuSVDQuantArgs, repr=False
@@ -386,9 +383,6 @@ class ServerArgs(DisaggServerArgsMixin):
         }
     )
 
-    # # DMD parameters
-    # dmd_denoising_steps: List[int] | None = field(default=None)
-
     # MoE parameters used by Wan2.2
     boundary_ratio: float | None = None