chore: import upstream snapshot with attribution

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

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from sglang.multimodal_gen.runtime.server_args import server_args as _server_args
+from sglang.multimodal_gen.runtime.server_args.server_args import (
+    BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS,
+    BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS,
+    DEFAULT_BCG_TEXT_BUCKETS,
+    LORA_MERGE_MODES,
+    LTX2_TWO_STAGE_DEVICE_MODE_CHOICES,
+    Backend,
+    PortArgs,
+    ServerArgs,
+    _normalize_ltx2_two_stage_device_mode,
+    get_global_server_args,
+    is_ltx2_two_stage_pipeline_name,
+    prepare_server_args,
+    set_global_server_args,
+)
+
+__all__ = [
+    "Backend",
+    "BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS",
+    "BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS",
+    "DEFAULT_BCG_TEXT_BUCKETS",
+    "LORA_MERGE_MODES",
+    "LTX2_TWO_STAGE_DEVICE_MODE_CHOICES",
+    "PortArgs",
+    "ServerArgs",
+    "_normalize_ltx2_two_stage_device_mode",
+    "get_global_server_args",
+    "is_ltx2_two_stage_pipeline_name",
+    "prepare_server_args",
+    "set_global_server_args",
+]
+
+
+def __getattr__(name: str):
+    if name == "_global_server_args":
+        return _server_args._global_server_args
+    raise AttributeError(name)

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

@@ -0,0 +1,620 @@
+"""
+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.managers.memory_managers.layerwise_offload_components import (
+    LAYERWISE_OFFLOAD_DIT_GROUP,
+    LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP,
+    LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP,
+    LAYERWISE_OFFLOAD_VAE_GROUP,
+)
+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.server_args import ServerArgs
+
+logger = init_logger(__name__)
+
+PERFORMANCE_MODES = ("manual", "auto", "speed", "memory")
+
+DEFAULT_LAYERWISE_COMPONENT_ARG_NAMES = (
+    (LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP, "text_encoder_cpu_offload"),
+    (LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP, "image_encoder_cpu_offload"),
+    (LAYERWISE_OFFLOAD_VAE_GROUP, "vae_cpu_offload"),
+)
+
+# task-type defaults for keep_resident_min_available_gb when a model does not pin
+# one: image vae is tiny so any datacenter gpu keeps it resident, video vae is
+# larger so it only stays resident on very-high-memory gpus
+IMAGE_GEN_KEEP_RESIDENT_MIN_AVAILABLE_GB = 45.0
+DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB = 120.0
+
+
+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()
+        self._explicit_layerwise_replacement_policy = (
+            self._has_explicit_layerwise_replacement_policy()
+        )
+
+    def _deployment_config(self) -> ModelDeploymentConfig:
+        return self.server_args.pipeline_config.get_model_deployment_config()
+
+    def _resolve_keep_resident_min_available_gb(
+        self, deployment_config: ModelDeploymentConfig
+    ) -> float | None:
+        # explicit per-model > task-type default > global default
+        explicit = deployment_config.keep_resident_min_available_gb
+        if explicit is not None:
+            return explicit
+        if self.server_args.pipeline_config.task_type.is_image_gen():
+            return IMAGE_GEN_KEEP_RESIDENT_MIN_AVAILABLE_GB
+        return DEFAULT_KEEP_RESIDENT_MIN_AVAILABLE_GB
+
+    def adjust_based_on_performance_mode(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 not args.enable_torch_compile and not args.is_arg_explicitly_set(
+                "enable_torch_compile"
+            ):
+                # speed means fastest: compile by default. An explicit
+                # --enable-torch-compile false still wins (e.g. models where
+                # compile measures slower, like short-step Z-Image runs).
+                args.enable_torch_compile = True
+                logger.info(
+                    "performance_mode=speed enables torch.compile "
+                    "(pass --enable-torch-compile false to opt out)"
+                )
+            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)
+                if (
+                    args.layerwise_offload_components is None
+                    and self._can_apply_default_layerwise_offload_policy()
+                ):
+                    args.layerwise_offload_components = (
+                        self._default_layerwise_components_for_unset_placement() or None
+                    )
+                return
+            args.use_fsdp_inference = False
+            if self._can_apply_default_layerwise_offload_policy():
+                # apply default 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()
+        deployment_config = self._deployment_config()
+        disable_threshold_gb = self._resolve_keep_resident_min_available_gb(
+            deployment_config
+        )
+        if (
+            min_available_gb is not None
+            and disable_threshold_gb is not None
+            and min_available_gb >= disable_threshold_gb
+        ):
+            changed = []
+            components = deployment_config.keep_resident_components
+            if (
+                args.layerwise_offload_components is not None
+                and not args.is_arg_explicitly_set("layerwise_offload_components")
+            ):
+                layerwise_components = [
+                    component_name
+                    for component_name in args.layerwise_offload_components
+                    if component_name not in components
+                ]
+                if layerwise_components != args.layerwise_offload_components:
+                    args.layerwise_offload_components = layerwise_components or None
+                    changed.append(
+                        f"layerwise_offload_components={args.layerwise_offload_components}"
+                    )
+            if (
+                args.dit_cpu_offload
+                and "dit" in components
+                and not args.is_arg_explicitly_set("dit_cpu_offload")
+            ):
+                args.dit_cpu_offload = False
+                changed.append("dit_cpu_offload=False")
+            if (
+                args.text_encoder_cpu_offload
+                and LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP in components
+                and not args.is_arg_explicitly_set("text_encoder_cpu_offload")
+            ):
+                args.text_encoder_cpu_offload = False
+                changed.append("text_encoder_cpu_offload=False")
+            if (
+                args.image_encoder_cpu_offload
+                and LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP in components
+                and not args.is_arg_explicitly_set("image_encoder_cpu_offload")
+            ):
+                args.image_encoder_cpu_offload = False
+                changed.append("image_encoder_cpu_offload=False")
+            if (
+                args.vae_cpu_offload
+                and LAYERWISE_OFFLOAD_VAE_GROUP in components
+                and not args.is_arg_explicitly_set("vae_cpu_offload")
+            ):
+                args.vae_cpu_offload = False
+                changed.append("vae_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),
+                )
+
+        self._maybe_keep_ltx23_resident_aux_components_resident()
+
+    def _maybe_keep_ltx23_resident_aux_components_resident(self) -> None:
+        args = self.server_args
+        if not args._uses_ltx23_high_memory_resident_two_stage_mode():
+            return
+
+        changed: list[str] = []
+        if (
+            args.layerwise_offload_components is not None
+            and not args.is_arg_explicitly_set("layerwise_offload_components")
+        ):
+            args.layerwise_offload_components = None
+            changed.append("layerwise_offload_components=None")
+
+        # high-memory resident mode keeps both DiTs on GPU; unset auxiliary
+        # placement should stay resident instead of using default layerwise
+        for arg_name in (
+            "text_encoder_cpu_offload",
+            "image_encoder_cpu_offload",
+            "vae_cpu_offload",
+        ):
+            if getattr(args, arg_name) and not args.is_arg_explicitly_set(arg_name):
+                setattr(args, arg_name, False)
+                changed.append(f"{arg_name}=False")
+
+        if changed:
+            logger.info(
+                "Keeping LTX2 high-memory two-stage auxiliary components resident: %s",
+                ", ".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_default_layerwise_offload(self) -> None:
+        """Enable verified layerwise defaults for unset component placement."""
+        args = self.server_args
+        if args.performance_mode != "auto":
+            return
+        if not self.could_override_server_args():
+            return
+        if (
+            args.layerwise_offload_components is not None
+            or args.dit_layerwise_offload is True
+        ):
+            return
+        if not current_platform.is_cuda():
+            return
+
+        layerwise_components = self._default_layerwise_components_for_unset_placement()
+        if not layerwise_components:
+            return
+
+        logger.info(
+            "Auto memory policy for %s selected layerwise offload components: %s",
+            args.pipeline_config.__class__.__name__,
+            ", ".join(layerwise_components),
+        )
+        args.layerwise_offload_components = layerwise_components
+
+    def maybe_replace_cpu_offloaded_components_with_layerwise(self) -> None:
+        args = self.server_args
+        if (
+            not self.could_override_server_args()
+            or self._explicit_layerwise_replacement_policy
+            or current_platform.is_cpu()
+            or not current_platform.is_cuda()
+            or envs.SGLANG_CACHE_DIT_ENABLED
+            or args.use_fsdp_inference
+            or args.layerwise_offload_components is not None
+        ):
+            return
+
+        layerwise_components: list[str] = []
+        if args.dit_layerwise_offload:
+            layerwise_components.append(LAYERWISE_OFFLOAD_DIT_GROUP)
+
+        changed: list[str] = []
+        if args.text_encoder_cpu_offload and not args.is_arg_explicitly_set(
+            "text_encoder_cpu_offload"
+        ):
+            layerwise_components.append(LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP)
+            changed.append(LAYERWISE_OFFLOAD_TEXT_ENCODER_GROUP)
+        if args.image_encoder_cpu_offload and not args.is_arg_explicitly_set(
+            "image_encoder_cpu_offload"
+        ):
+            layerwise_components.append(LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP)
+            changed.append(LAYERWISE_OFFLOAD_IMAGE_ENCODER_GROUP)
+        if args.vae_cpu_offload and not args.is_arg_explicitly_set("vae_cpu_offload"):
+            layerwise_components.append(LAYERWISE_OFFLOAD_VAE_GROUP)
+            changed.append(LAYERWISE_OFFLOAD_VAE_GROUP)
+
+        if not changed:
+            return
+
+        args.layerwise_offload_components = layerwise_components
+        logger.info(
+            "Automatically replacing CPU offload with layerwise offload for components: %s",
+            ", ".join(changed),
+        )
+
+    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.layerwise_offload_components is None:
+            args.layerwise_offload_components = (
+                self._default_layerwise_components_for_unset_placement() or None
+            )
+        if args.dit_cpu_offload is None:
+            args.dit_cpu_offload = True
+        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 _can_apply_default_layerwise_offload_policy(self) -> bool:
+        return current_platform.is_cuda()
+
+    def _default_layerwise_components_for_unset_placement(self) -> list[str]:
+        args = self.server_args
+        if args.pipeline_config.task_type.is_action_gen():
+            return []
+        if (
+            args.is_arg_explicitly_set("layerwise_offload_components")
+            or args.dit_layerwise_offload is True
+        ):
+            # The legacy --dit-layerwise-offload flag is a DiT-only selector.
+            # Do not merge implicit defaults into that explicit mode.
+            return []
+
+        # `*_cpu_offload` is the component placement knob. If a user explicitly
+        # set it to either true or false, keep that component out of default
+        # layerwise selection.
+        components = [
+            component_name
+            for component_name, arg_name in DEFAULT_LAYERWISE_COMPONENT_ARG_NAMES
+            if not args.is_arg_explicitly_set(arg_name)
+        ]
+        components = self._filter_high_memory_resident_components(components)
+        if self._should_auto_enable_dit_layerwise_offload():
+            components.insert(0, LAYERWISE_OFFLOAD_DIT_GROUP)
+            self._set_default_wan_dit_offload_prefetch_size()
+        return components
+
+    def _filter_high_memory_resident_components(
+        self, components: list[str]
+    ) -> list[str]:
+        args = self.server_args
+        if args.performance_mode != "auto" or current_platform.is_cpu():
+            return components
+
+        deployment_config = self._deployment_config()
+        threshold_gb = self._resolve_keep_resident_min_available_gb(deployment_config)
+        if threshold_gb is None:
+            return components
+
+        min_available_gb = self._get_min_available_device_memory_gb()
+        if min_available_gb is None or min_available_gb < threshold_gb:
+            return components
+
+        resident_components = set(deployment_config.keep_resident_components)
+        filtered_components = [
+            component
+            for component in components
+            if component not in resident_components
+        ]
+        skipped_components = [
+            component for component in components if component in resident_components
+        ]
+        if skipped_components:
+            logger.info(
+                "Keeping default layerwise components resident for %s because minimum available memory on selected GPUs is %.2f GiB: %s",
+                args.pipeline_config.__class__.__name__,
+                min_available_gb,
+                ", ".join(skipped_components),
+            )
+        return filtered_components
+
+    def _should_auto_enable_dit_layerwise_offload(self) -> bool:
+        args = self.server_args
+
+        # only for wan for now
+        if not self._is_wan_pipeline_config():
+            return False
+        if not self._deployment_config().auto_dit_layerwise_offload:
+            return False
+
+        if (
+            args.pipeline_config.dmd_denoising_steps is not None
+            or not current_platform.enable_dit_layerwise_offload_for_wan_by_default()
+            or envs.SGLANG_CACHE_DIT_ENABLED
+            or args.use_fsdp_inference
+            or args.is_arg_explicitly_set("dit_cpu_offload")
+        ):
+            return False
+
+        # memory mode is memory-first: keep the broad Wan DiT layerwise policy
+        # unless a guard above says it conflicts with another placement path
+        if args.performance_mode == "memory":
+            return True
+
+        # auto mode is performance-first: profiling only showed clear wins for
+        # Wan2.2 A14B, where coarse DiT CPU offload creates large step spikes
+        return (
+            args.performance_mode == "auto" and self._is_wan2_2_a14b_pipeline_config()
+        )
+
+    def _is_wan2_2_a14b_pipeline_config(self) -> bool:
+        config_name = self.server_args.pipeline_config.__class__.__name__
+        return config_name.startswith("Wan2_2_") and "A14B" in config_name
+
+    def _set_default_wan_dit_offload_prefetch_size(self) -> None:
+        args = self.server_args
+        if (
+            args.performance_mode == "auto"
+            and self._is_wan2_2_a14b_pipeline_config()
+            and not args.is_arg_explicitly_set("dit_offload_prefetch_size")
+        ):
+            # p2 was the fastest stable default in the Wan2.2 A14B sweep
+            args.dit_offload_prefetch_size = 2
+
+    def _is_wan_pipeline_config(self) -> bool:
+        return any(
+            cls.__module__.endswith(".wan")
+            for cls in self.server_args.pipeline_config.__class__.mro()
+        )
+
+    def _auto_uses_dit_offload(self) -> bool:
+        args = self.server_args
+        return bool(
+            args.dit_cpu_offload
+            or args.dit_layerwise_offload
+            or args.is_dit_layerwise_offload_selected
+        )
+
+    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 any(
+            args.is_arg_explicitly_set(arg_name)
+            for arg_name in (
+                "use_fsdp_inference",
+                "dit_cpu_offload",
+                "dit_layerwise_offload",
+                "layerwise_offload_components",
+            )
+        )
+
+    def _has_explicit_layerwise_replacement_policy(self) -> bool:
+        args = self.server_args
+        return any(
+            args.is_arg_explicitly_set(arg_name)
+            for arg_name in (
+                "dit_layerwise_offload",
+                "layerwise_offload_components",
+            )
+        )
+
+    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
+        deployment_config = self._deployment_config()
+        if (
+            deployment_config.auto_enable_cfg_parallel
+            and 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/runtime/server_args/disagg.py

@@ -0,0 +1,242 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Disaggregated diffusion server argument helpers."""
+
+from __future__ import annotations
+
+from typing import ClassVar, Literal
+
+from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
+from sglang.multimodal_gen.runtime.utils.common import (
+    format_tcp_endpoint,
+    parse_tcp_host_port,
+)
+from sglang.multimodal_gen.utils import FlexibleArgumentParser
+
+
+class DisaggServerArgsMixin:
+    DISAGG_RESULT_PORT_OFFSETS: ClassVar[dict[RoleType, int]] = {
+        RoleType.ENCODER: 1,
+        RoleType.DENOISER: 2,
+        RoleType.DECODER: 3,
+    }
+
+    def get_role_parallelism(self, role_type: RoleType) -> dict[str, int | None]:
+        _none = {
+            "tp_size": None,
+            "sp_degree": None,
+            "ulysses_degree": None,
+            "ring_degree": None,
+        }
+        if role_type == RoleType.ENCODER:
+            return {**_none, "tp_size": self.encoder_tp}
+        if role_type == RoleType.DENOISER:
+            return {
+                "tp_size": self.denoiser_tp,
+                "sp_degree": self.denoiser_sp,
+                "ulysses_degree": self.denoiser_ulysses,
+                "ring_degree": self.denoiser_ring,
+            }
+        if role_type == RoleType.DECODER:
+            return {**_none, "sp_degree": self.decoder_sp}
+        return _none
+
+    def derive_pool_result_endpoint(self) -> str:
+        host, base_port = parse_tcp_host_port(
+            self.disagg_server_addr, "disagg_server_addr"
+        )
+        role = (
+            self.disagg_role
+            if isinstance(self.disagg_role, RoleType)
+            else RoleType.from_string(self.disagg_role)
+        )
+        try:
+            offset = self.DISAGG_RESULT_PORT_OFFSETS[role]
+        except KeyError as exc:
+            raise ValueError(
+                "pool result endpoints are only defined for encoder, denoiser, "
+                f"and decoder roles, got {role.value!r}"
+            ) from exc
+        return format_tcp_endpoint(host, base_port + offset, "pool_result_endpoint")
+
+    def derive_pool_work_endpoint(self) -> str:
+        return format_tcp_endpoint("0.0.0.0", self.scheduler_port, "pool_work_endpoint")
+
+    def derive_pool_control_endpoint(self) -> str:
+        return format_tcp_endpoint(
+            "0.0.0.0", self.scheduler_port + 1, "pool_control_endpoint"
+        )
+
+    def derive_pool_control_advertised_endpoint(self) -> str:
+        host = self.host or self.disagg_p2p_hostname or "127.0.0.1"
+        if host == "0.0.0.0":
+            host = self.disagg_p2p_hostname or "127.0.0.1"
+        return format_tcp_endpoint(
+            host, self.scheduler_port + 1, "pool_control_advertised_endpoint"
+        )
+
+    def resolved_role_device(self) -> Literal["cpu", "cuda"]:
+        if self.disagg_role_device == "auto":
+            return "cpu" if self.num_gpus <= 0 else "cuda"
+        return self.disagg_role_device
+
+    @classmethod
+    def add_disagg_cli_args(cls, parser: FlexibleArgumentParser) -> None:
+        role_default = (
+            cls.disagg_role.value
+            if isinstance(cls.disagg_role, RoleType)
+            else cls.disagg_role
+        )
+        parser.add_argument(
+            "--disagg-role",
+            type=str,
+            default=role_default,
+            choices=RoleType.choices(),
+            help="Role for disaggregated pipeline.",
+        )
+        parser.add_argument(
+            "--disagg-timeout",
+            type=int,
+            default=cls.disagg_timeout,
+            help="Timeout in seconds for pending disagg requests. "
+            f"Default: {cls.disagg_timeout}.",
+        )
+        parser.add_argument(
+            "--disagg-downstream-wait-timeout",
+            type=int,
+            default=cls.disagg_downstream_wait_timeout,
+            help="Timeout in seconds while waiting for a downstream role slot. "
+            f"Default: {cls.disagg_downstream_wait_timeout}.",
+        )
+        parser.add_argument(
+            "--disagg-dispatch-policy",
+            type=str,
+            default=cls.disagg_dispatch_policy,
+            choices=["round_robin", "max_free_slots"],
+            help="Dispatch policy for pool mode disagg routing.",
+        )
+        parser.add_argument(
+            "--disagg-instance-id",
+            type=int,
+            default=cls.disagg_instance_id,
+            help="Stable per-role instance ID used by DiffusionServer registration.",
+        )
+        parser.add_argument(
+            "--disagg-max-slots-per-instance",
+            type=int,
+            default=cls.disagg_max_slots_per_instance,
+            help="Maximum concurrent transfer/computation slots tracked per instance.",
+        )
+        parser.add_argument(
+            "--disagg-transfer-redundancy",
+            type=float,
+            default=cls.disagg_transfer_redundancy,
+            help="Redundancy factor used when sizing transfer buffers from warmup.",
+        )
+        parser.add_argument(
+            "--disagg-role-device",
+            type=str,
+            default=cls.disagg_role_device,
+            choices=["auto", "cpu", "cuda"],
+            help=(
+                "Per-role device override. 'cpu' is intended for same-machine "
+                "encoder roles."
+            ),
+        )
+        parser.add_argument(
+            "--disagg-transfer-backend",
+            type=str,
+            default=cls.disagg_transfer_backend,
+            choices=["auto", "mock", "mooncake"],
+            help="Transfer backend for multimodal diffusion disaggregation.",
+        )
+        parser.add_argument(
+            "--disagg-transfer-pool-size",
+            type=int,
+            default=cls.disagg_transfer_pool_size,
+            help="Size of the P2P transfer buffer pool in bytes.",
+        )
+        parser.add_argument(
+            "--disagg-transfer-pin-memory",
+            type=str,
+            default=cls.disagg_transfer_pin_memory,
+            choices=["auto", "off", "required"],
+            help="CUDA host-register same-host shared-memory transfer buffers.",
+        )
+        parser.add_argument(
+            "--disagg-p2p-hostname",
+            type=str,
+            default=cls.disagg_p2p_hostname,
+            help="Hostname for P2P transfer engine.",
+        )
+        parser.add_argument(
+            "--disagg-ib-device",
+            type=str,
+            default=cls.disagg_ib_device,
+            help="InfiniBand device for P2P RDMA transfers.",
+        )
+        parser.add_argument(
+            "--disagg-server-addr",
+            type=str,
+            default=cls.disagg_server_addr,
+            help="DiffusionServer head node address for per-role launch mode.",
+        )
+        parser.add_argument(
+            "--encoder-urls",
+            type=str,
+            default=cls.encoder_urls,
+            help="Encoder instance work endpoints for DiffusionServer head mode.",
+        )
+        parser.add_argument(
+            "--denoiser-urls",
+            type=str,
+            default=cls.denoiser_urls,
+            help="Denoiser instance work endpoints for DiffusionServer head mode.",
+        )
+        parser.add_argument(
+            "--decoder-urls",
+            type=str,
+            default=cls.decoder_urls,
+            help="Decoder instance work endpoints for DiffusionServer head mode.",
+        )
+        parser.add_argument(
+            "--encoder-tp",
+            type=int,
+            default=cls.encoder_tp,
+            help="Tensor parallelism for encoder role.",
+        )
+        parser.add_argument(
+            "--denoiser-tp",
+            type=int,
+            default=cls.denoiser_tp,
+            help="Tensor parallelism for denoiser role.",
+        )
+        parser.add_argument(
+            "--denoiser-sp",
+            type=int,
+            default=cls.denoiser_sp,
+            help="Sequence parallelism for denoiser role.",
+        )
+        parser.add_argument(
+            "--denoiser-ulysses",
+            type=int,
+            default=cls.denoiser_ulysses,
+            help="Ulysses SP degree for denoiser role.",
+        )
+        parser.add_argument(
+            "--denoiser-ring",
+            type=int,
+            default=cls.denoiser_ring,
+            help="Ring SP degree for denoiser role.",
+        )
+        parser.add_argument(
+            "--decoder-sp",
+            type=int,
+            default=cls.decoder_sp,
+            help="Sequence parallelism for decoder role.",
+        )
+        parser.add_argument(
+            "--decoder-tp",
+            type=int,
+            default=cls.decoder_tp,
+            help="Deprecated alias for --decoder-sp.",
+        )