[diffusion] feat: add performance mode server args (#24491)

python/sglang/benchmark/utils.py

@@ -118,13 +118,16 @@ def download_and_cache_file(url: str, filename: Optional[str] = None):
     chunk_size = 1024  # Download in chunks of 1KB
 
     # Use tqdm to display the progress bar
-    with open(filename, "wb") as f, tqdm(
-        desc=filename,
-        total=total_size,
-        unit="B",
-        unit_scale=True,
-        unit_divisor=1024,
-    ) as bar:
+    with (
+        open(filename, "wb") as f,
+        tqdm(
+            desc=filename,
+            total=total_size,
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as bar,
+    ):
         for chunk in response.iter_content(chunk_size=chunk_size):
             f.write(chunk)
             bar.update(len(chunk))

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/qwen_image.py

@@ -120,7 +120,7 @@ class QwenImageEditExecutor(QwenImageExecutor):
         ref_latents=None,
         additional_t_cond=None,
         transformer_options={},
-        **kwargs
+        **kwargs,
     ):
         """Forward pass for QwenImageEdit model."""
         latents, orig_shape = self._pack_latents(x)

python/sglang/multimodal_gen/configs/models/dits/wanvideo.py

@@ -28,6 +28,7 @@ class WanVideoArchConfig(DiTArchConfig):
             r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
             r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
             r"^blocks\.(\d+)\.attn1\.attn_op\.local_attn\.proj_l\.(.*)$": r"blocks.\1.attn1.local_attn.proj_l.\2",
+            r"^blocks\.(\d+)\.attn2\.norm_added_q\.(.*)$": "",
             r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2",
             r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
             r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -22,6 +22,9 @@ from sglang.multimodal_gen.configs.models import (
 )
 from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput
 from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
 from sglang.multimodal_gen.configs.sample.sampling_params import DataType
 from sglang.multimodal_gen.configs.utils import update_config_from_args
 from sglang.multimodal_gen.runtime.distributed.cfg_policy import CFGPolicy
@@ -240,6 +243,9 @@ class PipelineConfig:
     # image encoding
     image_encoder_extra_args: dict = field(default_factory=lambda: {})
 
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig()
+
     def postprocess_image(self, image):
         return image.last_hidden_state
 

python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py

@@ -16,6 +16,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ModelTaskType,
     PipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
 from sglang.multimodal_gen.runtime.distributed import (
     get_sp_parallel_rank,
     get_sp_world_size,
@@ -187,6 +190,12 @@ class LTX2PipelineConfig(PipelineConfig):
     def vae_temporal_compression(self):
         return self.vae_config.arch_config.temporal_compression_ratio
 
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(
+            auto_disable_component_offload_min_available_memory_gb=70,
+            auto_disable_component_offload_components=("dit",),
+        )
+
     def prepare_latent_shape(self, batch, batch_size, num_frames):
         """Return unpacked latent shape [B, C, F, H, W]."""
         height = batch.height // self.vae_scale_factor

python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py

@@ -0,0 +1,26 @@
+"""
+ModelDeploymentConfig provides model-specific config on how to deploy a model optimally
+
+"""
+
+from dataclasses import dataclass
+from typing import Literal
+
+OffloadComponentName = Literal["dit", "text_encoder", "image_encoder"]
+
+
+@dataclass(frozen=True)
+class ModelDeploymentConfig:
+    auto_dit_layerwise_offload: bool = False
+    # if the available memory is bigger than this value, keep dit resident instead of apply layerwise-offload
+    auto_dit_layerwise_offload_high_memory_disable_gb: float | None = None
+    auto_disable_component_offload_min_available_memory_gb: float | None = None
+    # keep this explicit because large encoders can OOM even when DiT fits resident
+    auto_disable_component_offload_components: tuple[OffloadComponentName, ...] = (
+        "dit",
+        "text_encoder",
+        "image_encoder",
+    )
+    fsdp_auto_min_available_memory_gb: float | None = None
+    fsdp_auto_requires_cfg: bool = True
+    fsdp_auto_requires_default_parallelism: bool = True

python/sglang/multimodal_gen/configs/pipeline_configs/mova.py

@@ -17,6 +17,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ModelTaskType,
     PipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.wan import t5_postprocess_text
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -52,6 +55,12 @@ class MOVAPipelineConfig(PipelineConfig):
     time_division_factor: int = 4
     time_division_remainder: int = 1
 
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(
+            auto_dit_layerwise_offload=True,
+            auto_dit_layerwise_offload_high_memory_disable_gb=130,
+        )
+
     def _center_crop_and_resize(
         self, image: torch.Tensor | Image.Image, target_height: int, target_width: int
     ) -> torch.Tensor | Image.Image:

python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py

@@ -168,6 +168,7 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig
     postprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
         default_factory=lambda: (qwen_image_postprocess_text,)
     )
+
     text_encoder_extra_args: list[dict] = field(
         default_factory=lambda: [
             dict(

python/sglang/multimodal_gen/configs/pipeline_configs/wan.py

@@ -18,6 +18,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ModelTaskType,
     PipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 logger = init_logger(__name__)
@@ -91,6 +94,12 @@ class WanT2V480PConfig(PipelineConfig):
         self.vae_config.load_encoder = False
         self.vae_config.load_decoder = True
 
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(
+            auto_dit_layerwise_offload=True,
+            auto_dit_layerwise_offload_high_memory_disable_gb=130,
+        )
+
 
 @dataclass
 class TurboWanT2V480PConfig(WanT2V480PConfig):
@@ -136,6 +145,12 @@ class WanI2V480PConfig(WanT2V480PConfig, WanI2VCommonConfig):
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
 
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(
+            auto_dit_layerwise_offload=True,
+            auto_dit_layerwise_offload_high_memory_disable_gb=130,
+        )
+
 
 @dataclass
 class WanI2V720PConfig(WanI2V480PConfig):

python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py

@@ -17,6 +17,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     TextConditioningOutput,
     pad_text_embeddings_with_mask,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
 from sglang.multimodal_gen.configs.post_training.pipeline_configs import (
     ZImageRolloutPipelineMixin,
 )
@@ -80,6 +83,9 @@ class ZImagePipelineConfig(ZImageRolloutPipelineMixin, ImagePipelineConfig):
     PATCH_SIZE: int = 2
     F_PATCH_SIZE: int = 1
 
+    def get_model_deployment_config(self) -> ModelDeploymentConfig:
+        return ModelDeploymentConfig(fsdp_auto_min_available_memory_gb=40)
+
     def tokenize_prompt(self, prompts: list[str], tokenizer, tok_kwargs) -> dict:
         rendered_prompts = [
             tokenizer.apply_chat_template(

python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py

@@ -456,6 +456,7 @@ class DiffGenerator:
         lora_path: Union[str, None, List[Union[str, None]]] = None,
         target: Union[str, List[str]] = "all",
         strength: Union[float, List[float]] = 1.0,
+        merge_mode: str | None = None,
     ) -> None:
         """
         Set LoRA adapter(s) for the specified transformer(s).
@@ -471,12 +472,14 @@ class DiffGenerator:
                 - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2)
                 - "critic": Apply only to the critic model
             strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats.
+            merge_mode: Optional LoRA merge mode: "auto", "merge", or "dynamic".
         """
         req = SetLoraReq(
             lora_nickname=lora_nickname,
             lora_path=lora_path,
             target=target,
             strength=strength,
+            merge_mode=merge_mode,
         )
         nickname_str, target_str, strength_str = format_lora_message(
             lora_nickname, target, strength

python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py

@@ -66,6 +66,7 @@ async def set_lora(
     lora_path: Optional[Union[str, List[Optional[str]]]] = Body(None, embed=True),
     target: Union[str, List[str]] = Body("all", embed=True),
     strength: Union[float, List[float]] = Body(1.0, embed=True),
+    merge_mode: Optional[str] = Body(None, embed=True),
 ):
     """
     Set LoRA adapter(s) for the specified transformer(s).
@@ -84,12 +85,14 @@ async def set_lora(
         strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats.
             If a list, must match the length of lora_nickname. Values < 1.0 reduce the effect,
             values > 1.0 amplify the effect.
+        merge_mode: Optional LoRA merge mode: "auto", "merge", or "dynamic".
     """
     req = SetLoraReq(
         lora_nickname=lora_nickname,
         lora_path=lora_path,
         target=target,
         strength=strength,
+        merge_mode=merge_mode,
     )
     nickname_str, target_str, strength_str = format_lora_message(
         lora_nickname, target, strength

python/sglang/multimodal_gen/runtime/entrypoints/utils.py

@@ -48,6 +48,7 @@ class SetLoraReq:
     lora_path: Optional[Union[str, List[Optional[str]]]] = None
     target: Union[str, List[str]] = "all"
     strength: Union[float, List[float]] = 1.0
+    merge_mode: Optional[str] = None
 
 
 @dataclass

python/sglang/multimodal_gen/runtime/layers/lora/linear.py

@@ -403,6 +403,9 @@ class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA):
         super().__init__(base_layer, lora_rank, lora_alpha)
 
     def forward(self, input_: torch.Tensor) -> torch.Tensor:
+        if self.merged or self.disable_lora:
+            return self.base_layer(input_)
+
         lora_A = self.lora_A
         lora_B = self.lora_B
         if isinstance(self.lora_B, DTensor):

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -769,6 +769,7 @@ class GPUWorker:
         lora_path: Union[str, None, List[Union[str, None]]] = None,
         target: Union[str, List[str]] = "all",
         strength: Union[float, List[float]] = 1.0,
+        merge_mode: str | None = None,
     ) -> OutputBatch:
         """
         Set the LoRA adapter(s) for the pipeline.
@@ -779,10 +780,13 @@ class GPUWorker:
             lora_path: Path(s) to the LoRA adapter(s). Can be a string, None, or a list of strings/None.
             target: Which transformer(s) to apply the LoRA to. Can be a string or a list of strings.
             strength: LoRA strength(s) for merge, default 1.0. Can be a float or a list of floats.
+            merge_mode: Optional per-request LoRA merge mode.
         """
         if not isinstance(self.pipeline, LoRAPipeline):
             return OutputBatch(error="Lora is not enabled")
-        self.pipeline.set_lora(lora_nickname, lora_path, target, strength)
+        self.pipeline.set_lora(
+            lora_nickname, lora_path, target, strength, merge_mode=merge_mode
+        )
         return OutputBatch()
 
     def merge_lora_weights(
@@ -868,16 +872,22 @@ class GPUWorker:
         return checksums
 
 
-OOM_MSG = f"""
+OOM_MSG = """
 OOM detected. Possible solutions:
   - If the OOM occurs during loading:
-    1. Enable CPU offload for memory-intensive components, or use `--dit-layerwise-offload` for DiT
+    1. Check available memory on every selected GPU, not only total capacity.
+       In multi-GPU runs, the least-free selected GPU is the bottleneck.
+    2. For single-GPU deployment, use `--performance-mode memory`, component CPU offload,
+       or `--dit-layerwise-offload` for supported Wan/MOVA DiTs.
+    3. For multi-GPU deployment, keep the default `--performance-mode auto` or set
+       `--use-fsdp-inference true` to shard DiT weights with FSDP. FSDP is not a
+       single-GPU substitute for CPU offload.
   - If the OOM occurs during runtime:
-    1. Enable SP and/or TP (in a multi-GPU setup)
-    2. Reduce the number of output tokens by lowering resolution or decreasing `--num-frames`
-    3. Opt for a sparse-attention backend
-    4. Enable FSDP by `--use-fsdp-inference` (in a multi-GPU setup)
-    5. Enable quantization (e.g. nunchaku)
+    1. Reduce resolution, `--num-frames`, or batch size.
+    2. Use `--performance-mode memory` for lower memory usage.
+    3. Enable SP/Ulysses/Ring for sequence-heavy workloads in multi-GPU setups.
+    4. Use FSDP, with CFG parallelism when supported, for validated multi-GPU workloads.
+    5. Use a lower-memory attention backend or quantization when available.
   Or, open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose
 """
 

python/sglang/multimodal_gen/runtime/managers/scheduler.py

@@ -189,7 +189,11 @@ class Scheduler(SchedulerDisaggMixin):
         # TODO: return with SetLoRAResponse or something more appropriate
         req = reqs[0]
         return self.worker.set_lora(
-            req.lora_nickname, req.lora_path, req.target, req.strength
+            req.lora_nickname,
+            req.lora_path,
+            req.target,
+            req.strength,
+            req.merge_mode,
         )
 
     def _handle_merge_lora(self, reqs: List[Any]):

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

@@ -559,6 +559,25 @@ class LTX2SnapshotResidencyStrategy(LTX2TwoStageResidencyStrategy):
         self.manager._sync_refinement_stage_transformer("stage1")
         self.manager._active_phase = "stage1"
 
+    def finish_request(
+        self,
+        module: torch.nn.Module,
+        use: ComponentUse,
+        state: ResidencyState,
+        *,
+        preferred: bool,
+    ) -> None:
+        if (
+            preferred
+            and state.batch_is_warmup
+            and self._snapshot_low_vram_mode
+            and self._phase(use) == "stage1"
+        ):
+            # keep the text encoder warm, but avoid stage1 DiT overlap before the first real request
+            self.manager._active_phase = None
+            return
+        super().finish_request(module, use, state, preferred=preferred)
+
     def finish_use(
         self,
         module: torch.nn.Module,
@@ -618,6 +637,13 @@ class LTX2SnapshotResidencyStrategy(LTX2TwoStageResidencyStrategy):
         if not self.server_args.dit_cpu_offload:
             return True
         phase = self._phase(use)
+        if (
+            self._snapshot_low_vram_mode
+            and phase == "stage1"
+            and state.current_use is not None
+            and state.current_use.component_name.startswith("text_encoder")
+        ):
+            return False
         if phase == "stage2":
             if self._snapshot_strategy.is_ready("transformer_2"):
                 return True

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

@@ -10,6 +10,7 @@ from typing import Any
 import torch
 import torch.distributed as dist
 from safetensors.torch import load_file
+from torch.distributed.tensor import DTensor
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
 from sglang.multimodal_gen.runtime.layers.lora.linear import (
@@ -24,7 +25,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import
 from sglang.multimodal_gen.runtime.pipelines_core.lora_format_adapter import (
     normalize_lora_state_dict,
 )
-from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.server_args import LORA_MERGE_MODES, ServerArgs
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_lora
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -362,6 +363,7 @@ class LoRAPipeline(ComposedPipelineBase):
         module_name: str,
         target_nicknames: list[str],
         target_strengths: list[float],
+        target_merge_weights: bool,
         adapter_updated: bool,
     ) -> bool:
         """
@@ -376,11 +378,12 @@ class LoRAPipeline(ComposedPipelineBase):
         Returns:
             True if the configuration matches exactly (including order and strength), False otherwise.
         """
-        if not self.is_lora_merged.get(module_name, False):
-            return False
         if adapter_updated:
             return False  # Adapter was updated, need to reapply
 
+        if self.is_lora_merged.get(module_name, False) != target_merge_weights:
+            return False
+
         stored_config = self.cur_adapter_config.get(module_name)
         if stored_config is None:
             return False
@@ -392,6 +395,68 @@ class LoRAPipeline(ComposedPipelineBase):
             and stored_strengths == target_strengths
         )
 
+    @staticmethod
+    def _uses_dtensor_weights(lora_layers: dict[str, BaseLayerWithLoRA]) -> bool:
+        return any(isinstance(layer.weight, DTensor) for layer in lora_layers.values())
+
+    @staticmethod
+    def _has_active_unmerged_lora(
+        lora_layers: dict[str, BaseLayerWithLoRA],
+    ) -> bool:
+        return any(
+            not layer.merged and not layer.disable_lora
+            for layer in lora_layers.values()
+        )
+
+    def _is_lora_effective_for_module(
+        self,
+        module_name: str,
+        lora_layers: dict[str, BaseLayerWithLoRA],
+    ) -> bool:
+        return self.is_lora_merged.get(
+            module_name, False
+        ) or self._has_active_unmerged_lora(lora_layers)
+
+    def _resolve_lora_merge_mode(
+        self,
+        merge_weights: bool | None,
+        merge_mode: str | None,
+    ) -> str:
+        if merge_mode is None:
+            if merge_weights is not None:
+                merge_mode = "merge" if merge_weights else "dynamic"
+            else:
+                merge_mode = self.server_args.lora_merge_mode
+        if merge_mode not in LORA_MERGE_MODES:
+            raise ValueError(
+                f"Invalid LoRA merge mode: {merge_mode}. Valid modes: {LORA_MERGE_MODES}"
+            )
+        return merge_mode
+
+    def _should_merge_lora_for_layers(
+        self,
+        module_name: str,
+        lora_layers: dict[str, BaseLayerWithLoRA],
+        merge_mode: str,
+    ) -> bool:
+        if merge_mode == "dynamic":
+            return False
+        uses_dtensor_weights = self._uses_dtensor_weights(lora_layers)
+        if merge_mode == "auto":
+            if uses_dtensor_weights:
+                logger.info(
+                    "Using dynamic LoRA for %s because FSDP-sharded weights would require a full-gather merge.",
+                    module_name,
+                )
+                return False
+            return True
+        if uses_dtensor_weights:
+            logger.warning(
+                "Merging LoRA for %s with FSDP-sharded weights may require full-gather and can OOM.",
+                module_name,
+            )
+        return True
+
     def _apply_lora_to_layers(
         self,
         lora_layers: dict[str, BaseLayerWithLoRA],
@@ -516,14 +581,18 @@ class LoRAPipeline(ComposedPipelineBase):
 
     def is_lora_effective(self, target: str = "all") -> bool:
         """
-        Check if LoRA is currently effective (merged) for the specified target.
+        Check if LoRA is currently effective for the specified target.
 
         Args:
-            target: Which transformer to check. "all" returns True if any is merged.
+            target: Which transformer to check. "all" returns True if any is effective.
         """
-        if target == "all":
-            return any(self.is_lora_merged.values())
-        return self.is_lora_merged.get(target, False)
+        target_modules, error = self._get_target_lora_layers(target)
+        if error:
+            logger.warning("is_lora_effective: %s", error)
+        return any(
+            self._is_lora_effective_for_module(module_name, lora_layers_dict)
+            for module_name, lora_layers_dict in target_modules
+        )
 
     def is_lora_set(self, target: str = "all") -> bool:
         """
@@ -622,12 +691,15 @@ class LoRAPipeline(ComposedPipelineBase):
         lora_path: str | None | list[str | None] = None,
         target: str | list[str] = "all",
         strength: float | list[float] = 1.0,
-        merge_weights: bool = True,
+        merge_weights: bool | None = None,
+        merge_mode: str | None = None,
     ):  # type: ignore
         """
         Load LoRA adapter(s) into the pipeline and apply them to the specified transformer(s).
         Supports both single LoRA (backward compatible) and multiple LoRA adapters.
         """
+        merge_mode = self._resolve_lora_merge_mode(merge_weights, merge_mode)
+
         # Normalize inputs to lists for multi-LoRA support
         lora_nicknames, lora_paths, strengths, targets = self._normalize_lora_params(
             lora_nickname, lora_path, strength, target
@@ -700,15 +772,34 @@ class LoRAPipeline(ComposedPipelineBase):
 
                 # Skip if LoRA configuration matches exactly (including order and strength)
                 # Since all modules for the same target apply the same config, checking one is sufficient
-                first_module_name, _ = target_modules[0]
+                first_module_name, first_lora_layers_dict = target_modules[0]
+                first_effective_merge_weights = self._should_merge_lora_for_layers(
+                    first_module_name, first_lora_layers_dict, merge_mode
+                )
+                if not first_effective_merge_weights and len(tgt_nicknames) > 1:
+                    raise ValueError(
+                        "Dynamic LoRA currently supports only one adapter per target. "
+                        "Use merge_mode='merge' for multiple adapters."
+                    )
                 if self._check_lora_config_matches(
-                    first_module_name, tgt_nicknames, tgt_strengths, adapter_updated
+                    first_module_name,
+                    tgt_nicknames,
+                    tgt_strengths,
+                    first_effective_merge_weights,
+                    adapter_updated,
                 ):
                     logger.info("LoRA configuration matches exactly, skipping")
                     continue
 
                 # Apply LoRA to modules for this target
                 for module_name, lora_layers_dict in target_modules:
+                    effective_merge_weights = (
+                        first_effective_merge_weights
+                        if module_name == first_module_name
+                        else self._should_merge_lora_for_layers(
+                            module_name, lora_layers_dict, merge_mode
+                        )
+                    )
                     count = self._apply_lora_to_layers(
                         lora_layers_dict,
                         tgt_nicknames,
@@ -716,7 +807,7 @@ class LoRAPipeline(ComposedPipelineBase):
                         rank,
                         tgt_strengths,
                         clear_existing=True,
-                        merge_weights=merge_weights,
+                        merge_weights=effective_merge_weights,
                     )
                     adapted_count += count
                     self.cur_adapter_name[module_name] = merged_name
@@ -724,7 +815,7 @@ class LoRAPipeline(ComposedPipelineBase):
                         str(p or self.loaded_adapter_paths.get(n, ""))
                         for n, p in zip(tgt_nicknames, tgt_paths)
                     )
-                    self.is_lora_merged[module_name] = merge_weights
+                    self.is_lora_merged[module_name] = effective_merge_weights
                     self.cur_adapter_strength[module_name] = tgt_strengths[0]
                     # Store full configuration for multi-LoRA support (preserves order and all strengths)
                     self.cur_adapter_config[module_name] = (
@@ -733,7 +824,7 @@ class LoRAPipeline(ComposedPipelineBase):
                     )
 
         logger.info(
-            "Rank %d: LoRA adapter(s) %s applied to %d layers (targets: %s, strengths: %s, merge_weights=%s)",
+            "Rank %d: LoRA adapter(s) %s applied to %d layers (targets: %s, strengths: %s, merge_mode=%s)",
             rank,
             ", ".join(map(str, lora_paths)) if lora_paths else None,
             adapted_count,
@@ -743,7 +834,7 @@ class LoRAPipeline(ComposedPipelineBase):
                 if len(strengths) > 1
                 else f"{strengths[0]:.2f}"
             ),
-            merge_weights,
+            merge_mode,
         )
 
     def deactivate_lora_weights(self, target: str = "all") -> None:
@@ -799,6 +890,24 @@ class LoRAPipeline(ComposedPipelineBase):
         # Disable layerwise offload if enabled: load all layers to GPU
         with self._temporarily_disable_offload(target_modules=target_modules):
             for module_name, lora_layers_dict in target_modules:
+                if not self._should_merge_lora_for_layers(
+                    module_name, lora_layers_dict, self.server_args.lora_merge_mode
+                ):
+                    for layer in lora_layers_dict.values():
+                        if layer.lora_A is None:
+                            continue
+                        if layer.merged:
+                            layer.unmerge_lora_weights()
+                        layer.disable_lora = False
+                        layer.strength = strength
+                    self.is_lora_merged[module_name] = False
+                    self.cur_adapter_strength[module_name] = strength
+                    logger.info(
+                        "Dynamic LoRA activated for %s (strength: %s)",
+                        module_name,
+                        strength,
+                    )
+                    continue
                 if self.is_lora_merged.get(module_name, False):
                     # Check if strength is the same - if so, skip (idempotent)
                     if self.cur_adapter_strength.get(module_name) == strength:
@@ -815,13 +924,9 @@ class LoRAPipeline(ComposedPipelineBase):
                     )
                 for name, layer in lora_layers_dict.items():
                     # Only re-enable LoRA for layers that actually have LoRA weights
-                    has_lora_weights = (
-                        hasattr(layer, "lora_A") and layer.lora_A is not None
-                    )
-                    if not has_lora_weights:
+                    if layer.lora_A is None:
                         continue
-                    if hasattr(layer, "disable_lora"):
-                        layer.disable_lora = False
+                    layer.disable_lora = False
                     try:
                         layer.merge_lora_weights(strength=strength)
                     except Exception as e:
@@ -854,9 +959,17 @@ class LoRAPipeline(ComposedPipelineBase):
 
         for module_name, lora_layers_dict in target_modules:
             if not self.is_lora_merged.get(module_name, False):
-                logger.warning(
-                    "LoRA weights are not merged for %s, skipping", module_name
-                )
+                if self._has_active_unmerged_lora(lora_layers_dict):
+                    for layer in lora_layers_dict.values():
+                        if not layer.disable_lora:
+                            layer.disable_lora = True
+                    self.cur_adapter_strength.pop(module_name, None)
+                    self.cur_adapter_config.pop(module_name, None)
+                    logger.info("Unmerged LoRA weights deactivated for %s", module_name)
+                else:
+                    logger.warning(
+                        "LoRA weights are not merged for %s, skipping", module_name
+                    )
                 continue
             with self._temporarily_disable_offload(target_modules=target_modules):
                 for name, layer in lora_layers_dict.items():
@@ -897,7 +1010,14 @@ class LoRAPipeline(ComposedPipelineBase):
 
         def _module_status(module_name: str) -> list[dict] | None:
             # return list of dict to support multi-lora in the future
-            if not self.is_lora_merged.get(module_name, False):
+            if module_name == "transformer":
+                lora_layers = self.lora_layers
+            elif module_name == "transformer_2":
+                lora_layers = self.lora_layers_transformer_2
+            else:
+                lora_layers = self.lora_layers_critic
+
+            if not self._is_lora_effective_for_module(module_name, lora_layers):
                 return None
             else:
                 return [
@@ -905,6 +1025,11 @@ class LoRAPipeline(ComposedPipelineBase):
                         "nickname": self.cur_adapter_name.get(module_name, None),
                         "path": self.cur_adapter_path.get(module_name, None),
                         "merged": self.is_lora_merged.get(module_name, False),
+                        "mode": (
+                            "merged"
+                            if self.is_lora_merged.get(module_name, False)
+                            else "unmerged"
+                        ),
                         "strength": self.cur_adapter_strength.get(module_name, None),
                     }
                 ]

python/sglang/multimodal_gen/runtime/server_args.py

@@ -41,6 +41,10 @@ from sglang.multimodal_gen.runtime.platforms import (
     AttentionBackendEnum,
     current_platform,
 )
+from sglang.multimodal_gen.runtime.server_args_auto_tune import (
+    PERFORMANCE_MODES,
+    ServerArgsAutoTuner,
+)
 from sglang.multimodal_gen.runtime.utils.common import (
     is_port_available,
     is_valid_ipv6_address,
@@ -62,21 +66,11 @@ from sglang.multimodal_gen.utils import (
 
 logger = init_logger(__name__)
 
-# Derived from single-H200 benchmarking (~140.4 GiB total) at the maximum
-# supported 720p workloads with dit_layerwise_offload=False and
-# num_inference_steps=1:
-# - Wan-AI/Wan2.2-T2V-A14B-Diffusers, 1280x720, 81 frames:
-#   peak_reserved=108076 MB (~105.5 GiB), peak_allocated=97665 MB (~95.4 GiB)
-# - OpenMOSS-Team/MOVA-720p, 1280x720, 193 frames:
-#   peak_reserved=130264 MB (~127.2 GiB), peak_allocated=108819 MB (~106.3 GiB)
-# Also, on H200, enabling dit_layerwise_offload regressed latency noticeably on
-# our validated Wan/MOVA workloads, so use a 130 GiB cutoff to keep H200-class
-# GPUs on the faster no-offload default while preserving some headroom.
-WAN_LAYERWISE_OFFLOAD_AUTO_DISABLE_MEM_GB = 130
 LTX2_TWO_STAGE_DEVICE_MODES = ("original", "snapshot", "resident")
 LTX2_TWO_STAGE_PIPELINE_NAMES = ("LTX2TwoStagePipeline", "LTX2TwoStageHQPipeline")
 # H200-class GPUs (>=130 GiB total) can usually keep both LTX2 DiTs resident.
 LTX2_RESIDENT_AUTO_ENABLE_MEM_GB = 130
+LORA_MERGE_MODES = ("auto", "merge", "dynamic")
 
 
 def _normalize_ltx2_two_stage_device_mode(mode: str | None) -> str | None:
@@ -148,6 +142,7 @@ class ServerArgs(DisaggArgsMixin):
 
     # Parallelism
     num_gpus: int = 1
+    performance_mode: str = "auto"
     tp_size: Optional[int] = None
     sp_degree: Optional[int] = None
     # sequence parallelism
@@ -179,6 +174,7 @@ class ServerArgs(DisaggArgsMixin):
     lora_path: str | None = None
     lora_nickname: str = "default"  # for swapping adapters in the pipeline
     lora_scale: float = 1.0  # LoRA scale for merging (e.g., 0.125 for Hyper-SD)
+    lora_merge_mode: str = "auto"
     lora_weight_name: str | None = None
 
     # Component path overrides (key = model_index.json component name, value = path)
@@ -197,7 +193,7 @@ class ServerArgs(DisaggArgsMixin):
     text_encoder_cpu_offload: bool | None = None
     image_encoder_cpu_offload: bool | None = None
     vae_cpu_offload: bool | None = False
-    use_fsdp_inference: bool = False
+    use_fsdp_inference: bool | None = None
     pin_cpu_memory: bool = True
     ltx2_two_stage_device_mode: str | None = None
 
@@ -320,8 +316,15 @@ class ServerArgs(DisaggArgsMixin):
 
     def _adjust_parameters(self):
         """set defaults and normalize values."""
-        self._adjust_offload()
+        auto_tuner = ServerArgsAutoTuner(self)
+        auto_tuner.adjust()
+        if auto_tuner.could_override_server_args():
+            self._adjust_offload()
+            auto_tuner.maybe_adjust_auto_dit_layerwise_offload()
         self._adjust_ltx2_two_stage_device_mode()
+        if auto_tuner.could_override_server_args():
+            auto_tuner.maybe_adjust_auto_component_residency_after_offload()
+            auto_tuner.maybe_adjust_auto_fsdp_with_offload_enabled()
         self._adjust_path()
         self._adjust_quant_config()
         self._adjust_warmup()
@@ -331,6 +334,7 @@ class ServerArgs(DisaggArgsMixin):
         self._adjust_attention_backend()
         self._adjust_platform_specific()
         self._adjust_autocast()
+        auto_tuner.finalize_auto_flags()
         self.adjust_pipeline_config()
 
     def _validate_parameters(self):
@@ -477,6 +481,12 @@ class ServerArgs(DisaggArgsMixin):
             or is_ltx23_native_variant(self.pipeline_config.vae_config.arch_config)
         )
 
+    def _uses_ltx23_snapshot_two_stage_residency(self) -> bool:
+        return (
+            self.ltx2_two_stage_device_mode == "snapshot"
+            and self._is_ltx23_two_stage_pipeline()
+        )
+
     def _adjust_attention_backend(self):
         if self.attention_backend in ["fa3", "fa4"]:
             self.attention_backend = "fa"
@@ -667,13 +677,12 @@ class ServerArgs(DisaggArgsMixin):
             self.master_port = self.settle_port(self.master_port, 37)
 
     def _adjust_parallelism(self):
-        tp_unspecified = self.tp_size is None
         sp_unspecified = self.sp_degree is None
         ulysses_unspecified = self.ulysses_degree is None
         ring_unspecified = self.ring_degree is None
         cfg_unspecified = self.enable_cfg_parallel is None
 
-        if current_platform.is_cpu() and self.tp_size > 1:
+        if current_platform.is_cpu() and (self.tp_size or 1) > 1:
             # CPU platform reuse num_gpus to represent num cpu numa nodes as devices
             self.num_gpus = self.tp_size
 
@@ -698,7 +707,8 @@ class ServerArgs(DisaggArgsMixin):
         if cfg_unspecified:
             cfg_group_size = self.dp_size * self.tp_size * 2
             if (
-                self.num_gpus >= 2
+                self.performance_mode != "manual"
+                and self.num_gpus >= 2
                 and self.num_gpus % cfg_group_size == 0
                 and sp_unspecified
                 and ulysses_unspecified
@@ -787,38 +797,6 @@ class ServerArgs(DisaggArgsMixin):
             self.use_fsdp_inference = False
             self.dit_layerwise_offload = False
 
-        # automatically enable dit_layerwise_offload for Wan/MOVA models if appropriate
-        if not envs.SGLANG_CACHE_DIT_ENABLED:
-            pipeline_name_lower = self.pipeline_config.__class__.__name__.lower()
-            if (
-                "wan" in pipeline_name_lower or "mova" in pipeline_name_lower
-            ) and self.dit_layerwise_offload is None:
-                auto_enable_layerwise_offload = (
-                    current_platform.enable_dit_layerwise_offload_for_wan_by_default()
-                )
-                if auto_enable_layerwise_offload and current_platform.is_cuda():
-                    device_total_memory_gb = (
-                        current_platform.get_device_total_memory() / BYTES_PER_GB
-                    )
-                    if (
-                        device_total_memory_gb
-                        >= WAN_LAYERWISE_OFFLOAD_AUTO_DISABLE_MEM_GB
-                    ):
-                        logger.info(
-                            "Skipping automatic dit_layerwise_offload for %s on a high-memory CUDA GPU (e.g. H200/B200/B300-class, %.2f GiB total)",
-                            self.pipeline_config.__class__.__name__,
-                            device_total_memory_gb,
-                        )
-                        auto_enable_layerwise_offload = False
-                        self.dit_layerwise_offload = False
-
-                if auto_enable_layerwise_offload:
-                    logger.info(
-                        f"Automatically enable dit_layerwise_offload for {self.pipeline_config.__class__.__name__} "
-                        "for low memory and performance balance"
-                    )
-                    self.dit_layerwise_offload = True
-
     def _adjust_autocast(self):
         if self.disable_autocast is None:
             self.disable_autocast = not self.pipeline_config.enable_autocast
@@ -967,6 +945,21 @@ class ServerArgs(DisaggArgsMixin):
             default=ServerArgs.num_gpus,
             help="The number of GPUs to use.",
         )
+        parser.add_argument(
+            "--performance-mode",
+            "--mode",
+            type=str,
+            choices=PERFORMANCE_MODES,
+            default=ServerArgs.performance_mode,
+            help=(
+                "Preset for performance and memory defaults. "
+                "'manual' keeps performance-related server args under explicit user control; "
+                "'auto' keeps safe defaults and applies high-confidence FSDP/CFG improvements; "
+                "'speed' favors GPU-resident execution for lower latency and higher throughput, and may OOM; "
+                "'memory' favors lower GPU memory usage; "
+                "Explicit offload/FSDP/parallelism flags take precedence."
+            ),
+        )
 
         parser.add_argument(
             "--tp-size",
@@ -994,9 +987,9 @@ class ServerArgs(DisaggArgsMixin):
         )
         parser.add_argument(
             "--enable-cfg-parallel",
-            action="store_true",
+            action=StoreBoolean,
             default=None,
-            help="Enable cfg parallel at degree 2. Auto-enabled when num_gpus >= 2 and no SP flags are set.",
+            help="Enable cfg parallel at degree 2. Auto-enabled when num_gpus >= 2 and no SP flags are set. Use false to disable it explicitly.",
         )
         parser.add_argument(
             "--cfg-parallel-size",
@@ -1106,7 +1099,7 @@ class ServerArgs(DisaggArgsMixin):
         parser.add_argument(
             "--use-fsdp-inference",
             action=StoreBoolean,
-            help="Use FSDP for inference by sharding the model weights. Latency is very low due to prefetch--enable if run out of memory.",
+            help="Use FSDP inference to shard DiT weights across GPUs. For single-GPU memory pressure, prefer CPU or layerwise offload.",
         )
         parser.add_argument(
             "--text-encoder-cpu-offload",
@@ -1270,6 +1263,17 @@ class ServerArgs(DisaggArgsMixin):
             default=ServerArgs.lora_scale,
             help="LoRA scale for merging (e.g., 0.125 for Hyper-SD). Same as lora_scale in Diffusers",
         )
+        parser.add_argument(
+            "--lora-merge-mode",
+            type=str,
+            choices=LORA_MERGE_MODES,
+            default=ServerArgs.lora_merge_mode,
+            help=(
+                "How LoRA is applied: auto keeps static merge for regular weights "
+                "and uses dynamic LoRA for FSDP-sharded weights to avoid full-gather; "
+                "merge always merges into base weights; dynamic always applies LoRA at forward time."
+            ),
+        )
         parser.add_argument(
             "--lora-weight-name",
             type=str,
@@ -1562,6 +1566,8 @@ class ServerArgs(DisaggArgsMixin):
                 # For '--arg=value', this gets 'arg'; for '--arg', this also gets 'arg'.
                 arg_name = arg.split("=", 1)[0].replace("-", "_").lstrip("_")
                 provided_arg_names.add(arg_name)
+        if "mode" in provided_arg_names:
+            provided_arg_names.add("performance_mode")
 
         # Populate provided_args if the argument from the namespace was on the command line.
         for k, v in vars(args).items():

python/sglang/multimodal_gen/runtime/server_args_auto_tune.py

@@ -0,0 +1,370 @@
+"""
+ServerArgsAutoTuner tunes the ServerArgs based on the desired performance mode
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from sglang.multimodal_gen import envs
+from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
+    ModelDeploymentConfig,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+logger = init_logger(__name__)
+
+PERFORMANCE_MODES = ("manual", "auto", "speed", "memory")
+
+
+class ServerArgsAutoTuner:
+    """Auto-tunes the server-arg for the given performance-mode, based on practical deployment experience with different model architectures"""
+
+    def __init__(self, server_args: "ServerArgs"):
+        self.server_args = server_args
+        self._explicit_memory_policy = self._has_explicit_memory_policy()
+
+    def _deployment_config(self) -> ModelDeploymentConfig:
+        return self.server_args.pipeline_config.get_model_deployment_config()
+
+    def adjust(self) -> None:
+        """Adjust the server args based on the performance mode"""
+        args = self.server_args
+        args.performance_mode = self._normalize_performance_mode()
+
+        if current_platform.is_cpu():
+            return
+
+        if args.performance_mode == "speed":
+            logger.info("Applying performance_mode=speed")
+            if args.num_gpus >= 2 and self._can_apply_fsdp_policy(
+                require_memory_headroom=False
+            ):
+                self._set_gpu_resident_defaults(use_fsdp=True)
+                self._enable_cfg_parallel_if_supported()
+            else:
+                self._set_gpu_resident_defaults(use_fsdp=False)
+            return
+
+        if args.performance_mode == "memory":
+            logger.info("Applying performance_mode=memory")
+            if args.use_fsdp_inference:
+                self._set_gpu_resident_defaults(use_fsdp=True)
+                return
+            args.use_fsdp_inference = False
+            if self._can_apply_dit_layerwise_offload_policy():
+                # apply dit layerwise offload to save VRAM during denoising stage
+                self._set_layerwise_offload_defaults()
+            else:
+                self._set_component_offload_defaults()
+            return
+
+    def maybe_adjust_auto_component_residency_after_offload(self) -> None:
+        args = self.server_args
+        if (
+            args.performance_mode != "auto"
+            or self._explicit_memory_policy
+            or current_platform.is_cpu()
+        ):
+            return
+
+        min_available_gb = self._get_min_available_device_memory_gb()
+        disable_threshold_gb = (
+            self._deployment_config().auto_disable_component_offload_min_available_memory_gb
+        )
+        if (
+            min_available_gb is not None
+            and disable_threshold_gb is not None
+            and min_available_gb >= disable_threshold_gb
+        ):
+            changed = []
+            components = (
+                self._deployment_config().auto_disable_component_offload_components
+            )
+            if args._uses_ltx23_snapshot_two_stage_residency():
+                # ltx2 snapshot mode uses DiT offload to release/prefetch stage DiTs between phases
+                components = tuple(
+                    component for component in components if component != "dit"
+                )
+            if args.dit_cpu_offload and "dit" in components:
+                args.dit_cpu_offload = False
+                changed.append("dit_cpu_offload=False")
+            if args.text_encoder_cpu_offload and "text_encoder" in components:
+                args.text_encoder_cpu_offload = False
+                changed.append("text_encoder_cpu_offload=False")
+            if args.image_encoder_cpu_offload and "image_encoder" in components:
+                args.image_encoder_cpu_offload = False
+                changed.append("image_encoder_cpu_offload=False")
+            if changed:
+                logger.info(
+                    "Disabling component offload for %s because minimum available memory on selected GPUs is %.2f GiB: %s",
+                    args.pipeline_config.__class__.__name__,
+                    min_available_gb,
+                    ", ".join(changed),
+                )
+
+    def maybe_adjust_auto_fsdp_with_offload_enabled(self) -> None:
+        args = self.server_args
+        if (
+            args.performance_mode == "auto"
+            and args.num_gpus >= 2
+            and not self._explicit_memory_policy
+            and self._auto_uses_dit_offload()
+            and self._can_apply_fsdp_policy(require_memory_headroom=True)
+        ):
+            logger.info(
+                "Automatically selecting FSDP defaults for multi-GPU %s to replace DiT offload",
+                args.pipeline_config.__class__.__name__,
+            )
+            args.use_fsdp_inference = True
+            if args.dit_cpu_offload:
+                args.dit_cpu_offload = False
+            if args.dit_layerwise_offload:
+                args.dit_layerwise_offload = False
+            self._enable_cfg_parallel_if_supported()
+
+    def maybe_adjust_auto_dit_layerwise_offload(self) -> None:
+        args = self.server_args
+        if not self.could_override_server_args():
+            return
+        if self._explicit_memory_policy:
+            return
+        deployment_config = self._deployment_config()
+        if envs.SGLANG_CACHE_DIT_ENABLED:
+            return
+        if (
+            not deployment_config.auto_dit_layerwise_offload
+            or args.dit_layerwise_offload is not None
+        ):
+            return
+        if args.use_fsdp_inference:
+            # if fsdp is enabled, layerwise-offload is weakened since the parameter has already been sharded
+            args.dit_layerwise_offload = False
+            return
+
+        auto_enable_layerwise_offload = (
+            current_platform.enable_dit_layerwise_offload_for_wan_by_default()
+        )
+        disable_threshold_gb = (
+            deployment_config.auto_dit_layerwise_offload_high_memory_disable_gb
+        )
+        if (
+            auto_enable_layerwise_offload
+            and current_platform.is_cuda()
+            and disable_threshold_gb is not None
+        ):
+            # auto turn off layerwise-offload if we have sufficient VRAM headroom
+            device_total_memory_gb = current_platform.get_device_total_memory() / (
+                1 << 30
+            )
+            if device_total_memory_gb >= disable_threshold_gb:
+                logger.info(
+                    "Skipping automatic dit_layerwise_offload for %s on a high-memory CUDA GPU (e.g. H200/B200/B300-class, %.2f GiB total)",
+                    args.pipeline_config.__class__.__name__,
+                    device_total_memory_gb,
+                )
+                auto_enable_layerwise_offload = False
+                args.dit_layerwise_offload = False
+
+        if auto_enable_layerwise_offload:
+            logger.info(
+                "Automatically enable dit_layerwise_offload for %s for low memory and performance balance",
+                args.pipeline_config.__class__.__name__,
+            )
+            args.dit_layerwise_offload = True
+            args.dit_cpu_offload = False
+
+    def finalize_auto_flags(self) -> None:
+        """if some args are unset after all the adjustment, set them to defaults"""
+        if not self.could_override_server_args():
+            return
+        args = self.server_args
+        if args.use_fsdp_inference is None:
+            args.use_fsdp_inference = False
+        if args.dit_cpu_offload is None:
+            args.dit_cpu_offload = False
+        if args.dit_layerwise_offload is None:
+            args.dit_layerwise_offload = False
+        if args.text_encoder_cpu_offload is None:
+            args.text_encoder_cpu_offload = False
+        if args.image_encoder_cpu_offload is None:
+            args.image_encoder_cpu_offload = False
+
+    def _normalize_performance_mode(self) -> str:
+        args = self.server_args
+        mode = (args.performance_mode or "auto").lower()
+        if mode not in PERFORMANCE_MODES:
+            valid_modes = PERFORMANCE_MODES
+            raise ValueError(
+                f"Invalid performance_mode={args.performance_mode!r}. "
+                f"Expected one of {valid_modes}."
+            )
+        return mode
+
+    def could_override_server_args(self) -> bool:
+        return self.server_args.performance_mode != "manual"
+
+    def _set_gpu_resident_defaults(self, *, use_fsdp: bool) -> None:
+        """set all components to be resident"""
+        args = self.server_args
+        changed = []
+        if args.use_fsdp_inference is None:
+            args.use_fsdp_inference = use_fsdp
+            changed.append(f"use_fsdp_inference={use_fsdp}")
+        if args.dit_cpu_offload is None:
+            args.dit_cpu_offload = False
+            changed.append("dit_cpu_offload=False")
+        if args.dit_layerwise_offload is None:
+            args.dit_layerwise_offload = False
+            changed.append("dit_layerwise_offload=False")
+        if args.text_encoder_cpu_offload is None:
+            args.text_encoder_cpu_offload = False
+            changed.append("text_encoder_cpu_offload=False")
+        if args.image_encoder_cpu_offload is None:
+            args.image_encoder_cpu_offload = False
+            changed.append("image_encoder_cpu_offload=False")
+
+        if changed:
+            logger.debug(
+                "Applied GPU-resident performance defaults: %s", ", ".join(changed)
+            )
+
+    def _set_component_offload_defaults(self) -> None:
+        args = self.server_args
+        changed = []
+        if args.dit_cpu_offload is None:
+            args.dit_cpu_offload = True
+            changed.append("dit_cpu_offload=True")
+        if args.text_encoder_cpu_offload is None:
+            args.text_encoder_cpu_offload = True
+            changed.append("text_encoder_cpu_offload=True")
+        if args.image_encoder_cpu_offload is None:
+            args.image_encoder_cpu_offload = True
+            changed.append("image_encoder_cpu_offload=True")
+        if args.use_fsdp_inference is None:
+            args.use_fsdp_inference = False
+            changed.append("use_fsdp_inference=False")
+
+        if changed:
+            logger.info(
+                "Applied low-memory component offload defaults: %s",
+                ", ".join(changed),
+            )
+
+    def _set_layerwise_offload_defaults(self) -> None:
+        args = self.server_args
+        if args.dit_layerwise_offload is None:
+            args.dit_layerwise_offload = True
+        if args.dit_cpu_offload is None:
+            args.dit_cpu_offload = False
+        if args.text_encoder_cpu_offload is None:
+            args.text_encoder_cpu_offload = True
+        if args.image_encoder_cpu_offload is None:
+            args.image_encoder_cpu_offload = True
+
+    def _can_apply_dit_layerwise_offload_policy(self) -> bool:
+        return (
+            self._deployment_config().auto_dit_layerwise_offload
+            and not envs.SGLANG_CACHE_DIT_ENABLED
+            and current_platform.enable_dit_layerwise_offload_for_wan_by_default()
+        )
+
+    def _auto_uses_dit_offload(self) -> bool:
+        args = self.server_args
+        return bool(args.dit_cpu_offload or args.dit_layerwise_offload)
+
+    def _get_min_available_device_memory_gb(self) -> float | None:
+        args = self.server_args
+        if current_platform.is_cpu():
+            return None
+
+        # Multi-GPU defaults are limited by the least-free selected GPU.
+        return min(
+            current_platform.get_available_gpu_memory(
+                device_id=device_id,
+                empty_cache=False,
+            )
+            for device_id in range(
+                args.base_gpu_id, args.base_gpu_id + max(1, args.num_gpus)
+            )
+        )
+
+    def _has_explicit_memory_policy(self) -> bool:
+        args = self.server_args
+        return (
+            args.use_fsdp_inference is not None
+            or args.dit_cpu_offload is not None
+            or args.dit_layerwise_offload is not None
+            or args.text_encoder_cpu_offload is not None
+            or args.image_encoder_cpu_offload is not None
+        )
+
+    def _has_explicit_parallel_policy(self) -> bool:
+        args = self.server_args
+        return (
+            args.tp_size is not None
+            or args.sp_degree is not None
+            or args.ulysses_degree is not None
+            or args.ring_degree is not None
+            or args.enable_cfg_parallel is not None
+        )
+
+    def _enable_cfg_parallel_if_supported(self) -> None:
+        args = self.server_args
+        if (
+            args.enable_cfg_parallel is None
+            and not self._has_explicit_parallel_policy()
+            and args._model_default_uses_cfg()
+        ):
+            args.enable_cfg_parallel = True
+
+    def _supports_high_confidence_fsdp(self) -> bool:
+        deployment_config = self._deployment_config()
+        return deployment_config.fsdp_auto_min_available_memory_gb is not None and (
+            not deployment_config.fsdp_auto_requires_cfg
+            or self.server_args._model_default_uses_cfg()
+        )
+
+    def _has_enough_available_memory_for_fsdp(self) -> bool:
+        args = self.server_args
+        min_available_gb = self._get_min_available_device_memory_gb()
+        if min_available_gb is None:
+            return True
+
+        required_gb = self._deployment_config().fsdp_auto_min_available_memory_gb
+        if required_gb is None:
+            return False
+        if min_available_gb < required_gb:
+            logger.info(
+                "Skipping automatic FSDP defaults: minimum available memory on selected GPUs %.2f GiB is below %.2f GiB for %s",
+                min_available_gb,
+                required_gb,
+                args.pipeline_config.__class__.__name__,
+            )
+            return False
+        return True
+
+    def _can_apply_fsdp_policy(self, *, require_memory_headroom: bool) -> bool:
+        args = self.server_args
+        deployment_config = self._deployment_config()
+        if not self._supports_high_confidence_fsdp():
+            return False
+        if envs.SGLANG_CACHE_DIT_ENABLED:
+            logger.info("Skipping automatic FSDP defaults because cache-dit is enabled")
+            return False
+        if (
+            args.performance_mode == "auto"
+            and deployment_config.fsdp_auto_requires_default_parallelism
+            and self._has_explicit_parallel_policy()
+        ):
+            logger.info(
+                "Skipping automatic FSDP defaults because an explicit parallel policy is set"
+            )
+            return False
+        return (
+            not require_memory_headroom or self._has_enough_available_memory_for_fsdp()
+        )

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -1105,7 +1105,7 @@
                 "TimestepPreparationStage": 7.85,
                 "LTX2AVLatentPreparationStage": 0.34,
                 "LTX2ImageEncodingStage": 0.02,
-                "LTX2AVDenoisingStage": 7744.44,
+                "LTX2AVDenoisingStage": 6000.0,
                 "LTX2UpsampleStage": 2.98,
                 "LTX2RefinementStage": 666.08,
                 "LTX2AVDecodingStage": 338.87,
@@ -1156,9 +1156,9 @@
                 "41": 223.51,
                 "42": 221.07
             },
-            "expected_e2e_ms": 10601.1,
-            "expected_avg_denoise_ms": 195.04,
-            "expected_median_denoise_ms": 191.34,
+            "expected_e2e_ms": 8500.0,
+            "expected_avg_denoise_ms": 150.0,
+            "expected_median_denoise_ms": 135.0,
             "estimated_full_test_time_s": 345.4
         },
         "wan2_2_ti2v_5b": {
@@ -2384,7 +2384,7 @@
                 "LTX2SigmaPreparationStage": 0.26,
                 "TimestepPreparationStage": 13.63,
                 "LTX2AVLatentPreparationStage": 0.12,
-                "LTX2AVDenoisingStage": 24658.05,
+                "LTX2AVDenoisingStage": 18500.0,
                 "LTX2AVDecodingStage": 383.25
             },
             "denoise_step_ms": {
@@ -2419,9 +2419,9 @@
                 "28": 815.91,
                 "29": 803.65
             },
-            "expected_e2e_ms": 26917.71,
-            "expected_avg_denoise_ms": 791.24,
-            "expected_median_denoise_ms": 791.0,
+            "expected_e2e_ms": 21000.0,
+            "expected_avg_denoise_ms": 620.0,
+            "expected_median_denoise_ms": 620.0,
             "estimated_full_test_time_s": 144.2
         },
         "ltx_2.3_two_stage_t2v_2gpus": {

python/sglang/multimodal_gen/test/server/test_server_common.py

@@ -2,7 +2,7 @@
 Config-driven diffusion generation test with pytest parametrization.
 
 
-If the actual run is significantly better than the baseline, the improved cases with their updated baseline will be printed
+Each collected request prints a performance log before validation.
 """
 
 from __future__ import annotations
@@ -225,13 +225,11 @@ class DiffusionServerBase:
     """
 
     _perf_results: list[dict[str, Any]] = []
-    _improved_baselines: list[dict[str, Any]] = []
     _pytest_config = None  # Store pytest config for stash access
 
     @classmethod
     def setup_class(cls):
         cls._perf_results = []
-        cls._improved_baselines = []
 
     @classmethod
     def teardown_class(cls):
@@ -251,20 +249,6 @@ class DiffusionServerBase:
                 "[DEBUG teardown_class] No pytest_config available, skipping stash update"
             )
 
-        if cls._improved_baselines:
-            import json
-
-            output = """
---- POTENTIAL BASELINE IMPROVEMENTS DETECTED ---
-The following test cases performed significantly better than their baselines.
-Consider updating perf_baselines.json with the snippets below:
-"""
-            for item in cls._improved_baselines:
-                output += (
-                    f'\n"{item["id"]}": {json.dumps(item["baseline"], indent=4)},\n'
-                )
-            print(output)
-
     @pytest.fixture(autouse=True)
     def _capture_pytest_config(self, request):
         """Capture pytest config for use in teardown_class."""
@@ -351,6 +335,7 @@ Consider updating perf_baselines.json with the snippets below:
         )
 
         summary = validator.collect_metrics(perf_record)
+        self._print_performance_log(case, summary, scenario)
 
         if case.run_perf_check:
             if is_baseline_generation_mode:
@@ -365,8 +350,6 @@ Consider updating perf_baselines.json with the snippets below:
                     )
                 return
 
-            self._check_for_improvement(case, summary, scenario)
-
             # only run performance validation if run_perf_check is True
             try:
                 validator.validate(perf_record, case.sampling_params.num_frames)
@@ -400,83 +383,43 @@ Consider updating perf_baselines.json with the snippets below:
             f"[DEBUG _validate_and_record] Appended result for {case.id}, class {self.__class__.__name__} now has {len(self.__class__._perf_results)} results"
         )
 
-    def _check_for_improvement(
+    def _print_performance_log(
         self,
         case: DiffusionTestCase,
         summary: PerformanceSummary,
-        scenario: "ScenarioConfig",
+        scenario: ScenarioConfig | None,
     ) -> None:
-        """Check for potential significant performance improvements and record them."""
-        is_improved = False
-        threshold = BASELINE_CONFIG.improvement_threshold
-
-        def is_sig_faster(actual, expected):
-            if expected == 0 or expected is None:
-                return False
-            return actual < expected * (1 - threshold)
-
-        def safe_get_metric(metric_dict, key):
-            val = metric_dict.get(key)
-            return val if val is not None else float("inf")
-
-        # Check for any significant improvement
-        if (
-            is_sig_faster(summary.e2e_ms, scenario.expected_e2e_ms)
-            or is_sig_faster(summary.avg_denoise_ms, scenario.expected_avg_denoise_ms)
-            or is_sig_faster(
-                summary.median_denoise_ms, scenario.expected_median_denoise_ms
+        lines = [
+            "",
+            f"--- Performance Log: {case.id} ---",
+            (
+                f"  e2e={summary.e2e_ms:.2f}ms, "
+                f"avg_denoise={summary.avg_denoise_ms:.2f}ms, "
+                f"median_denoise={summary.median_denoise_ms:.2f}ms"
+            ),
+        ]
+        if scenario is not None:
+            lines.append(
+                "  baseline: "
+                f"e2e={scenario.expected_e2e_ms:.2f}ms, "
+                f"avg_denoise={scenario.expected_avg_denoise_ms:.2f}ms, "
+                f"median_denoise={scenario.expected_median_denoise_ms:.2f}ms"
             )
-        ):
-            is_improved = True
-        # Combine metrics, always taking the better (lower) value
-        new_stages = {
-            stage: min(
-                safe_get_metric(summary.stage_metrics, stage),
-                safe_get_metric(scenario.stages_ms, stage),
+        if summary.stage_metrics:
+            stages = ", ".join(
+                f"{name}={duration:.2f}ms"
+                for name, duration in summary.stage_metrics.items()
             )
-            for stage in set(summary.stage_metrics) | set(scenario.stages_ms)
-        }
-        new_denoise_steps = {
-            step: min(
-                safe_get_metric(summary.all_denoise_steps, step),
-                safe_get_metric(scenario.denoise_step_ms, step),
+            lines.append(f"  stages: {stages}")
+        if summary.all_denoise_steps:
+            # ci retries need the exact outlier, not only sampled checkpoints
+            steps = ", ".join(
+                f"{idx}={duration:.2f}ms"
+                for idx, duration in sorted(summary.all_denoise_steps.items())
             )
-            for step in set(summary.all_denoise_steps.keys())
-            | set(scenario.denoise_step_ms)
-        }
-
-        # Check for stage-level improvements
-        if not is_improved:
-            for stage, new_val in new_stages.items():
-                if is_sig_faster(new_val, scenario.stages_ms.get(stage, float("inf"))):
-                    is_improved = True
-                    break
-        if not is_improved:
-            for step, new_val in new_denoise_steps.items():
-                if is_sig_faster(
-                    new_val, scenario.denoise_step_ms.get(step, float("inf"))
-                ):
-                    is_improved = True
-                    break
-
-        if is_improved:
-            new_baseline = {
-                "stages_ms": {k: round(v, 2) for k, v in new_stages.items()},
-                "denoise_step_ms": {
-                    str(k): round(v, 2) for k, v in new_denoise_steps.items()
-                },
-                "expected_e2e_ms": round(
-                    min(summary.e2e_ms, scenario.expected_e2e_ms), 2
-                ),
-                "expected_avg_denoise_ms": round(
-                    min(summary.avg_denoise_ms, scenario.expected_avg_denoise_ms), 2
-                ),
-                "expected_median_denoise_ms": round(
-                    min(summary.median_denoise_ms, scenario.expected_median_denoise_ms),
-                    2,
-                ),
-            }
-            self._improved_baselines.append({"id": case.id, "baseline": new_baseline})
+            lines.append(f"  denoise_steps: {steps}")
+        lines.append(f"--- End Performance Log: {case.id} ---")
+        print("\n".join(lines), flush=True)
 
     def _dump_baseline_for_testcase(
         self,

python/sglang/multimodal_gen/test/unit/test_resolve_prompts.py

@@ -51,9 +51,10 @@ class TestResolvePrompts(unittest.TestCase):
             os.unlink(path)
 
     def test_prompt_path_takes_priority_over_server_args(self):
-        with tempfile.NamedTemporaryFile(
-            "w", suffix=".txt", delete=False
-        ) as f1, tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as f2:
+        with (
+            tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as f1,
+            tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as f2,
+        ):
             f1.write("from prompt_path\n")
             f2.write("from server_args\n")
             path1, path2 = f1.name, f2.name

python/sglang/multimodal_gen/test/unit/test_server_args.py

@@ -12,9 +12,13 @@ from sglang.multimodal_gen.configs.pipeline_configs.base import (
     ModelTaskType,
     PipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     QwenImagePipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.wan import WanT2V480PConfig
+from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipelineConfig
 from sglang.multimodal_gen.registry import _get_config_info
 from sglang.multimodal_gen.runtime.models.dits.qwen_image import (
     QwenImageTransformer2DModel,
@@ -110,12 +114,34 @@ class TestServerArgsPathExpansion(unittest.TestCase):
             "torch_sdpa",
         ]
 
-        with patch.object(sys, "argv", ["sglang"] + argv):
-            args, unknown_args = parser.parse_known_args(argv)
-            with patch.object(
+        with (
+            patch.object(sys, "argv", ["sglang"] + argv),
+            patch.object(
                 PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
-            ):
-                server_args = ServerArgs.from_cli_args(args, unknown_args)
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.is_cpu",
+                return_value=False,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.is_mps",
+                return_value=False,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.is_cuda",
+                return_value=True,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.get_device_total_memory",
+                return_value=80 * 1024**3,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.get_available_gpu_memory",
+                return_value=80,
+            ),
+        ):
+            args, unknown_args = parser.parse_known_args(argv)
+            server_args = ServerArgs.from_cli_args(args, unknown_args)
 
         self.assertEqual(
             server_args.component_attention_backends, {"text_encoder": "torch_sdpa"}
@@ -123,6 +149,50 @@ class TestServerArgsPathExpansion(unittest.TestCase):
 
 
 class TestOffloadDefaults(unittest.TestCase):
+    def _from_dict_with_pipeline_config(
+        self,
+        pipeline_config,
+        *,
+        memory_gb=80,
+        available_memory_gb=None,
+        kwargs=None,
+    ):
+        def get_available_gpu_memory(device_id=0, **_kwargs):
+            if isinstance(available_memory_gb, dict):
+                return available_memory_gb[device_id]
+            if available_memory_gb is not None:
+                return available_memory_gb
+            return memory_gb
+
+        with (
+            patch.object(PipelineConfig, "from_kwargs", return_value=pipeline_config),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.is_cpu",
+                return_value=False,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.is_mps",
+                return_value=False,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.is_cuda",
+                return_value=True,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.enable_dit_layerwise_offload_for_wan_by_default",
+                return_value=True,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.get_device_total_memory",
+                return_value=memory_gb * 1024**3,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.get_available_gpu_memory",
+                side_effect=get_available_gpu_memory,
+            ),
+        ):
+            return ServerArgs.from_dict({"model_path": "/fake", **(kwargs or {})})
+
     def _from_dict_with_task_type(
         self,
         task_type,
@@ -142,6 +212,10 @@ class TestOffloadDefaults(unittest.TestCase):
                 "sglang.multimodal_gen.runtime.server_args.current_platform.get_device_total_memory",
                 return_value=memory_gb * 1024**3,
             ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.get_available_gpu_memory",
+                return_value=memory_gb,
+            ),
         ):
             return ServerArgs.from_dict({"model_path": "/fake", **(kwargs or {})})
 
@@ -151,7 +225,11 @@ class TestOffloadDefaults(unittest.TestCase):
         self.assertFalse(args.vae_cpu_offload)
 
     def test_vae_cpu_offload_defaults_false_on_low_memory_gpu(self):
-        args = self._from_dict_with_task_type(ModelTaskType.T2V, memory_gb=16)
+        args = self._from_dict_with_task_type(
+            ModelTaskType.T2V,
+            memory_gb=16,
+            kwargs={"performance_mode": "memory"},
+        )
 
         self.assertFalse(args.vae_cpu_offload)
         self.assertTrue(args.dit_cpu_offload)
@@ -166,6 +244,353 @@ class TestOffloadDefaults(unittest.TestCase):
 
         self.assertTrue(args.vae_cpu_offload)
 
+    def test_pipeline_configs_declare_auto_tune_hints(self):
+        qwen_deployment = QwenImagePipelineConfig().get_model_deployment_config()
+        wan_deployment = WanT2V480PConfig().get_model_deployment_config()
+        mova_deployment = MOVAPipelineConfig().get_model_deployment_config()
+        zimage_deployment = ZImagePipelineConfig().get_model_deployment_config()
+        ltx_deployment = LTX2PipelineConfig().get_model_deployment_config()
+
+        self.assertIsNone(qwen_deployment.fsdp_auto_min_available_memory_gb)
+        self.assertFalse(qwen_deployment.auto_dit_layerwise_offload)
+
+        self.assertIsNone(wan_deployment.fsdp_auto_min_available_memory_gb)
+        self.assertTrue(wan_deployment.auto_dit_layerwise_offload)
+
+        self.assertIsNone(mova_deployment.fsdp_auto_min_available_memory_gb)
+        self.assertTrue(mova_deployment.auto_dit_layerwise_offload)
+
+        self.assertEqual(zimage_deployment.fsdp_auto_min_available_memory_gb, 40)
+        self.assertTrue(zimage_deployment.fsdp_auto_requires_cfg)
+        self.assertFalse(zimage_deployment.auto_dit_layerwise_offload)
+
+        self.assertEqual(
+            ltx_deployment.auto_disable_component_offload_min_available_memory_gb, 70
+        )
+        self.assertEqual(
+            ltx_deployment.auto_disable_component_offload_components, ("dit",)
+        )
+
+    def test_manual_mode_preserves_unset_performance_args(self):
+        args = self._from_dict_with_pipeline_config(
+            QwenImagePipelineConfig(),
+            kwargs={
+                "model_path": "Qwen/Qwen-Image",
+                "num_gpus": 2,
+                "performance_mode": "manual",
+            },
+        )
+
+        self.assertEqual(args.performance_mode, "manual")
+        self.assertIsNone(args.use_fsdp_inference)
+        self.assertIsNone(args.dit_cpu_offload)
+        self.assertIsNone(args.dit_layerwise_offload)
+        self.assertIsNone(args.text_encoder_cpu_offload)
+        self.assertIsNone(args.image_encoder_cpu_offload)
+        self.assertFalse(args.enable_cfg_parallel)
+
+    def test_default_auto_keeps_legacy_single_gpu_offload_defaults(self):
+        args = self._from_dict_with_pipeline_config(
+            QwenImagePipelineConfig(),
+            kwargs={"model_path": "Qwen/Qwen-Image"},
+        )
+
+        self.assertEqual(args.performance_mode, "auto")
+        self.assertFalse(args.use_fsdp_inference)
+        self.assertTrue(args.dit_cpu_offload)
+        self.assertFalse(args.dit_layerwise_offload)
+        self.assertTrue(args.text_encoder_cpu_offload)
+        self.assertFalse(args.image_encoder_cpu_offload)
+
+    def test_auto_ltx_snapshot_keeps_dit_offload_with_headroom(self):
+        args = self._from_dict_with_pipeline_config(
+            LTX2PipelineConfig(),
+            available_memory_gb=76,
+            kwargs={
+                "model_path": "Lightricks/LTX-2.3",
+                "pipeline_class_name": "LTX2TwoStageHQPipeline",
+                "ltx2_two_stage_device_mode": "snapshot",
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertEqual(args.ltx2_two_stage_device_mode, "snapshot")
+        self.assertTrue(args.dit_cpu_offload)
+        self.assertTrue(args.text_encoder_cpu_offload)
+        self.assertTrue(args.image_encoder_cpu_offload)
+
+    def test_auto_wan_layerwise_offload_is_enabled_without_fsdp(self):
+        args = self._from_dict_with_pipeline_config(
+            WanT2V480PConfig(),
+            kwargs={"performance_mode": "auto"},
+        )
+
+        self.assertTrue(args.dit_layerwise_offload)
+        self.assertFalse(args.use_fsdp_inference)
+
+    def test_memory_wan_layerwise_offload_is_enabled_without_fsdp(self):
+        args = self._from_dict_with_pipeline_config(
+            WanT2V480PConfig(),
+            kwargs={"performance_mode": "memory"},
+        )
+
+        self.assertTrue(args.dit_layerwise_offload)
+        self.assertFalse(args.use_fsdp_inference)
+
+    def test_auto_wan_layerwise_offload_does_not_disable_explicit_fsdp(self):
+        args = self._from_dict_with_pipeline_config(
+            WanT2V480PConfig(),
+            kwargs={
+                "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+                "num_gpus": 2,
+                "performance_mode": "auto",
+                "use_fsdp_inference": True,
+            },
+        )
+
+        self.assertFalse(args.dit_layerwise_offload)
+        self.assertTrue(args.use_fsdp_inference)
+
+    def test_auto_multi_gpu_wan_uses_layerwise_offload_without_cfg(self):
+        with patch.object(ServerArgs, "_model_default_uses_cfg", return_value=False):
+            args = self._from_dict_with_pipeline_config(
+                WanT2V480PConfig(),
+                kwargs={
+                    "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+                    "num_gpus": 2,
+                    "performance_mode": "auto",
+                },
+            )
+
+        self.assertFalse(args.use_fsdp_inference)
+        self.assertFalse(args.enable_cfg_parallel)
+        self.assertFalse(args.dit_cpu_offload)
+        self.assertTrue(args.dit_layerwise_offload)
+
+    def test_auto_multi_gpu_qwen_keeps_legacy_offload_with_cfg(self):
+        args = self._from_dict_with_pipeline_config(
+            QwenImagePipelineConfig(),
+            kwargs={
+                "model_path": "Qwen/Qwen-Image",
+                "num_gpus": 2,
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertFalse(args.use_fsdp_inference)
+        self.assertTrue(args.enable_cfg_parallel)
+        self.assertTrue(args.dit_cpu_offload)
+        self.assertFalse(args.dit_layerwise_offload)
+        self.assertTrue(args.text_encoder_cpu_offload)
+        self.assertFalse(args.image_encoder_cpu_offload)
+
+    def test_auto_multi_gpu_zimage_base_prefers_fsdp(self):
+        args = self._from_dict_with_pipeline_config(
+            ZImagePipelineConfig(),
+            kwargs={
+                "model_path": "Tongyi-MAI/Z-Image",
+                "num_gpus": 2,
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertTrue(args.use_fsdp_inference)
+        self.assertTrue(args.enable_cfg_parallel)
+
+    def test_auto_multi_gpu_zimage_turbo_skips_fsdp(self):
+        args = self._from_dict_with_pipeline_config(
+            ZImagePipelineConfig(),
+            kwargs={
+                "model_path": "Tongyi-MAI/Z-Image-Turbo",
+                "num_gpus": 2,
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertFalse(args.use_fsdp_inference)
+        self.assertFalse(args.enable_cfg_parallel)
+
+    def test_auto_multi_gpu_qwen_preserves_explicit_fsdp_false(self):
+        args = self._from_dict_with_pipeline_config(
+            QwenImagePipelineConfig(),
+            kwargs={
+                "model_path": "Qwen/Qwen-Image",
+                "num_gpus": 2,
+                "performance_mode": "auto",
+                "use_fsdp_inference": False,
+            },
+        )
+
+        self.assertFalse(args.use_fsdp_inference)
+        self.assertTrue(args.enable_cfg_parallel)
+        self.assertTrue(args.dit_cpu_offload)
+        self.assertTrue(args.text_encoder_cpu_offload)
+        self.assertFalse(args.image_encoder_cpu_offload)
+
+    def test_auto_multi_gpu_qwen_skips_fsdp_when_available_memory_is_low(self):
+        args = self._from_dict_with_pipeline_config(
+            QwenImagePipelineConfig(),
+            memory_gb=50,
+            kwargs={
+                "model_path": "Qwen/Qwen-Image",
+                "num_gpus": 2,
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertFalse(args.use_fsdp_inference)
+        self.assertTrue(args.enable_cfg_parallel)
+        self.assertTrue(args.dit_cpu_offload)
+        self.assertTrue(args.text_encoder_cpu_offload)
+        self.assertFalse(args.image_encoder_cpu_offload)
+
+    def test_auto_multi_gpu_qwen_uses_selected_gpu_min_available_memory(self):
+        args = self._from_dict_with_pipeline_config(
+            QwenImagePipelineConfig(),
+            available_memory_gb={1: 50, 2: 80},
+            kwargs={
+                "model_path": "Qwen/Qwen-Image",
+                "base_gpu_id": 1,
+                "num_gpus": 2,
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertFalse(args.use_fsdp_inference)
+        self.assertTrue(args.enable_cfg_parallel)
+
+    def test_auto_multi_gpu_qwen_keeps_legacy_offload_with_headroom(self):
+        args = self._from_dict_with_pipeline_config(
+            QwenImagePipelineConfig(),
+            available_memory_gb={1: 72, 2: 80},
+            kwargs={
+                "model_path": "Qwen/Qwen-Image",
+                "base_gpu_id": 1,
+                "num_gpus": 2,
+                "performance_mode": "auto",
+            },
+        )
+
+        self.assertFalse(args.use_fsdp_inference)
+        self.assertTrue(args.enable_cfg_parallel)
+        self.assertTrue(args.dit_cpu_offload)
+        self.assertTrue(args.text_encoder_cpu_offload)
+        self.assertFalse(args.image_encoder_cpu_offload)
+
+    def test_speed_mode_single_gpu_disables_offload(self):
+        args = self._from_dict_with_pipeline_config(
+            QwenImagePipelineConfig(),
+            kwargs={
+                "model_path": "Qwen/Qwen-Image",
+                "performance_mode": "speed",
+            },
+        )
+
+        self.assertEqual(args.performance_mode, "speed")
+        self.assertFalse(args.use_fsdp_inference)
+        self.assertFalse(args.dit_cpu_offload)
+        self.assertFalse(args.dit_layerwise_offload)
+        self.assertFalse(args.text_encoder_cpu_offload)
+        self.assertFalse(args.image_encoder_cpu_offload)
+
+    def test_speed_mode_preserves_explicit_offload(self):
+        args = self._from_dict_with_pipeline_config(
+            QwenImagePipelineConfig(),
+            kwargs={
+                "model_path": "Qwen/Qwen-Image",
+                "performance_mode": "speed",
+                "dit_cpu_offload": True,
+            },
+        )
+
+        self.assertEqual(args.performance_mode, "speed")
+        self.assertTrue(args.dit_cpu_offload)
+        self.assertFalse(args.text_encoder_cpu_offload)
+        self.assertFalse(args.image_encoder_cpu_offload)
+
+    def test_memory_mode_wan_uses_layerwise_offload(self):
+        args = self._from_dict_with_pipeline_config(
+            WanT2V480PConfig(),
+            kwargs={
+                "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+                "performance_mode": "memory",
+            },
+        )
+
+        self.assertFalse(args.use_fsdp_inference)
+        self.assertTrue(args.dit_layerwise_offload)
+        self.assertFalse(args.dit_cpu_offload)
+        self.assertTrue(args.text_encoder_cpu_offload)
+        self.assertTrue(args.image_encoder_cpu_offload)
+
+    def test_memory_mode_preserves_explicit_fsdp(self):
+        args = self._from_dict_with_pipeline_config(
+            WanT2V480PConfig(),
+            kwargs={
+                "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+                "num_gpus": 2,
+                "performance_mode": "memory",
+                "use_fsdp_inference": True,
+            },
+        )
+
+        self.assertTrue(args.use_fsdp_inference)
+        self.assertFalse(args.dit_layerwise_offload)
+        self.assertFalse(args.dit_cpu_offload)
+
+    def test_invalid_performance_mode_raises(self):
+        with self.assertRaises(ValueError):
+            self._from_dict_with_pipeline_config(
+                QwenImagePipelineConfig(),
+                kwargs={"performance_mode": "turbo"},
+            )
+
+    def test_cfg_parallel_cli_can_be_disabled_explicitly(self):
+        parser = FlexibleArgumentParser()
+        ServerArgs.add_cli_args(parser)
+        argv = [
+            "--model-path",
+            "Qwen/Qwen-Image",
+            "--num-gpus",
+            "2",
+            "--performance-mode",
+            "auto",
+            "--enable-cfg-parallel",
+            "false",
+        ]
+
+        with (
+            patch.object(sys, "argv", ["sglang"] + argv),
+            patch.object(
+                PipelineConfig, "from_kwargs", return_value=QwenImagePipelineConfig()
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.is_cpu",
+                return_value=False,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.is_mps",
+                return_value=False,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.is_cuda",
+                return_value=True,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.get_device_total_memory",
+                return_value=80 * 1024**3,
+            ),
+            patch(
+                "sglang.multimodal_gen.runtime.server_args.current_platform.get_available_gpu_memory",
+                return_value=80,
+            ),
+        ):
+            args, unknown_args = parser.parse_known_args(argv)
+            server_args = ServerArgs.from_cli_args(args, unknown_args)
+
+        self.assertFalse(server_args.use_fsdp_inference)
+        self.assertFalse(server_args.enable_cfg_parallel)
+
 
 class TestFSDPShardConditions(unittest.TestCase):
     def test_helpers_match_only_direct_block_entries(self):

python/sglang/srt/configs/chatglm.py

@@ -43,7 +43,7 @@ class ChatGLMConfig(PretrainedConfig):
         quantization_bit=0,
         pre_seq_len=None,
         prefix_projection=False,
-        **kwargs
+        **kwargs,
     ):
         self.num_layers = num_layers
         self.vocab_size = padded_vocab_size

python/sglang/srt/configs/dots_ocr.py

@@ -16,7 +16,7 @@ class DotsOCRConfig(Qwen2Config):
         video_token_id=151656,
         vision_config: Optional[dict] = None,
         *args,
-        **kwargs
+        **kwargs,
     ):
         super().__init__(*args, **kwargs)
         self.image_token_id = image_token_id
@@ -42,7 +42,7 @@ class DotsVLProcessor(Qwen2_5_VLProcessor):
         tokenizer=None,
         video_processor=None,
         chat_template=None,
-        **kwargs
+        **kwargs,
     ):
         if video_processor is None:
             video_processor = DummyVideoProcessor()

python/sglang/srt/configs/exaone.py

@@ -161,7 +161,7 @@ class ExaoneConfig(PretrainedConfig):
         bos_token_id=0,
         eos_token_id=2,
         tie_word_embeddings=True,
-        **kwargs
+        **kwargs,
     ):
         self.vocab_size = vocab_size
         self.max_position_embeddings = max_position_embeddings
@@ -192,5 +192,5 @@ class ExaoneConfig(PretrainedConfig):
             bos_token_id=bos_token_id,
             eos_token_id=eos_token_id,
             tie_word_embeddings=tie_word_embeddings,
-            **kwargs
+            **kwargs,
         )

python/sglang/srt/configs/kimi_vl.py

@@ -18,7 +18,7 @@ class KimiVLConfig(PretrainedConfig):
         ignore_index: int = -100,
         media_placeholder_token_id: int = 163605,
         pad_token_id: int = 0,
-        **kwargs
+        **kwargs,
     ):
         if vision_config is None:
             vision_config = MoonViTConfig()

python/sglang/srt/distributed/parallel_state.py

@@ -1559,9 +1559,10 @@ def graph_capture(stream: Optional[torch.cuda.Stream] = None):
     in order to explicitly distinguish the kernels to capture
     from other kernels possibly launched on background in the default stream.
     """
-    with get_tp_group().graph_capture(
-        stream=stream
-    ) as context, get_pp_group().graph_capture(context):
+    with (
+        get_tp_group().graph_capture(stream=stream) as context,
+        get_pp_group().graph_capture(context),
+    ):
         with contextlib.ExitStack() as stack:
             seen = {id(_TP)}
             for group in (_MOE_EP, _MOE_TP):

python/sglang/srt/entrypoints/engine.py

@@ -621,8 +621,9 @@ class Engine(EngineScoreMixin, EngineBase):
                                 writer,
                             ),
                         )
-                        with memory_saver_adapter.configure_subprocess(), numa_utils.configure_subprocess(
-                            server_args, gpu_id
+                        with (
+                            memory_saver_adapter.configure_subprocess(),
+                            numa_utils.configure_subprocess(server_args, gpu_id),
                         ):
                             proc.start()
 

python/sglang/srt/entrypoints/openai/tool_server.py

@@ -19,9 +19,10 @@ logger = logging.getLogger(__name__)
 
 async def list_server_and_tools(server_url: str):
 
-    async with sse_client(url=server_url) as streams, ClientSession(
-        *streams
-    ) as session:
+    async with (
+        sse_client(url=server_url) as streams,
+        ClientSession(*streams) as session,
+    ):
         initialize_response = await session.initialize()
         list_tools_response = await session.list_tools()
         return initialize_response, list_tools_response
@@ -131,9 +132,10 @@ class MCPToolServer(ToolServer):
     async def get_tool_session(self, tool_name: str):
         url = self.urls.get(tool_name)
         if url:
-            async with sse_client(url=url) as streams, ClientSession(
-                *streams
-            ) as session:
+            async with (
+                sse_client(url=url) as streams,
+                ClientSession(*streams) as session,
+            ):
                 await session.initialize()
                 yield session
         else:

python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py

@@ -108,11 +108,14 @@ class NPUGraphRunner(CudaGraphRunner):
         else:
             skip_guard_context = empty_context()
 
-        with skip_guard_context, torch.npu.graph(
-            graph,
-            pool=pool,
-            stream=stream,
-            auto_dispatch_capture=True,
+        with (
+            skip_guard_context,
+            torch.npu.graph(
+                graph,
+                pool=pool,
+                stream=stream,
+                auto_dispatch_capture=True,
+            ),
         ):
             out = run_once_fn()
         return out

python/sglang/srt/layers/attention/triton_ops/rocm_mla_decode_rope.py

@@ -395,7 +395,7 @@ def _decode_grouped_att_m_fwd_rope(
         IS_NEOX_STYLE=is_neox_style,
         num_warps=4,
         num_stages=num_stages,
-        **extra_kargs
+        **extra_kargs,
     )
 
 

python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py

@@ -47,7 +47,7 @@ class QuarkW4A4MXFP4(QuarkLinearScheme):
         input_size_per_partition: int,
         params_dtype: torch.dtype,
         weight_loader: Callable,
-        **kwargs
+        **kwargs,
     ):
         output_size_per_partition = sum(output_partition_sizes)
         layer.logical_widths = output_partition_sizes

python/sglang/srt/managers/data_parallel_controller.py

@@ -545,8 +545,9 @@ class DataParallelController:
                             writer,
                         ),
                     )
-                    with memory_saver_adapter.configure_subprocess(), numa_utils.configure_subprocess(
-                        server_args, gpu_id
+                    with (
+                        memory_saver_adapter.configure_subprocess(),
+                        numa_utils.configure_subprocess(server_args, gpu_id),
                     ):
                         proc.start()
                 self.scheduler_procs.append(proc)

python/sglang/srt/mem_cache/memory_pool.py

@@ -249,10 +249,13 @@ class MambaPool:
             maybe_init_custom_mem_pool(device=self.device)
         )
 
-        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE), (
-            torch.cuda.use_mem_pool(self.custom_mem_pool)
-            if self.enable_custom_mem_pool
-            else nullcontext()
+        with (
+            self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE),
+            (
+                torch.cuda.use_mem_pool(self.custom_mem_pool)
+                if self.enable_custom_mem_pool
+                else nullcontext()
+            ),
         ):
             conv_state = [
                 torch.zeros(

python/sglang/srt/model_executor/breakable_cuda_graph_runner.py

@@ -308,9 +308,11 @@ class BreakableCudaGraphRunner:
 
     def _capture_all(self):
         """Capture breakable CUDA graphs for all token sizes."""
-        with freeze_gc(
-            self.model_runner.server_args.enable_cudagraph_gc
-        ), graph_capture() as graph_capture_context, enable_breakable_cuda_graph():
+        with (
+            freeze_gc(self.model_runner.server_args.enable_cudagraph_gc),
+            graph_capture() as graph_capture_context,
+            enable_breakable_cuda_graph(),
+        ):
             stream = graph_capture_context.stream
             pool = get_global_graph_memory_pool()
 

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -488,9 +488,10 @@ class PiecewiseCudaGraphRunner:
         # Trigger CUDA graph capture for specific shapes.
         # Capture the large shapes first so that the smaller shapes
         # can reuse the memory pool allocated for the large shapes.
-        with freeze_gc(
-            self.model_runner.server_args.enable_cudagraph_gc
-        ), graph_capture() as graph_capture_context:
+        with (
+            freeze_gc(self.model_runner.server_args.enable_cudagraph_gc),
+            graph_capture() as graph_capture_context,
+        ):
             stream = graph_capture_context.stream
             with set_pcg_capture_stream(stream):
                 avail_mem = get_available_gpu_memory(

python/sglang/srt/multimodal/processors/deepseek_vl_v2.py

@@ -42,7 +42,7 @@ class DeepseekVL2ImageProcessor(BaseMultimodalProcessor):
         request_obj,
         max_req_input_len,
         *args,
-        **kwargs
+        **kwargs,
     ):
         base_output = self.load_mm_data(
             input_text,

python/sglang/srt/speculative/eagle_worker.py

@@ -218,9 +218,11 @@ class EAGLEWorker(TpModelWorker):
             self.eagle_use_aux_hidden_state = eagle_config.get(
                 "use_aux_hidden_state", True
             )
-        with self.draft_tp_context(
-            self.draft_model_runner.tp_group
-        ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+        with (
+            self.draft_tp_context(self.draft_model_runner.tp_group),
+            speculative_moe_backend_context(),
+            speculative_moe_a2a_backend_context(),
+        ):
             self.init_attention_backend()
             self.init_cuda_graphs()
             if self.adaptive_controller is not None:
@@ -459,9 +461,11 @@ class EAGLEWorker(TpModelWorker):
                 seq_lens_cpu,
                 can_run_cuda_graph,
             ) = self.forward_target_extend(batch)
-            with self.draft_tp_context(
-                self.draft_model_runner.tp_group
-            ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            with (
+                self.draft_tp_context(self.draft_model_runner.tp_group),
+                speculative_moe_backend_context(),
+                speculative_moe_a2a_backend_context(),
+            ):
                 self.forward_draft_extend(
                     batch,
                     logits_output.hidden_states,
@@ -478,9 +482,11 @@ class EAGLEWorker(TpModelWorker):
         else:
             set_time_batch(batch.reqs, "set_spec_draft_start_time", trace_only=True)
 
-            with self.draft_tp_context(
-                self.draft_model_runner.tp_group
-            ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            with (
+                self.draft_tp_context(self.draft_model_runner.tp_group),
+                speculative_moe_backend_context(),
+                speculative_moe_a2a_backend_context(),
+            ):
                 verify_input = self.draft(batch)
 
             set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True)
@@ -502,9 +508,11 @@ class EAGLEWorker(TpModelWorker):
                 batch.reqs, "set_spec_draft_extend_start_time", trace_only=True
             )
 
-            with self.draft_tp_context(
-                self.draft_model_runner.tp_group
-            ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            with (
+                self.draft_tp_context(self.draft_model_runner.tp_group),
+                speculative_moe_backend_context(),
+                speculative_moe_a2a_backend_context(),
+            ):
                 # NOTE: We should use `check_forward_draft_extend_after_decode`
                 # when DP attention is enabled, but it is slow. Skip it for now.
                 draft_extend_input = verify_output.draft_extend_input

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -181,9 +181,11 @@ class EagleDraftWorker(BaseDraftWorker):
         self.draft_tp_context = (
             draft_tp_context if server_args.enable_dp_attention else empty_context
         )
-        with self.draft_tp_context(
-            self.draft_runner.tp_group
-        ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+        with (
+            self.draft_tp_context(self.draft_runner.tp_group),
+            speculative_moe_backend_context(),
+            speculative_moe_a2a_backend_context(),
+        ):
             self.init_attention_backend()
             self.init_cuda_graphs()
 
@@ -706,9 +708,13 @@ class EAGLEWorkerV2(BaseSpecWorker):
 
         # Build adaptive runtime states (must be after draft worker is fully initialized)
         if self.adaptive_controller is not None:
-            with self._draft_worker.draft_tp_context(
-                self._draft_worker.draft_runner.tp_group
-            ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            with (
+                self._draft_worker.draft_tp_context(
+                    self._draft_worker.draft_runner.tp_group
+                ),
+                speculative_moe_backend_context(),
+                speculative_moe_a2a_backend_context(),
+            ):
                 self.adaptive_controller.register(
                     SpecRuntimeState(
                         speculative_num_steps=self.speculative_num_steps,
@@ -758,9 +764,13 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 else CaptureHiddenMode.LAST
             )
             model_worker_batch.capture_hidden_mode = draft_capture_mode
-            with self.draft_worker.draft_tp_context(
-                self.draft_worker.draft_runner.tp_group
-            ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            with (
+                self.draft_worker.draft_tp_context(
+                    self.draft_worker.draft_runner.tp_group
+                ),
+                speculative_moe_backend_context(),
+                speculative_moe_a2a_backend_context(),
+            ):
                 batch_output.next_draft_input = (
                     self.draft_worker._draft_extend_for_prefill(
                         model_worker_batch,
@@ -784,9 +794,13 @@ class EAGLEWorkerV2(BaseSpecWorker):
                     topk=self.topk,
                     capture_hidden_mode=capture_mode,
                 )
-            with self.draft_worker.draft_tp_context(
-                self.draft_worker.draft_runner.tp_group
-            ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            with (
+                self.draft_worker.draft_tp_context(
+                    self.draft_worker.draft_runner.tp_group
+                ),
+                speculative_moe_backend_context(),
+                speculative_moe_a2a_backend_context(),
+            ):
                 verify_input: EagleVerifyInput = self.draft_worker.draft(
                     model_worker_batch
                 )
@@ -800,9 +814,13 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 self._draft_done_event.record()
             model_worker_batch.spec_info = verify_input
             batch_output = self.verify(model_worker_batch)
-            with self.draft_worker.draft_tp_context(
-                self.draft_worker.draft_runner.tp_group
-            ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            with (
+                self.draft_worker.draft_tp_context(
+                    self.draft_worker.draft_runner.tp_group
+                ),
+                speculative_moe_backend_context(),
+                speculative_moe_a2a_backend_context(),
+            ):
                 self.draft_worker._draft_extend_for_decode(
                     model_worker_batch, batch_output
                 )

python/sglang/srt/speculative/frozen_kv_mtp_worker.py

@@ -179,9 +179,11 @@ class FrozenKVMTPWorker(TpModelWorker):
         self.draft_model_runner.draft_attn_backend = self.draft_attn_backend
         self.cuda_graph_runner = None
 
-        with self.draft_tp_context(
-            self.draft_model_runner.tp_group
-        ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+        with (
+            self.draft_tp_context(self.draft_model_runner.tp_group),
+            speculative_moe_backend_context(),
+            speculative_moe_a2a_backend_context(),
+        ):
             self.init_cuda_graphs()
 
     @property
@@ -422,9 +424,11 @@ class FrozenKVMTPWorker(TpModelWorker):
                 seq_lens_cpu,
                 can_run_cuda_graph,
             ) = self.forward_target_extend(batch)
-            with self.draft_tp_context(
-                self.draft_model_runner.tp_group
-            ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+            with (
+                self.draft_tp_context(self.draft_model_runner.tp_group),
+                speculative_moe_backend_context(),
+                speculative_moe_a2a_backend_context(),
+            ):
                 self.forward_draft_extend(
                     batch,
                     logits_output.hidden_states,
@@ -440,9 +444,11 @@ class FrozenKVMTPWorker(TpModelWorker):
             )
 
         set_time_batch(batch.reqs, "set_spec_draft_start_time", trace_only=True)
-        with self.draft_tp_context(
-            self.draft_model_runner.tp_group
-        ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+        with (
+            self.draft_tp_context(self.draft_model_runner.tp_group),
+            speculative_moe_backend_context(),
+            speculative_moe_a2a_backend_context(),
+        ):
             verify_input = self.draft(batch)
         set_time_batch(batch.reqs, "set_spec_draft_end_time", trace_only=True)
         set_time_batch(batch.reqs, "set_spec_verify_start_time", trace_only=True)
@@ -458,9 +464,11 @@ class FrozenKVMTPWorker(TpModelWorker):
                 )
 
         set_time_batch(batch.reqs, "set_spec_draft_extend_start_time", trace_only=True)
-        with self.draft_tp_context(
-            self.draft_model_runner.tp_group
-        ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+        with (
+            self.draft_tp_context(self.draft_model_runner.tp_group),
+            speculative_moe_backend_context(),
+            speculative_moe_a2a_backend_context(),
+        ):
             draft_extend_input = verify_output.draft_extend_input
             if (
                 self.server_args.enable_dp_attention

python/sglang/srt/speculative/multi_layer_eagle_worker.py

@@ -193,9 +193,10 @@ class MultiLayerEagleWorker(TpModelWorker):
         self.draft_tp_context = (
             draft_tp_context if server_args.enable_dp_attention else empty_context
         )
-        with self.draft_tp_context(
-            self.mtp_model_runner(0).tp_group
-        ), speculative_moe_backend_context():
+        with (
+            self.draft_tp_context(self.mtp_model_runner(0).tp_group),
+            speculative_moe_backend_context(),
+        ):
             self.init_attention_backend()
             self.init_cuda_graphs()
 
@@ -265,9 +266,10 @@ class MultiLayerEagleWorker(TpModelWorker):
                 seq_lens_cpu,
                 can_run_cuda_graph,
             ) = self.forward_target_extend(batch)
-            with self.draft_tp_context(
-                self.mtp_model_runner(0).tp_group
-            ), speculative_moe_backend_context():
+            with (
+                self.draft_tp_context(self.mtp_model_runner(0).tp_group),
+                speculative_moe_backend_context(),
+            ):
                 self.forward_draft_extend(
                     batch, logits_output.hidden_states, next_token_ids, seq_lens_cpu
                 )
@@ -278,16 +280,18 @@ class MultiLayerEagleWorker(TpModelWorker):
                 can_run_cuda_graph=can_run_cuda_graph,
             )
         else:
-            with self.draft_tp_context(
-                self.mtp_model_runner(0).tp_group
-            ), speculative_moe_backend_context():
+            with (
+                self.draft_tp_context(self.mtp_model_runner(0).tp_group),
+                speculative_moe_backend_context(),
+            ):
                 verify_input = self.draft(batch)
             batch.spec_info = verify_input
             logits_output, verify_output, can_run_cuda_graph = self.verify(batch)
 
-            with self.draft_tp_context(
-                self.mtp_model_runner(0).tp_group
-            ), speculative_moe_backend_context():
+            with (
+                self.draft_tp_context(self.mtp_model_runner(0).tp_group),
+                speculative_moe_backend_context(),
+            ):
                 # NOTE: We should use `check_forward_draft_extend_after_decode`
                 # when DP attention is enabled, but it is slow. Skip it for now.
                 draft_extend_input = verify_output.draft_extend_input

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -172,9 +172,10 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
         self.draft_tp_context = (
             draft_tp_context if server_args.enable_dp_attention else empty_context
         )
-        with self.draft_tp_context(
-            self.draft_runner_list[0].tp_group
-        ), speculative_moe_backend_context():
+        with (
+            self.draft_tp_context(self.draft_runner_list[0].tp_group),
+            speculative_moe_backend_context(),
+        ):
             self.init_attention_backend()
             self.init_cuda_graphs()
 

python/sglang/srt/speculative/standalone_worker.py

@@ -77,7 +77,11 @@ class StandaloneWorker(EAGLEWorker):
             self.hot_token_id = None
 
         # Init draft worker
-        with empty_context(), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+        with (
+            empty_context(),
+            speculative_moe_backend_context(),
+            speculative_moe_a2a_backend_context(),
+        ):
             TpModelWorker.__init__(
                 self,
                 server_args=server_args,
@@ -102,9 +106,11 @@ class StandaloneWorker(EAGLEWorker):
         self.draft_tp_context = (
             draft_tp_context if server_args.enable_dp_attention else empty_context
         )
-        with self.draft_tp_context(
-            self.draft_model_runner.tp_group
-        ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
+        with (
+            self.draft_tp_context(self.draft_model_runner.tp_group),
+            speculative_moe_backend_context(),
+            speculative_moe_a2a_backend_context(),
+        ):
             self.init_attention_backend()
             self.init_cuda_graphs()
 

python/sglang/srt/speculative/standalone_worker_v2.py

@@ -116,9 +116,10 @@ class StandaloneDraftWorker(EagleDraftWorker):
         self.draft_tp_context = (
             draft_tp_context if server_args.enable_dp_attention else empty_context
         )
-        with self.draft_tp_context(
-            self.draft_runner.tp_group
-        ), speculative_moe_backend_context():
+        with (
+            self.draft_tp_context(self.draft_runner.tp_group),
+            speculative_moe_backend_context(),
+        ):
             self.init_attention_backend()
             self.init_cuda_graphs()
         self.tree_mask_mode = TreeMaskMode.FULL_MASK

python/sglang/test/server_fixtures/eagle_fixture.py

@@ -37,9 +37,10 @@ class EagleServerBase(CustomTestCase):
     @classmethod
     def setUpClass(cls):
         cls.base_url = DEFAULT_URL_FOR_TEST
-        with envs.SGLANG_SPEC_NAN_DETECTION.override(
-            True
-        ), envs.SGLANG_SPEC_OOB_DETECTION.override(True):
+        with (
+            envs.SGLANG_SPEC_NAN_DETECTION.override(True),
+            envs.SGLANG_SPEC_OOB_DETECTION.override(True),
+        ):
             cls.process = popen_launch_server(
                 cls.target_model,
                 cls.base_url,

python/sglang/test/simple_eval_common.py

@@ -528,13 +528,16 @@ def download_dataset(path, url):
         total_size = int(response.headers.get("content-length", 0))
         block_size = 8192
 
-        with open(path, "wb") as f, tqdm(
-            desc="Downloading",
-            total=total_size,
-            unit="iB",
-            unit_scale=True,
-            unit_divisor=1024,
-        ) as progress_bar:
+        with (
+            open(path, "wb") as f,
+            tqdm(
+                desc="Downloading",
+                total=total_size,
+                unit="iB",
+                unit_scale=True,
+                unit_divisor=1024,
+            ) as progress_bar,
+        ):
             for data in response.iter_content(block_size):
                 size = f.write(data)
                 progress_bar.update(size)

python/sglang/test/test_mm_utils.py

@@ -30,12 +30,12 @@ class TestMultimodalInputsFromDict(unittest.TestCase):
             model_specific_data={"image_grid_thw": [[1, 1, 1], [1, 1, 1]]},
         )
 
-        with patch.object(
-            schedule_batch.torch.cuda, "is_available", return_value=True
-        ), patch.object(
-            schedule_batch.torch.cuda, "current_device", return_value=0
-        ), patch.object(
-            schedule_batch.envs.SGLANG_MM_BUFFER_SIZE_MB, "get", return_value=0
+        with (
+            patch.object(schedule_batch.torch.cuda, "is_available", return_value=True),
+            patch.object(schedule_batch.torch.cuda, "current_device", return_value=0),
+            patch.object(
+                schedule_batch.envs.SGLANG_MM_BUFFER_SIZE_MB, "get", return_value=0
+            ),
         ):
             mm_inputs = MultimodalInputs.from_dict({"mm_items": [mm_item]})
 

python/sglang/utils.py

@@ -426,13 +426,16 @@ def download_and_cache_file(url: str, filename: Optional[str] = None):
     chunk_size = 1024  # Download in chunks of 1KB
 
     # Use tqdm to display the progress bar
-    with open(filename, "wb") as f, tqdm(
-        desc=filename,
-        total=total_size,
-        unit="B",
-        unit_scale=True,
-        unit_divisor=1024,
-    ) as bar:
+    with (
+        open(filename, "wb") as f,
+        tqdm(
+            desc=filename,
+            total=total_size,
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as bar,
+    ):
         for chunk in response.iter_content(chunk_size=chunk_size):
             f.write(chunk)
             bar.update(len(chunk))