sgl-project--sglang/python/sglang/multimodal_gen/runtime/server_args/server_args.py

# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo

# SPDX-License-Identifier: Apache-2.0
# Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py
"""The arguments of sglang-diffusion Inference."""

import argparse
import dataclasses
import json
import math
import os
import random
import sys
import tempfile
from dataclasses import field
from enum import Enum
from typing import Any, List, Literal, Optional

import addict
import yaml

from sglang.multimodal_gen import envs
from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
    LTX2PipelineConfig,
    is_ltx23_native_variant,
)
from sglang.multimodal_gen.configs.quantization.nunchaku import NunchakuSVDQuantArgs
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
    NunchakuConfig,
)
from sglang.multimodal_gen.runtime.loader.utils import BYTES_PER_GB
from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload_components import (
    LAYERWISE_OFFLOAD_ALL_COMPONENTS,
    LAYERWISE_OFFLOAD_DIT_GROUP,
    cpu_offload_flags_for_layerwise_components,
    layerwise_component_matches_any_selection,
    normalize_layerwise_offload_components,
)
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.server_args.disagg import DisaggServerArgsMixin
from sglang.multimodal_gen.runtime.utils.common import (
    is_port_available,
    is_valid_ipv6_address,
    normalize_gpu_ids,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import (
    _sanitize_for_logging,
    configure_logger,
    init_logger,
)
from sglang.multimodal_gen.utils import (
    FlexibleArgumentParser,
    StoreBoolean,
    expand_path_fields,
)

logger = init_logger(__name__)

LTX2_TWO_STAGE_DEVICE_MODES = ("original", "resident")
LTX2_TWO_STAGE_DEVICE_MODE_CHOICES = (*LTX2_TWO_STAGE_DEVICE_MODES, "snapshot")
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:
    if mode is None:
        return None
    mode = mode.lower()
    if mode == "snapshot":
        logger.warning(
            "ltx2_two_stage_device_mode=snapshot is deprecated and is treated "
            "as original. Please use ltx2_two_stage_device_mode=original or "
            "resident instead. This alias may be removed after two release cycles."
        )
        return "original"
    return mode


def is_ltx2_two_stage_pipeline_name(pipeline_class_name: str | None) -> bool:
    return pipeline_class_name in LTX2_TWO_STAGE_PIPELINE_NAMES


class Backend(str, Enum):
    """
    Enumeration for different model backends.
    - AUTO: Automatically select backend (prefer sglang native, fallback to diffusers)
    - SGLANG: Use sglang's native optimized implementation
    - DIFFUSERS: Use vanilla diffusers pipeline (supports all diffusers models)
    """

    AUTO = "auto"
    SGLANG = "sglang"
    DIFFUSERS = "diffusers"

    @classmethod
    def from_string(cls, value: str) -> "Backend":
        """Convert string to Backend enum."""
        try:
            return cls(value.lower())
        except ValueError:
            raise ValueError(
                f"Invalid backend: {value}. Must be one of: {', '.join([m.value for m in cls])}"
            ) from None

    @classmethod
    def choices(cls) -> list[str]:
        """Get all available choices as strings for argparse."""
        return [backend.value for backend in cls]


WARMUP_MODES = ("off", "request", "server")

# Default prompt sequence-length buckets for breakable CUDA graph (BCG) padding.
# Prompt-conditioning is padded up to the smallest bucket that fits so prompts
# of different lengths share one captured graph.
DEFAULT_BCG_TEXT_BUCKETS = (64, 128, 256, 512, 1024)

BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS = frozenset(
    {
        "comfy-org/ideogram-4",
        "glm-image",
        "ideogram-4",
        "ideogram-4-fp8",
        "ideogram-4-nf4",
        "ideogram-ai/ideogram-4-fp8",
        "ideogram-ai/ideogram-4-nf4",
        "qwen/qwen-image",
        "qwen/qwen-image-2512",
        "qwen-image",
        "qwen-image-2512",
        "tongyi-mai/z-image",
        "tongyi-mai/z-image-turbo",
        "zai-org/glm-image",
        "z-image",
        "z-image-turbo",
    }
)

BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS = frozenset(
    {
        "GlmImagePipelineConfig",
        "Ideogram4PipelineConfig",
        "QwenImagePipelineConfig",
        "ZImagePipelineConfig",
    }
)


def _normalized_bcg_model_refs(model_ref: str | None) -> set[str]:
    if not model_ref:
        return set()

    normalized = str(model_ref).strip().rstrip("/").lower()
    refs = {normalized, os.path.basename(normalized)}

    if "models--" in normalized:
        hf_cache_name = normalized.split("models--", 1)[1].split("/", 1)[0]
        refs.add(hf_cache_name.replace("--", "/"))

    return refs


@dataclasses.dataclass
class ServerArgs(DisaggServerArgsMixin):
    # Model and path configuration (for convenience)
    model_path: str

    # explicit model ID override (e.g. "Qwen-Image")
    model_id: str | None = None

    # Model backend (sglang native or diffusers)
    backend: Backend = Backend.AUTO

    # Attention
    attention_backend: str = None
    attention_backend_config: addict.Dict | None = None
    component_attention_backends: dict[str, str] | str | None = field(
        default_factory=dict
    )
    cache_dit_config: str | dict[str, Any] | None = (
        None  # cache-dit config for diffusers
    )

    # Distributed executor backend
    nccl_port: Optional[int] = None

    # HuggingFace specific parameters
    trust_remote_code: bool = False
    revision: str | None = None

    # Parallelism
    num_gpus: int = 1
    performance_mode: str = "auto"
    base_gpu_id: int = 0
    gpu_ids: list[int] | None = None
    tp_size: Optional[int] = None
    sp_degree: Optional[int] = None
    # sequence parallelism
    ulysses_degree: Optional[int] = None
    ring_degree: Optional[int] = None
    # data parallelism
    # number of data parallelism groups
    dp_size: int = 1
    # number of gpu in a dp group
    dp_degree: int = 1
    # cfg parallel (None = auto-decide based on num_gpus)
    enable_cfg_parallel: Optional[bool] = None
    # number of GPUs in each CFG parallel group (None = auto, 1 = disabled, N > 1 = enabled)
    cfg_parallel_degree: Optional[int] = None

    hsdp_replicate_dim: int = 1
    hsdp_shard_dim: Optional[int] = None
    dist_timeout: int | None = 3600  # 1 hour

    pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False)

    # Pipeline override
    pipeline_class_name: str | None = (
        None  # Override pipeline class from model_index.json
    )

    # LoRA parameters
    # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated.
    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)
    component_paths: dict[str, str] = field(default_factory=dict)

    # path to pre-quantized transformer weights (single .safetensors or directory).
    transformer_weights_path: str | None = None
    # Per-component transformer weight overrides (key = model_index.json component name).
    # Pipelines use this when a checkpoint ships separate quantized weights for
    # secondary DiT components; the generic loader consumes it without model-specific
    # filename logic.
    component_transformer_weights_paths: dict[str, str] = field(default_factory=dict)

    # Quantization method for online quantization
    quantization: str | None = None
    # Layer name patterns to skip during online quantization
    quantization_ignored_layers: list[str] | None = None

    # can restrict layers to adapt, e.g. ["q_proj"]
    # Will adapt only q, k, v, o by default.
    lora_target_modules: list[str] | None = None

    # CPU offload parameters
    dit_cpu_offload: bool | None = None
    # if true, select the DiT layerwise group
    dit_layerwise_offload: bool | None = None
    layerwise_offload_components: list[str] | None = None
    dit_offload_prefetch_size: float = 0.0
    offload_during_compile: bool = True
    text_encoder_cpu_offload: bool | None = None
    image_encoder_cpu_offload: bool | None = None
    vae_cpu_offload: bool | None = False
    use_fsdp_inference: bool | None = None
    pin_cpu_memory: bool = True
    ltx2_two_stage_device_mode: str | None = None
    _explicit_arg_names: set[str] = field(default_factory=set, repr=False)

    # ComfyUI integration
    comfyui_mode: bool = False

    # Compilation
    enable_torch_compile: bool = False

    # Breakable CUDA graph (BCG): capture the DiT forward as CUDA-graph
    # segments split at attention modules (SP all-to-all / dynamic attention
    # stay eager). Mutually exclusive with --enable-torch-compile and
    # Cache-DiT; BCG takes priority when more than one is requested.
    #
    # BCG graphs are resolution-specific, so --warmup-resolutions is required
    # when BCG is enabled: every requested resolution is captured at warmup so
    # serving never triggers a fresh capture.
    enable_breakable_cuda_graph: bool = False
    # Text/prompt sequence-length padding budget for BCG. Prompt-conditioning
    # inputs are padded up to the smallest bucket that fits, so prompts of
    # different lengths reuse one captured graph. Warmup captures one graph per
    # bucket; a prompt longer than the largest bucket falls back to eager.
    # ``None`` resolves to DEFAULT_BCG_TEXT_BUCKETS.
    bcg_text_buckets: list[int] = None

    # NVTX profiling
    enable_layerwise_nvtx_marker: bool = False

    # warmup
    # `warmup_mode` is the canonical knob: one of WARMUP_MODES
    #   - "off":     no warmup.
    #   - "server":  server-based warmup — a synthetic request right after the
    #                server is ready, before real traffic
    #   - "request": request-based warmup — warm on the first real request(s).
    #                This is a BENCHMARK aid
    # existing consumers keep working) and as deprecated CLI aliases. None means
    # "derive the mode from the legacy booleans"; _adjust_warmup resolves it.
    warmup_mode: str | None = None

    # deprecated: warmup and server_warmup
    warmup: bool = False
    server_warmup: bool = False

    warmup_resolutions: list[str] = None
    warmup_steps: int = 1

    disable_autocast: bool | None = None

    # Explicit quantization method override (e.g. "mxfp8", "fp8", "modelslim").
    # When set, the transformer loader will use this instead of auto-detection.
    quantization: str | None = None

    # Quantization / Nunchaku SVDQuant configuration
    nunchaku_config: NunchakuSVDQuantArgs | NunchakuConfig | None = field(
        default_factory=NunchakuSVDQuantArgs, repr=False
    )

    # Master port for distributed inference
    master_port: int = 30005

    # http server endpoint config
    host: str | None = "127.0.0.1"
    port: int | None = 30000

    # TODO: webui and their endpoint, check if webui_port is available.
    webui: bool = False
    webui_port: int | None = 12312

    scheduler_port: int = 5555
    batching_mode: str = "dynamic"
    batching_max_size: int = 1
    batching_delay_ms: float = 0.0
    batching_config: str | None = None
    enable_batching_metrics: bool = False

    # Strict port mode: fail if requested port is unavailable instead of auto-selecting
    strict_ports: bool = False

    output_path: str | None = "outputs/"
    input_save_path: str | None = "inputs/uploads"

    # Prompt text file for batch processing
    prompt_file_path: str | None = None

    # model paths for correct deallocation
    model_paths: dict[str, str] = field(default_factory=dict)
    model_loaded: dict[str, bool] = field(
        default_factory=lambda: {
            "transformer": True,
            "vae": True,
            "video_vae": True,
            "audio_vae": True,
            "video_dit": True,
            "audio_dit": True,
            "dual_tower_bridge": True,
        }
    )

    # # DMD parameters
    # dmd_denoising_steps: List[int] | None = field(default=None)

    # MoE parameters used by Wan2.2
    boundary_ratio: float | None = None

    # Disaggregation (pool mode only — launched via launch_pool_disagg_server())
    disagg_role: RoleType = RoleType.MONOLITHIC
    disagg_timeout: int = 3600
    disagg_downstream_wait_timeout: int = 1800
    disagg_dispatch_policy: str = "round_robin"
    disagg_mode: bool = False
    disagg_instance_id: int = 0
    disagg_max_slots_per_instance: int = 8
    disagg_transfer_redundancy: float = 1.25
    disagg_role_device: Literal["auto", "cpu", "cuda"] = "auto"
    disagg_transfer_backend: Literal["auto", "mock", "mooncake"] = "auto"
    disagg_transfer_pool_size: int = 256 * 1024 * 1024
    disagg_transfer_pin_memory: Literal["auto", "off", "required"] = "auto"
    disagg_p2p_hostname: str = "127.0.0.1"
    disagg_ib_device: str | None = None
    disagg_server_addr: str | None = None
    encoder_urls: str | None = None
    denoiser_urls: str | None = None
    decoder_urls: str | None = None
    encoder_tp: int | None = None
    denoiser_tp: int | None = None
    denoiser_sp: int | None = None
    denoiser_ulysses: int | None = None
    denoiser_ring: int | None = None
    decoder_sp: int | None = None
    decoder_tp: int | None = None
    pool_work_endpoint: str | None = None
    pool_result_endpoint: str | None = None
    pool_control_endpoint: str | None = None
    pool_control_advertised_endpoint: str | None = None

    # Logging
    log_level: str = "info"
    log_requests: bool = False
    log_requests_level: int = 2
    log_requests_format: str = "text"
    log_requests_target: Optional[List[str]] = None
    uvicorn_access_log_exclude_prefixes: list[str] = field(default_factory=list)

    # Tracing
    enable_trace: bool = False
    otlp_traces_endpoint: str = "localhost:4317"

    # SGLang backend for encoder stage
    srt_encoder_url: str | None = None
    srt_encoder_connect_timeout: int = 3.05
    srt_encoder_timeout: int = 100

    @property
    def broker_port(self) -> int:
        return self.port + 1

    @property
    def is_local_mode(self) -> bool:
        """
        If no server is running when a generation task begins, 'local_mode' will be enabled: a dedicated server will be launched
        """
        return self.host is None or self.port is None

    def _adjust_path(self):
        expand_path_fields(self)
        self._adjust_save_paths()

    def _adjust_parameters(self):
        """set defaults and normalize values."""
        auto_tuner = ServerArgsAutoTuner(self)
        auto_tuner.adjust_based_on_performance_mode()
        self._adjust_disagg_parallelism_aliases()
        if auto_tuner.could_override_server_args():
            self._adjust_offload()
            auto_tuner.maybe_adjust_auto_default_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()
            auto_tuner.maybe_replace_cpu_offloaded_components_with_layerwise()
        self._adjust_path()
        self._adjust_quant_config()
        self._adjust_breakable_cuda_graph_support()
        self._adjust_warmup()
        self._adjust_network_ports()
        # adjust parallelism before attention backend
        self._adjust_parallelism()
        self._adjust_attention_backend()
        self._adjust_platform_specific()
        self._adjust_layerwise_offload_components()
        self._adjust_autocast()
        auto_tuner.finalize_auto_flags()
        self.adjust_pipeline_config()

    def _adjust_disagg_parallelism_aliases(self):
        if self.decoder_tp is None:
            return
        if self.decoder_sp is not None and self.decoder_sp != self.decoder_tp:
            raise ValueError(
                "decoder_tp is deprecated in favor of decoder_sp; "
                "please set only one of them or keep the same value."
            )
        if self.decoder_sp is None:
            logger.warning(
                "decoder_tp is deprecated and is treated as decoder_sp for "
                "decoder/VAE parallel decode. Please use decoder_sp instead."
            )
            self.decoder_sp = self.decoder_tp

    def _validate_parameters(self):
        """check consistency and raise errors for invalid configs"""
        self._validate_pipeline()
        self._validate_offload()
        if not current_platform.is_cpu():
            self._validate_parallelism()
        self._validate_cfg_parallel()
        self._validate_batching()
        self._validate_breakable_cuda_graph()

    def resolved_bcg_text_buckets(self) -> tuple[int, ...]:
        """Sorted, de-duplicated, positive BCG text buckets.

        Falls back to :data:`DEFAULT_BCG_TEXT_BUCKETS` when ``--bcg-text-buckets``
        is unset, so both prompt padding and warmup capture share one source of
        truth instead of the legacy ``SGLANG_BCG_TEXT_BUCKETS`` env var.
        """
        raw = self.bcg_text_buckets
        if not raw:
            return DEFAULT_BCG_TEXT_BUCKETS
        buckets = sorted({int(b) for b in raw if int(b) > 0})
        return tuple(buckets) or DEFAULT_BCG_TEXT_BUCKETS

    def _validate_breakable_cuda_graph(self):
        if not self.enable_breakable_cuda_graph:
            return
        # BCG graphs are captured per resolution and only replay for that exact
        # latent shape, so the user must declare the resolutions up front. We
        # capture every one of them at warmup; serving then never re-captures.
        if not self.warmup_resolutions:
            raise ValueError(
                "--enable-breakable-cuda-graph requires --warmup-resolutions: "
                "diffusion CUDA graphs only replay for a fixed resolution, so "
                "every served resolution must be declared and captured at "
                "warmup, e.g. --warmup-resolutions 1024x1024 1328x1328."
            )
        if self.bcg_text_buckets is not None and not any(
            int(b) > 0 for b in self.bcg_text_buckets
        ):
            raise ValueError(
                "--bcg-text-buckets must contain at least one positive integer."
            )

    def _adjust_breakable_cuda_graph_support(self):
        if not self.enable_breakable_cuda_graph:
            return

        pipeline_config = getattr(self, "pipeline_config", None)
        pipeline_config_name = type(pipeline_config).__name__
        if (
            pipeline_config_name in BREAKABLE_CUDA_GRAPH_SUPPORTED_PIPELINE_CONFIGS
            and self._is_breakable_cuda_graph_supported_model()
        ):
            return

        logger.warning(
            "[Diffusion BCG] disabled for %s: only Ideogram-4, Qwen/Qwen-Image, "
            "Qwen/Qwen-Image-2512, Tongyi-MAI/Z-Image/Z-Image-Turbo, "
            "and zai-org/GLM-Image are currently supported.",
            pipeline_config_name,
        )
        self.enable_breakable_cuda_graph = False

    def _is_breakable_cuda_graph_supported_model(self) -> bool:
        refs = _normalized_bcg_model_refs(self.model_id)
        refs.update(_normalized_bcg_model_refs(self.model_path))
        return bool(refs & BREAKABLE_CUDA_GRAPH_SUPPORTED_MODEL_IDS)

    def _adjust_save_paths(self):
        """Normalize empty-string save paths to None (disabled)."""
        if self.output_path is not None and self.output_path.strip() == "":
            self.output_path = None
        if self.input_save_path is not None and self.input_save_path.strip() == "":
            self.input_save_path = None

    def _adjust_quant_config(self):
        """
        resolve, validate and adjust quantization config

        handles only nunchaku for now
        """

        ncfg = self.nunchaku_config
        if ncfg is None or isinstance(ncfg, NunchakuConfig):
            return

        resolution = ncfg.resolve_runtime_config()
        if resolution.transformer_weights_path:
            self.transformer_weights_path = resolution.transformer_weights_path
        self.nunchaku_config = resolution.nunchaku_config

    def adjust_pipeline_config(self):
        # 1. adjust for encoder parallel folding
        tp_size = self.tp_size or 1
        dp_size = self.dp_size or 1
        sp_degree = self.sp_degree or 1
        # one replica = all its GPUs
        replica_size = (self.num_gpus or tp_size) // dp_size
        fold_world = dp_size == 1 and not self.disagg_mode and replica_size > tp_size

        if fold_world:
            mode = "world"
        elif tp_size == 1 and sp_degree > 1:
            # Preserve prior behavior for dp>1 / disaggregated SP runs.
            mode = "sp"
        else:
            return

        # Propose the fold group from the parallelism for every encoder. The
        # loader keeps it only for encoders wide enough to benefit at their real
        # (post-load) size and whose dims divide the group -- see
        # finalize_encoder_folding. Deciding on real size (not architecture)
        # handles the same encoder family at different parameter counts.
        encoder_configs = list(self.pipeline_config.text_encoder_configs) + list(
            getattr(self.pipeline_config, "image_encoder_configs", ()) or ()
        )
        for encoder_config in encoder_configs:
            encoder_config.parallel_folding_mode = mode

        logger.info(
            "Proposed encoder parallel folding (mode=%s) for %s "
            "(tp=%s sp=%s cfg=%s replica=%s); the loader keeps it for encoders "
            "wide enough to benefit.",
            mode,
            self.__class__.__name__,
            tp_size,
            sp_degree,
            self.cfg_parallel_degree or 1,
            replica_size,
        )

    def _adjust_offload(self):
        if current_platform.is_cpu():
            # CPU platform does not need offload
            return

        if self.pipeline_config.task_type.is_action_gen():
            if self.dit_cpu_offload is None:
                self.dit_cpu_offload = False
            if self.text_encoder_cpu_offload is None:
                self.text_encoder_cpu_offload = False
            if self.image_encoder_cpu_offload is None:
                self.image_encoder_cpu_offload = False
            if self.vae_cpu_offload is None:
                self.vae_cpu_offload = False
            return

        # TODO: to be handled by each platform
        if current_platform.get_device_total_memory() / BYTES_PER_GB < 30:
            logger.info(
                "Enabling large component offloading for GPU with low device memory"
            )
            if self.dit_cpu_offload is None:
                self.dit_cpu_offload = True
            if self.text_encoder_cpu_offload is None:
                self.text_encoder_cpu_offload = True
            if self.image_encoder_cpu_offload is None:
                self.image_encoder_cpu_offload = True
        elif self.pipeline_config.task_type.is_image_gen():
            if self.dit_cpu_offload is None:
                self.dit_cpu_offload = True
            if self.text_encoder_cpu_offload is None:
                self.text_encoder_cpu_offload = True
            if self.image_encoder_cpu_offload is None:
                self.image_encoder_cpu_offload = False
        else:
            if self.dit_cpu_offload is None:
                self.dit_cpu_offload = True
            if self.text_encoder_cpu_offload is None:
                self.text_encoder_cpu_offload = True
            if self.image_encoder_cpu_offload is None:
                self.image_encoder_cpu_offload = True

    def _adjust_ltx2_two_stage_device_mode(self):
        if not self._is_ltx23_two_stage_pipeline():
            return

        mode = self.ltx2_two_stage_device_mode
        if mode is None:
            env_mode = os.getenv("SGLANG_LTX2_TWO_STAGE_DEVICE_MODE")
            mode = (
                _normalize_ltx2_two_stage_device_mode(env_mode)
                if env_mode
                else self._resolve_default_ltx2_two_stage_device_mode()
            )
        else:
            mode = _normalize_ltx2_two_stage_device_mode(mode)

        if mode not in LTX2_TWO_STAGE_DEVICE_MODES:
            raise ValueError(
                f"Invalid ltx2_two_stage_device_mode={mode!r}. "
                f"Expected one of {LTX2_TWO_STAGE_DEVICE_MODE_CHOICES}."
            )

        self.ltx2_two_stage_device_mode = mode

    def _resolve_default_ltx2_two_stage_device_mode(self) -> str:
        if not current_platform.is_cuda():
            logger.info(
                "Automatically set ltx2_two_stage_device_mode=original on non-CUDA platform"
            )
            return "original"

        device_name = str(current_platform.get_device_name(0)).upper()
        device_total_memory_gb = (
            current_platform.get_device_total_memory() / BYTES_PER_GB
        )
        if (
            "H200" in device_name
            or device_total_memory_gb >= LTX2_RESIDENT_AUTO_ENABLE_MEM_GB
        ):
            logger.info(
                "Automatically set ltx2_two_stage_device_mode=resident for high-memory CUDA GPU (%s, %.2f GiB total)",
                device_name,
                device_total_memory_gb,
            )
            return "resident"

        logger.info(
            "Automatically set ltx2_two_stage_device_mode=original for CUDA GPU (%s, %.2f GiB total)",
            device_name,
            device_total_memory_gb,
        )
        return "original"

    def _is_ltx23_two_stage_pipeline(self) -> bool:
        return is_ltx2_two_stage_pipeline_name(self.pipeline_class_name) and (
            self._is_ltx23_model_path(self.model_path)
            or is_ltx23_native_variant(self.pipeline_config.vae_config.arch_config)
        )

    def _uses_ltx23_high_memory_resident_two_stage_mode(self) -> bool:
        if (
            self.ltx2_two_stage_device_mode != "resident"
            or not self._is_ltx23_two_stage_pipeline()
            or not current_platform.is_cuda()
        ):
            return False
        return (
            current_platform.get_device_total_memory() / BYTES_PER_GB
            >= LTX2_RESIDENT_AUTO_ENABLE_MEM_GB
        )

    def _adjust_attention_backend(self):
        if self.attention_backend in ["fa3", "fa4"]:
            self.attention_backend = "fa"
        self.component_attention_backends = (
            self._normalize_component_attention_backends(
                self.component_attention_backends
            )
        )

        # attention_backend_config
        if self.attention_backend_config is None:
            self.attention_backend_config = addict.Dict()
        elif isinstance(self.attention_backend_config, str):
            self.attention_backend_config = addict.Dict(
                self._parse_attention_backend_config(self.attention_backend_config)
            )

        if self.backend != Backend.DIFFUSERS and isinstance(
            self.pipeline_config, LTX2PipelineConfig
        ):
            text_backend = self.component_attention_backends.get("text_encoder")
            if text_backend != "torch_sdpa":
                if text_backend is None:
                    logger.info(
                        "Automatically set torch_sdpa backend for component text_encoder to preserve LTX2 official attention semantics"
                    )
                else:
                    logger.warning(
                        "Overriding %s backend with torch_sdpa for component text_encoder to preserve LTX2 official attention semantics",
                        text_backend,
                    )
                self.component_attention_backends["text_encoder"] = "torch_sdpa"

        if self.ring_degree > 1:
            if self.attention_backend is not None and self.attention_backend not in (
                "fa",
                "sage_attn",
            ):
                raise ValueError(
                    "Ring Attention is only supported for flash attention or sage attention backend for now"
                )
            if self.attention_backend is None:
                self.attention_backend = "fa"
                logger.info(
                    "Ring Attention is currently only supported for flash attention or sage attention; "
                    "attention_backend has been automatically set to flash attention"
                )

        if self.attention_backend is None and self.backend != Backend.DIFFUSERS:
            if (
                current_platform.is_cuda()
                and self.pipeline_class_name is None
                and self.num_gpus == 1
                and self.tp_size == 1
                and self.sp_degree == 1
                and self.ulysses_degree == 1
                and self.ring_degree == 1
                and self._is_ltx23_model_path(self.model_path)
            ):
                self.attention_backend = "fa"
                logger.info(
                    "Automatically set attention_backend=fa for LTX-2.3 one-stage on 1 GPU to preserve precision"
                )
                return
            self._set_default_attention_backend()

    @staticmethod
    def _normalize_attention_backend_name(backend: str) -> str:
        if not isinstance(backend, str):
            raise ValueError("Attention backend name must be a string")
        normalized = backend.strip().lower()
        if normalized in ("fa3", "fa4"):
            normalized = "fa"
        try:
            return AttentionBackendEnum[normalized.upper()].name.lower()
        except KeyError:
            raise ValueError(
                f"Invalid attention backend '{backend}'. "
                f"Available options are: {[e.name.lower() for e in AttentionBackendEnum]}"
            ) from None

    @staticmethod
    def _parse_component_attention_backend_map(
        value: dict[str, str] | str | None,
    ) -> dict[str, str]:
        if value is None or value == "":
            return {}
        if isinstance(value, dict):
            return dict(value)
        if not isinstance(value, str):
            raise ValueError(
                "component_attention_backends must be a dict or a comma-separated component=backend string"
            )

        try:
            parsed = json.loads(value)
            if not isinstance(parsed, dict):
                raise ValueError
            return parsed
        except (json.JSONDecodeError, ValueError):
            pass

        result: dict[str, str] = {}
        for pair in value.split(","):
            pair = pair.strip()
            if not pair:
                continue
            if "=" not in pair:
                raise ValueError(
                    "component_attention_backends must use component=backend entries"
                )
            component, backend = pair.split("=", 1)
            result[component.strip()] = backend.strip()
        return result

    @classmethod
    def _normalize_component_attention_backends(
        cls, value: dict[str, str] | str | None
    ) -> dict[str, str]:
        raw = cls._parse_component_attention_backend_map(value)
        normalized: dict[str, str] = {}
        for component, backend in raw.items():
            if not isinstance(component, str):
                raise ValueError("Component attention backend key must be a string")
            component_name = component.strip().replace("-", "_")
            if not component_name:
                raise ValueError("Component attention backend key must not be empty")
            normalized[component_name] = cls._normalize_attention_backend_name(backend)
        return normalized

    def resolve_component_attention_backend(
        self, *component_names: str | None
    ) -> tuple[AttentionBackendEnum | None, str | None]:
        for component_name in component_names:
            if component_name is None:
                continue
            key = component_name.replace("-", "_")
            fallback_keys = [key]
            if key.endswith("_2"):
                # Secondary two-stage components inherit the base component
                # backend unless explicitly overridden.
                fallback_keys.append(key[:-2])
            for backend_key in fallback_keys:
                backend = self.component_attention_backends.get(backend_key)
                if backend is not None:
                    return AttentionBackendEnum[backend.upper()], backend_key
        return None, None

    def _adjust_warmup(self):
        #   --warmup-mode > --warmup/--server-warmup
        mode_explicit = self.is_arg_explicitly_set("warmup_mode")
        legacy_explicit = self.is_arg_explicitly_set(
            "warmup"
        ) or self.is_arg_explicitly_set("server_warmup")
        if self.warmup_mode is not None:
            if self.warmup_mode not in WARMUP_MODES:
                raise ValueError(
                    f"Invalid --warmup-mode {self.warmup_mode!r}; "
                    f"expected one of {WARMUP_MODES}."
                )
            if mode_explicit and legacy_explicit:
                logger.warning(
                    "Both --warmup-mode and the deprecated --warmup/--server-warmup "
                    "were set; --warmup-mode=%s takes precedence.",
                    self.warmup_mode,
                )
            if mode_explicit or not legacy_explicit:
                self.warmup = self.warmup_mode != "off"
                self.server_warmup = self.warmup_mode == "server"
            elif self.warmup:
                self.server_warmup = self.server_warmup or self.warmup_mode == "server"

        # Explicit resolutions imply warmup is on (request-based).
        if self.warmup_resolutions is not None:
            self.warmup = True

        if (
            self.enable_torch_compile
            and self.warmup_mode is None
            and not mode_explicit
            and not legacy_explicit
        ):
            self.warmup = True
            self.server_warmup = True
            logger.info(
                "Automatically enabled server warmup for torch.compile so first "
                "real requests do not pay compile latency. Set --warmup-mode off "
                "to disable this behavior."
            )

        # BCG captures every graph during a synthetic warmup forward at startup
        # so that serving never records a fresh graph. That requires
        # server-based warmup (a real warmup request issued at startup), not
        # request-based warmup which runs no forward until the first request.
        if self.enable_breakable_cuda_graph and self.disagg_role == RoleType.MONOLITHIC:
            self.warmup = True
            self.server_warmup = True

        if self.disagg_role != RoleType.MONOLITHIC:
            self.server_warmup = False

        if not self.warmup:
            self.server_warmup = False

        self.warmup_mode = (
            "off" if not self.warmup else "server" if self.server_warmup else "request"
        )

    @staticmethod
    def _require_port(port: int, name: str) -> None:
        """Raise if *port* is occupied (used under ``--strict-ports``)."""
        if not is_port_available(port):
            raise RuntimeError(
                f"{name} port {port} is unavailable and --strict-ports is enabled. "
                f"Either use a different port or disable --strict-ports."
            )

    def _adjust_network_ports(self):
        # Disagg role instances (encoder/denoiser/decoder) don't serve HTTP,
        # so skip settling the HTTP port to avoid unnecessary port collisions.
        needs_http = self.disagg_role in (
            RoleType.MONOLITHIC,
            RoleType.SERVER,
        )

        if self.strict_ports:
            requested_ports = []
            if needs_http:
                requested_ports.append((self.port, "HTTP"))
            requested_ports.append((self.scheduler_port, "Scheduler"))
            if self.master_port is not None:
                requested_ports.append((self.master_port, "Master"))
            seen_ports: dict[int, str] = {}
            for port, name in requested_ports:
                if port in seen_ports:
                    raise RuntimeError(
                        f"{name} port {port} duplicates {seen_ports[port]} port and "
                        "--strict-ports is enabled."
                    )
                seen_ports[port] = name
                self._require_port(port, name)
        else:
            settled_ports: set[int] = set()
            if needs_http:
                self.port = self.settle_port(self.port)
                settled_ports.add(self.port)
            initial_scheduler_port = self.scheduler_port + (
                random.randint(0, 100) if self.scheduler_port == 5555 else 0
            )
            self.scheduler_port = self.settle_port(
                initial_scheduler_port, avoid=settled_ports
            )
            settled_ports.add(self.scheduler_port)
            if self.master_port is not None:
                self.master_port = self.settle_port(
                    self.master_port, 37, avoid=settled_ports
                )

    def _adjust_parallelism(self):
        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 self.tp_size is None:
            self.tp_size = 1

        if current_platform.is_cpu() and self.tp_size > 1:
            # CPU platform reuse num_gpus to represent num cpu numa nodes as devices
            self.num_gpus = self.tp_size

        if self.hsdp_shard_dim is None:
            self.hsdp_shard_dim = self.num_gpus

        # --cfg-parallel-size takes precedence over --enable-cfg-parallel bool.
        if self.cfg_parallel_degree is not None:
            if self.cfg_parallel_degree == 1:
                self.enable_cfg_parallel = False
            elif self.cfg_parallel_degree > 1:
                self.enable_cfg_parallel = True
            cfg_unspecified = False

        # Auto-enable CFG parallel when user hasn't set any parallelism flags
        # and there are enough GPUs.  Only auto-enable for models whose default
        # SamplingParams use classifier-free guidance (negative_prompt is not None),
        # because non-CFG models (e.g. FLUX) crash when CFG parallel splits ranks.
        if cfg_unspecified:
            deployment_config = self.pipeline_config.get_model_deployment_config()
            auto_cfg_parallel_degree = deployment_config.get_auto_cfg_parallel_degree(
                self.num_gpus
            )
            if auto_cfg_parallel_degree < 1:
                self.enable_cfg_parallel = False
            else:
                cfg_group_size = self.dp_size * self.tp_size * auto_cfg_parallel_degree
                if (
                    self.performance_mode != "manual"
                    and deployment_config.auto_enable_cfg_parallel
                    and self.num_gpus >= 2
                    and self.num_gpus % cfg_group_size == 0
                    and sp_unspecified
                    and ulysses_unspecified
                    and ring_unspecified
                    and self._model_default_uses_cfg()
                ):
                    self.cfg_parallel_degree = auto_cfg_parallel_degree
                    self.enable_cfg_parallel = auto_cfg_parallel_degree > 1
                    if self.enable_cfg_parallel:
                        logger.info(
                            "Automatically enabled CFG parallel at degree %d for %d GPUs. "
                            "Use --sp-degree / --ulysses-degree to use sequence "
                            "parallelism instead.",
                            self.cfg_parallel_degree,
                            self.num_gpus,
                        )
                    else:
                        logger.info(
                            "Automatically disabled CFG parallel for %d GPUs based on model deployment config.",
                            self.num_gpus,
                        )
                else:
                    self.enable_cfg_parallel = False

        # Resolve cfg_parallel_degree to a concrete int now that enable_cfg_parallel is settled.
        if self.cfg_parallel_degree is None:
            self.cfg_parallel_degree = 2 if self.enable_cfg_parallel else 1

        # adjust sp_degree: allocate all remaining GPUs after TP and DP
        if self.sp_degree is None:
            num_gpus_per_group = self.dp_size * self.tp_size
            if self.enable_cfg_parallel:
                num_gpus_per_group *= self.cfg_parallel_degree
            if self.num_gpus % num_gpus_per_group == 0:
                self.sp_degree = self.num_gpus // num_gpus_per_group
            else:
                # Will be validated later
                self.sp_degree = 1

        if (
            self.ulysses_degree is None
            and self.ring_degree is None
            and self.sp_degree != 1
        ):
            self.ulysses_degree = self.sp_degree
            logger.info(
                f"Automatically set ulysses_degree=sp_degree={self.ulysses_degree} for best performance"
            )

        if self.ulysses_degree is None:
            self.ulysses_degree = 1
            logger.debug(
                f"Ulysses degree not set, using default value {self.ulysses_degree}"
            )

        if self.ring_degree is None:
            self.ring_degree = 1
            logger.debug(f"Ring degree not set, using default value {self.ring_degree}")

    def _model_default_uses_cfg(self) -> bool:
        """
        Check whether the model uses classifier-free guidance by default.

        CFG is active when *both* ``negative_prompt is not None`` and ``guidance_scale > 1``.
        """
        from sglang.multimodal_gen.registry import get_model_info

        model_info = get_model_info(self.model_path, self.backend, self.model_id)
        if model_info is None:
            return False
        default_params = model_info.sampling_param_cls()

        return (
            getattr(default_params, "negative_prompt", None) is not None
            and getattr(default_params, "guidance_scale", 0) > 1.0
        )

    @staticmethod
    def _is_ltx23_model_path(model_path: str | None) -> bool:
        if not model_path:
            return False
        normalized = model_path.lower()
        return any(
            token in normalized
            for token in (
                "lightricks/ltx-2.3",
                "models--lightricks--ltx-2.3",
                "lightricks__ltx-2.3",
            )
        )

    def _adjust_platform_specific(self):
        if current_platform.is_mps():
            self.use_fsdp_inference = False
            self.dit_layerwise_offload = False
            self.layerwise_offload_components = None
            if (
                self.dit_cpu_offload
                or self.text_encoder_cpu_offload
                or self.image_encoder_cpu_offload
                or self.vae_cpu_offload
            ):
                logger.warning(
                    "Disabling component CPU offload on MPS because CPU-to-MPS "
                    "module relocation can produce invalid diffusion outputs."
                )
            self.dit_cpu_offload = False
            self.text_encoder_cpu_offload = False
            self.image_encoder_cpu_offload = False
            self.vae_cpu_offload = False

    def is_arg_explicitly_set(self, arg_name: str) -> bool:
        return arg_name in self._explicit_arg_names

    def should_configure_layerwise_offload_for_lazy_component(
        self, component_name: str
    ) -> bool:
        """Return whether a lazy-loaded component should try layerwise offload.

        Lazy components are loaded after the normal pipeline-wide configuration
        pass, so they should only attempt layerwise configuration when their
        component name is covered by the selected layerwise scope.
        """
        component_names = normalize_layerwise_offload_components(
            self.layerwise_offload_components
        )
        if not component_names:
            return False
        if LAYERWISE_OFFLOAD_ALL_COMPONENTS in component_names:
            return True
        return layerwise_component_matches_any_selection(
            component_name, component_names
        )

    @property
    def is_dit_layerwise_offload_selected(self) -> bool:
        """returns if dit is selected to be layerwise-offload"""
        component_names = self.layerwise_offload_components
        return bool(
            component_names
            and "dit_cpu_offload"
            in cpu_offload_flags_for_layerwise_components(component_names)
        )

    def _adjust_layerwise_offload_components(self):
        explicitly_set_component_names = normalize_layerwise_offload_components(
            self.layerwise_offload_components
        )
        if self.dit_layerwise_offload:
            if explicitly_set_component_names is None:
                explicitly_set_component_names = [LAYERWISE_OFFLOAD_DIT_GROUP]
            elif LAYERWISE_OFFLOAD_DIT_GROUP not in explicitly_set_component_names:
                explicitly_set_component_names = [
                    LAYERWISE_OFFLOAD_DIT_GROUP,
                    *explicitly_set_component_names,
                ]

        if explicitly_set_component_names is not None:
            self.layerwise_offload_components = explicitly_set_component_names
            self._disable_non_dit_cpu_offload_for_layerwise_components(
                explicitly_set_component_names
            )
            return

    def _disable_non_dit_cpu_offload_for_layerwise_components(
        self, component_names: list[str]
    ) -> None:
        # non-DiT layerwise offload replaces the corresponding component-level CPU offload
        flag_names = cpu_offload_flags_for_layerwise_components(component_names)
        disabled_flag_names: list[str] = []

        if (
            "text_encoder_cpu_offload" in flag_names
            and self.text_encoder_cpu_offload is not False
        ):
            self.text_encoder_cpu_offload = False
            disabled_flag_names.append("text_encoder_cpu_offload")
        if (
            "image_encoder_cpu_offload" in flag_names
            and self.image_encoder_cpu_offload is not False
        ):
            self.image_encoder_cpu_offload = False
            disabled_flag_names.append("image_encoder_cpu_offload")
        if "vae_cpu_offload" in flag_names and self.vae_cpu_offload is not False:
            self.vae_cpu_offload = False
            disabled_flag_names.append("vae_cpu_offload")

        explicit_disabled_flag_names = [
            flag_name
            for flag_name in disabled_flag_names
            if self.is_arg_explicitly_set(flag_name)
        ]
        if explicit_disabled_flag_names:
            logger.info(
                "Ignoring explicit CPU-offload flags because layerwise offload "
                "manages the same component weights: %s",
                ", ".join(
                    f"{flag_name}=False" for flag_name in explicit_disabled_flag_names
                ),
            )

    def _adjust_autocast(self):
        if self.disable_autocast is None:
            self.disable_autocast = not self.pipeline_config.enable_autocast

    def _parse_attention_backend_config(self, config_str: str) -> dict[str, Any]:
        """parse attention backend config from string."""
        if not config_str:
            return {}

        # 1. treat as file path
        if os.path.exists(config_str):
            if config_str.endswith((".yaml", ".yml")):
                with open(config_str, "r") as f:
                    return yaml.safe_load(f)
            elif config_str.endswith(".json"):
                with open(config_str, "r") as f:
                    return json.load(f)

        # 2. treat as JSON string
        try:
            return json.loads(config_str)
        except json.JSONDecodeError:
            pass

        # 3. treat as k=v pairs (simple implementation). e.g., "sparsity=0.5,enable_x=true"
        try:
            config = {}
            pairs = config_str.split(",")
            for pair in pairs:
                k, v = pair.split("=", 1)
                k = k.strip()
                v = v.strip()
                if v.lower() == "true":
                    v = True
                elif v.lower() == "false":
                    v = False
                elif v.replace(".", "", 1).isdigit():
                    v = float(v) if "." in v else int(v)
                config[k] = v
            return config
        except Exception:
            raise ValueError(f"Could not parse attention backend config: {config_str}")

    def __post_init__(self):
        # configure logger before use
        configure_logger(server_args=self)

        # Convert string disagg_role to enum (from CLI/config)
        if isinstance(self.disagg_role, str):
            self.disagg_role = RoleType.from_string(self.disagg_role)
        self.gpu_ids = normalize_gpu_ids(self.gpu_ids)

        # 1. adjust parameters
        self._adjust_parameters()

        # 2. Validate parameters
        self._validate_parameters()

        # log clean server_args
        try:
            safe_args = _sanitize_for_logging(self, key_hint="server_args")
            logger.info("server_args: %s", json.dumps(safe_args, ensure_ascii=False))
        except Exception:
            # Fallback to default repr if sanitization fails
            logger.info(f"server_args: {self}")

    @staticmethod
    def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
        # Model and path configuration
        parser.add_argument(
            "--model-path",
            type=str,
            help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
        )
        parser.add_argument(
            "--model-id",
            type=str,
            default=ServerArgs.model_id,
            help=(
                "Override the model ID used for config resolution. "
                "Useful when --model-path is a local directory whose name does not match "
                "any registered HF repo name. Should be the repo name portion of the HF ID "
                "(e.g. 'Qwen-Image' for 'Qwen/Qwen-Image')."
            ),
        )
        parser.add_argument(
            "--pipeline",
            "--pipeline-class-name",
            dest="pipeline_class_name",
            type=str,
            default=ServerArgs.pipeline_class_name,
            help=(
                "Advanced override for pipeline class selection from the model registry "
                "or model_index.json. Must match a registered pipeline_name."
            ),
        )
        # attention
        parser.add_argument(
            "--attention-backend",
            type=str,
            default=None,
            help=(
                "The attention backend to use. For SGLang-native pipelines, use "
                "values like fa, torch_sdpa, sage_attn, etc. For diffusers pipelines, "
                "use diffusers attention backend names such as flash, _flash_3_hub, "
                "sage, or xformers."
            ),
        )
        parser.add_argument(
            "--attention-backend-config",
            type=str,
            default=None,
            help="Configuration for the attention backend. Can be a JSON string, a path to a JSON/YAML file, or key=value pairs.",
        )
        parser.add_argument(
            "--component-attention-backends",
            type=str,
            default=None,
            help=(
                "Per-component attention backend overrides for native pipelines. "
                "Use component names from model_index.json, e.g. "
                "'text_encoder=torch_sdpa,transformer=fa'."
            ),
        )
        parser.add_argument(
            "--cache-dit-config",
            type=str,
            default=ServerArgs.cache_dit_config,
            help="Path to a Cache-DiT YAML/JSON config. Enables cache-dit for diffusers backend.",
        )

        # HuggingFace specific parameters
        parser.add_argument(
            "--trust-remote-code",
            action=StoreBoolean,
            default=ServerArgs.trust_remote_code,
            help="Trust remote code when loading HuggingFace models",
        )
        parser.add_argument(
            "--revision",
            type=str,
            default=ServerArgs.revision,
            help="The specific model version to use (can be a branch name, tag name, or commit id)",
        )

        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, no adjustment is made; "
                "'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."
            ),
        )

        # Parallelism
        parser.add_argument(
            "--num-gpus",
            type=int,
            default=ServerArgs.num_gpus,
            help="The number of GPUs to use.",
        )
        parser.add_argument(
            "--base-gpu-id",
            type=int,
            default=ServerArgs.base_gpu_id,
            help="The starting GPU ID for this instance. Used with --disagg-role "
            "to place role instances on specific GPUs without CUDA_VISIBLE_DEVICES.",
        )
        parser.add_argument(
            "--gpu-ids",
            nargs="+",
            default=None,
            help=(
                "Physical GPU IDs for this instance, e.g. --gpu-ids 0 1 6 7 "
                "or --gpu-ids 0,1,6,7. Overrides --base-gpu-id for standalone "
                "disagg roles."
            ),
        )
        parser.add_argument(
            "--tp-size",
            type=int,
            default=None,
            help="The tensor parallelism size. Defaults to 1 if not specified.",
        )
        parser.add_argument(
            "--sp-degree",
            type=int,
            default=None,
            help="The sequence parallelism size. If not specified, will use all remaining GPUs after accounting for TP and DP.",
        )
        parser.add_argument(
            "--ulysses-degree",
            type=int,
            default=ServerArgs.ulysses_degree,
            help="Ulysses sequence parallel degree. Used in attention layer.",
        )
        parser.add_argument(
            "--ring-degree",
            type=int,
            default=ServerArgs.ring_degree,
            help="Ring sequence parallel degree. Used in attention layer.",
        )
        parser.add_argument(
            "--enable-cfg-parallel",
            action=StoreBoolean,
            default=None,
            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",
            dest="cfg_parallel_degree",
            type=int,
            default=None,
            help=(
                "Number of GPUs per CFG parallel group (1 = disabled, N > 1 = enabled at degree N). "
                "Supersedes --enable-cfg-parallel. Allows 4-branch CFG parallel (e.g., --cfg-parallel-size 4) "
                "for models with cond + neg + perturbed + modality branches."
            ),
        )
        parser.add_argument(
            "--data-parallel-size",
            "--dp-size",
            "--dp",
            dest="dp_size",
            type=int,
            default=ServerArgs.dp_size,
            help="The data parallelism size.",
        )

        parser.add_argument(
            "--hsdp-replicate-dim",
            type=int,
            default=ServerArgs.hsdp_replicate_dim,
            help="The data parallelism size.",
        )
        parser.add_argument(
            "--hsdp-shard-dim",
            type=int,
            default=None,
            help="The data parallelism shards. Defaults to num_gpus if not specified.",
        )
        parser.add_argument(
            "--dist-timeout",
            type=int,
            default=ServerArgs.dist_timeout,
            help="Timeout for torch.distributed operations in seconds. "
            "Increase this value if you encounter 'Connection closed by peer' errors after the service is idle. ",
        )

        ServerArgs.add_disagg_cli_args(parser)

        # Prompt text file for batch processing
        parser.add_argument(
            "--prompt-file-path",
            type=str,
            default=ServerArgs.prompt_file_path,
            help="Path to a text file containing prompts (one per line) for batch processing",
        )

        parser.add_argument(
            "--mask-strategy-file-path",
            type=str,
            help="Path to mask strategy JSON file for STA",
        )
        parser.add_argument(
            "--enable-torch-compile",
            action=StoreBoolean,
            default=ServerArgs.enable_torch_compile,
            help="Use torch.compile to speed up diffusion hot paths. "
            + "When no warmup mode is configured, this enables server warmup "
            + "so first real requests do not pay compile latency. "
            + "However, will likely cause precision drifts. See (https://github.com/pytorch/pytorch/issues/145213)",
        )
        parser.add_argument(
            "--offload-during-compile",
            action=StoreBoolean,
            default=ServerArgs.offload_during_compile,
            help="Offload components during the torch.compile warmup (the DiT layerwise) so max-autotune fits on tighter-memory GPUs, then restore the configured residency for serving. Skipped when the DiT is already layerwise-offloaded, or under cache-dit / FSDP.",
        )
        parser.add_argument(
            "--enable-breakable-cuda-graph",
            action=StoreBoolean,
            default=ServerArgs.enable_breakable_cuda_graph,
            help="Capture the DiT forward as breakable CUDA graph segments "
            "(split at attention; SP all-to-all / dynamic attention stay "
            "eager) to cut per-kernel launch overhead. Mutually exclusive "
            "with --enable-torch-compile and Cache-DiT (BCG takes priority). "
            "Requires --warmup-resolutions; all of them are captured at warmup.",
        )
        parser.add_argument(
            "--bcg-text-buckets",
            type=int,
            nargs="+",
            default=ServerArgs.bcg_text_buckets,
            help="Prompt sequence-length padding budget for breakable CUDA "
            "graph. Prompt-conditioning is padded up to the smallest bucket "
            "that fits so different prompt lengths reuse one captured graph; "
            "warmup captures one graph per bucket. Defaults to "
            f"{' '.join(map(str, DEFAULT_BCG_TEXT_BUCKETS))}. "
            "Replaces the legacy SGLANG_BCG_TEXT_BUCKETS env var.",
        )

        parser.add_argument(
            "--enable-layerwise-nvtx-marker",
            action=StoreBoolean,
            default=ServerArgs.enable_layerwise_nvtx_marker,
            help="Enable layerwise NVTX markers for profiling with Nsight Systems. "
            "Adds NVTX ranges around each pipeline stage, the denoising loop, "
            "every denoising step, the predict_noise / scheduler_step "
            "sub-operations, and every transformer submodule forward (recursive). "
            "Warmup steps are excluded to keep captured traces clean.",
        )

        # warmup
        parser.add_argument(
            "--warmup-mode",
            type=str,
            choices=list(WARMUP_MODES),
            default=ServerArgs.warmup_mode,
            help=(
                "Warmup mode (canonical knob). One of: "
                "`off` (no warmup); `request` (request-based: warm on real "
                "incoming requests); `server` (server-based: a synthetic warmup "
                "request right after the server is ready, before traffic). "
                "Takes precedence over the deprecated --warmup/--server-warmup. "
                "`sglang serve` defaults to `server`; other entrypoints default "
                "to request-based when warmup is enabled. When enabled, look for "
                "the line ending with `(with warmup excluded)` for actual "
                "processing time."
            ),
        )
        parser.add_argument(
            "--warmup",
            action=StoreBoolean,
            default=ServerArgs.warmup,
            help=(
                "[DEPRECATED: use --warmup-mode] Perform warmup before normal "
                "traffic. Maps to --warmup-mode request (or server, combined "
                "with --server-warmup). Recommended when benchmarking for fair "
                "comparison and best performance."
            ),
        )
        parser.add_argument(
            "--warmup-resolutions",
            type=str,
            nargs="+",
            default=ServerArgs.warmup_resolutions,
            help="Specify explicit warmup resolutions. e.g., `--warmup-resolutions 256x256 720x720`",
        )
        parser.add_argument(
            "--warmup-steps",
            type=int,
            default=ServerArgs.warmup_steps,
            help="The number of warmup steps to perform for each resolution.",
        )
        parser.add_argument(
            "--server-warmup",
            action=StoreBoolean,
            default=ServerArgs.server_warmup,
            help=(
                "[DEPRECATED: use --warmup-mode server] Send a synthetic warmup "
                "request after the server is ready (server-based warmup)."
            ),
        )

        # layerwise offload
        parser.add_argument(
            "--dit-cpu-offload",
            action=StoreBoolean,
            help="Use CPU offload for DiT inference. Enable if run out of memory with FSDP.",
        )
        parser.add_argument(
            "--dit-layerwise-offload",
            action=StoreBoolean,
            default=ServerArgs.dit_layerwise_offload,
            help="Enable layerwise CPU offload with async H2D prefetch overlap for DiTs. "
            "It selects only the DiT layerwise group. Cannot be used together with cache-dit "
            "(SGLANG_CACHE_DIT_ENABLED) or use_fsdp_inference. May be combined with "
            "--dit-cpu-offload, in which case DiT weights stay on host memory and only the "
            "layers needed for the current step are brought on-device (lowest peak GPU memory).",
        )
        parser.add_argument(
            "--layerwise-offload-components",
            "--layerwise-offload-modules",
            type=str,
            nargs="+",
            default=ServerArgs.layerwise_offload_components,
            help="Select pipeline components for layerwise offload. "
            "Use dit to select the DiT layerwise group, default for the default group "
            "(currently text_encoder, image_encoder, and vae), "
            "or all to select every layerwise-offloadable component. "
            "This option does not imply --dit-layerwise-offload. Example: "
            "--layerwise-offload-components text_encoder image_encoder vae.",
        )
        parser.add_argument(
            "--dit-offload-prefetch-size",
            type=float,
            default=ServerArgs.dit_offload_prefetch_size,
            help="The size of prefetch for dit-layerwise-offload. If the value is between 0.0 and 1.0, it is treated as a ratio of the total number of layers. If the value is >= 1, it is treated as the absolute number of layers. 0.0 means prefetch 1 layer (lowest memory). Values above 0.5 might have peak memory close to no offload but worse performance.",
        )

        # offload flags
        parser.add_argument(
            "--text-encoder-cpu-offload",
            action=StoreBoolean,
            help="Use CPU offload for text encoder. Enable if run out of memory.",
        )
        parser.add_argument(
            "--image-encoder-cpu-offload",
            action=StoreBoolean,
            help="Use CPU offload for image encoder. Enable if run out of memory.",
        )
        parser.add_argument(
            "--vae-cpu-offload",
            action=StoreBoolean,
            help="Use CPU offload for VAE. Enable if run out of memory.",
        )

        parser.add_argument(
            "--use-fsdp-inference",
            action=StoreBoolean,
            help="Use FSDP inference to shard DiT weights across GPUs. For single-GPU memory pressure, prefer CPU or layerwise offload.",
        )
        parser.add_argument(
            "--pin-cpu-memory",
            action=StoreBoolean,
            help='Pin memory for CPU offload. Only added as a temp workaround if it throws "CUDA error: invalid argument". '
            "Should be enabled in almost all cases",
        )
        parser.add_argument(
            "--ltx2-two-stage-device-mode",
            type=str,
            choices=LTX2_TWO_STAGE_DEVICE_MODE_CHOICES,
            default=ServerArgs.ltx2_two_stage_device_mode,
            help=(
                "LTX-2.3 two-stage device residency mode: "
                "'original' keeps official two-stage semantics without premerged stage2, "
                "'resident' keeps both transformers resident on GPU. "
                "'snapshot' is deprecated, treated as 'original', and may be "
                "removed after two release cycles. "
                "Default is auto: resident on H200/high-memory CUDA GPUs, otherwise original."
            ),
        )
        parser.add_argument(
            "--disable-autocast",
            action=StoreBoolean,
            help="Disable autocast for denoising loop and vae decoding in pipeline sampling",
        )

        # quantization
        parser.add_argument(
            "--quantization",
            type=str,
            default=ServerArgs.quantization,
            help=(
                "Quantization method for the transformer. If omitted, the method is "
                "auto-detected from the checkpoint config or safetensors metadata when "
                "possible. Use this flag to override auto-detection. "
                "Online (post-load) quantization from a BF16/FP16 checkpoint "
                "is supported for 'fp8' and 'mxfp4'. Other methods "
                "('modelopt', 'modelopt_fp8', 'modelopt_fp4', 'mxfp8', "
                "'mxfp4_npu', 'modelslim') require a pre-quantized checkpoint. "
                "Note: 'mxfp4' targets ROCm + MI350+ (gfx95x); "
                "'mxfp4_npu' / 'mxfp8' target Ascend NPU (A5 series for mxfp4_npu)."
            ),
        )
        parser.add_argument(
            "--quantization-ignored-layers",
            type=str,
            nargs="+",
            default=ServerArgs.quantization_ignored_layers,
            help=(
                "Layer name patterns to keep unquantized during online quantization "
                "(fp8/mxfp4). Each pattern is matched against the layer prefix. "
                "Example: --quantization-ignored-layers img_mod txt_mod to_out"
            ),
        )

        # Nunchaku SVDQuant quantization parameters
        NunchakuSVDQuantArgs.add_cli_args(parser)

        # Master port for distributed inference
        parser.add_argument(
            "--master-port",
            type=int,
            default=ServerArgs.master_port,
            help="Master port for distributed inference. If not set, a random free port will be used.",
        )
        parser.add_argument(
            "--scheduler-port",
            type=int,
            default=ServerArgs.scheduler_port,
            help="Port for the scheduler server.",
        )
        parser.add_argument(
            "--batching-mode",
            type=str,
            default=ServerArgs.batching_mode,
            choices=["dynamic"],
            help="Request batching scheduler mode. Currently only 'dynamic' is implemented.",
        )
        parser.add_argument(
            "--batching-max-size",
            type=int,
            default=ServerArgs.batching_max_size,
            help="Maximum number of compatible generation requests to merge into one batch.",
        )
        parser.add_argument(
            "--batching-delay-ms",
            type=float,
            default=ServerArgs.batching_delay_ms,
            help="Maximum time (in ms) to wait for forming a larger batch before dispatch.",
        )
        parser.add_argument(
            "--batching-config",
            type=str,
            default=ServerArgs.batching_config,
            help=(
                "Optional JSON file with {'schema_version': 1, 'rules': [...]} "
                "batching admission rules that can cap model/resolution shapes "
                "below --batching-max-size."
            ),
        )
        parser.add_argument(
            "--enable-batching-metrics",
            action="store_true",
            default=ServerArgs.enable_batching_metrics,
            help="Log periodic batch efficiency metrics such as realized batch size and queue wait time.",
        )
        parser.add_argument(
            "--host",
            type=str,
            default=ServerArgs.host,
            help="Host for the HTTP API server.",
        )
        parser.add_argument(
            "--port",
            type=int,
            default=ServerArgs.port,
            help="Port for the HTTP API server.",
        )
        parser.add_argument(
            "--strict-ports",
            action=StoreBoolean,
            default=ServerArgs.strict_ports,
            help="If enabled, fail when requested ports are unavailable instead of auto-selecting.",
        )
        parser.add_argument(
            "--webui",
            action=StoreBoolean,
            default=ServerArgs.webui,
            help="Whether to use webui for better display",
        )

        parser.add_argument(
            "--webui-port",
            type=int,
            default=ServerArgs.webui_port,
            help="Whether to use webui for better display",
        )
        parser.add_argument(
            "--output-path",
            type=str,
            default=ServerArgs.output_path,
            help='Directory path to save generated images/videos. Set to "" to disable persistent saving.',
        )
        parser.add_argument(
            "--input-save-path",
            type=str,
            default=ServerArgs.input_save_path,
            help='Directory path to save uploaded input images/videos. Set to "" to disable persistent saving.',
        )

        # LoRA
        parser.add_argument(
            "--lora-path",
            type=str,
            default=ServerArgs.lora_path,
            help="The path to the LoRA adapter weights (can be local file path or HF hub id) to launch with",
        )
        parser.add_argument(
            "--lora-nickname",
            type=str,
            default=ServerArgs.lora_nickname,
            help="The nickname for the LoRA adapter to launch with",
        )
        parser.add_argument(
            "--lora-scale",
            type=float,
            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,
            default=ServerArgs.lora_weight_name,
            help="Specific safetensors filename to load from a multi-file LoRA repo",
        )
        # Add pipeline configuration arguments
        PipelineConfig.add_cli_args(parser)

        # Logging
        parser.add_argument(
            "--log-level",
            type=str,
            default=ServerArgs.log_level,
            help="The logging level of all loggers.",
        )

        # Tracing
        parser.add_argument(
            "--enable-trace",
            action="store_true",
            default=False,
            help="Enable OpenTelemetry tracing.",
        )
        parser.add_argument(
            "--otlp-traces-endpoint",
            type=str,
            default=ServerArgs.otlp_traces_endpoint,
            help="OTLP collector endpoint when --enable-trace is set. Format: <host>:<port>",
        )
        parser.add_argument(
            "--log-requests",
            action="store_true",
            help="Log user-facing fields of all requests (default: False). "
            "Verbosity is controlled by --log-requests-level.",
        )
        parser.add_argument(
            "--log-requests-level",
            type=int,
            default=ServerArgs.log_requests_level,
            choices=[0, 1, 2, 3],
            help="Verbosity level for request logging. "
            "0: Log request metadata only (request_id). "
            "1: Log metadata + sampling config (seed, steps, guidance, resolution, frames, fps, ...). "
            "2: Log metadata + sampling config + prompt/negative prompt (truncated to 2 KiB). "
            "3: Log metadata + sampling config + full prompt/negative prompt.",
        )
        parser.add_argument(
            "--log-requests-format",
            type=str,
            default=ServerArgs.log_requests_format,
            choices=["text", "json"],
            help="Format for request logging: 'text' (human-readable) or 'json' (structured)",
        )
        parser.add_argument(
            "--log-requests-target",
            type=str,
            nargs="+",
            default=ServerArgs.log_requests_target,
            help="Target(s) for request logging: 'stdout' and/or directory path(s) for file output. "
            "Can specify multiple targets, e.g., '--log-requests-target stdout /my/path'. ",
        )
        parser.add_argument(
            "--uvicorn-access-log-exclude-prefixes",
            type=str,
            nargs="*",
            default=[],
            help="Exclude uvicorn access logs whose request path starts with any of these prefixes. "
            "Defaults to empty (disabled). "
            "Example: --uvicorn-access-log-exclude-prefixes /metrics /health",
        )
        parser.add_argument(
            "--backend",
            type=str,
            choices=Backend.choices(),
            default=ServerArgs.backend.value,
            help="The model backend to use. 'auto' prefers sglang native and falls back to diffusers. "
            "'sglang' uses native optimized implementation. 'diffusers' uses vanilla diffusers pipeline.",
        )

        # SGLang backend for encoder stage
        parser.add_argument(
            "--srt-encoder-url",
            type=str,
            default=ServerArgs.srt_encoder_url,
            help="Url of SGLang server for encoder stage",
        )
        parser.add_argument(
            "--srt-encoder-connection-timeout",
            type=int,
            default=ServerArgs.srt_encoder_connect_timeout,
            help="Timeout (in seconds) for establishing the initial TCP connection to the SGLang encoder server. "
            "Default value is 3.05.",
        )
        parser.add_argument(
            "--srt-encoder-timeout",
            type=int,
            default=ServerArgs.srt_encoder_timeout,
            help="Timeout (in seconds) for HTTP requests to the SGLang encoder server. "
            "Increase value if connection between diffusion server and AR model server is slow.",
        )

        return parser

    def url(self):
        host = self.host
        if not host or host == "0.0.0.0":
            host = "127.0.0.1"
        elif host == "::":
            host = "::1"
        if is_valid_ipv6_address(host):
            return f"http://[{host}]:{self.port}"
        else:
            return f"http://{host}:{self.port}"

    @property
    def scheduler_endpoint(self):
        """
        Internal endpoint for scheduler.
        Prefers the configured host but normalizes localhost -> 127.0.0.1 to avoid ZMQ issues.
        """
        scheduler_host = self.host
        if scheduler_host is None or scheduler_host == "localhost":
            scheduler_host = "127.0.0.1"
        return f"tcp://{scheduler_host}:{self.scheduler_port}"

    def settle_port(
        self,
        port: int,
        port_inc: int = 42,
        max_attempts: int = 100,
        avoid: set[int] | None = None,
    ) -> int:
        """
        Find an available port with retry logic.
        """
        attempts = 0
        original_port = port
        avoid = avoid or set()

        while attempts < max_attempts:
            if port not in avoid and is_port_available(port):
                if attempts > 0:
                    logger.info(
                        f"Port {original_port} was unavailable, using port {port} instead"
                    )
                return port

            attempts += 1
            if port < 60000:
                port += port_inc
            else:
                # Wrap around with randomization to avoid collision
                port = 5000 + random.randint(0, 1000)

        raise RuntimeError(
            f"Failed to find available port after {max_attempts} attempts "
            f"(started from port {original_port})"
        )

    @staticmethod
    def _extract_component_paths(
        unknown_args: list[str],
    ) -> tuple[dict[str, str], list[str]]:
        """
        Extract dynamic component path args from unrecognised CLI args.

        Supported forms:
        - ``--<component>-path /path/to/component``
        - ``--component-paths.<component> /path/to/component`` (expanded from config)
        """
        component_paths: dict[str, str] = {}
        remaining: list[str] = []
        i = 0
        while i < len(unknown_args):
            arg = unknown_args[i]
            key_part = arg.split("=", 1)[0] if "=" in arg else arg
            component = None
            if key_part.startswith("--component-paths."):
                component = key_part[len("--component-paths.") :].replace("-", "_")
            elif key_part.startswith("--component_paths."):
                component = key_part[len("--component_paths.") :].replace("-", "_")
            elif key_part.startswith("--") and key_part.endswith("-path"):
                component = key_part[2:-5].replace("-", "_")

            if component is not None:
                if "=" in arg:
                    component_paths[component] = arg.split("=", 1)[1]
                elif i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith(
                    "-"
                ):
                    i += 1
                    component_paths[component] = unknown_args[i]
                else:
                    remaining.append(arg)
                    i += 1
                    continue
            else:
                remaining.append(arg)
            i += 1

        # canonicalize and validate
        for component, path in component_paths.items():
            path = os.path.expanduser(path)
            component_paths[component] = path
        return component_paths, remaining

    @staticmethod
    def _extract_component_attention_backends(
        unknown_args: list[str],
    ) -> tuple[dict[str, str], list[str]]:
        component_attention_backends: dict[str, str] = {}
        remaining: list[str] = []
        i = 0
        while i < len(unknown_args):
            arg = unknown_args[i]
            key_part = arg.split("=", 1)[0] if "=" in arg else arg
            component = None
            if key_part.startswith("--component-attention-backends."):
                component = key_part[len("--component-attention-backends.") :].replace(
                    "-", "_"
                )
            elif key_part.startswith("--component_attention_backends."):
                component = key_part[len("--component_attention_backends.") :].replace(
                    "-", "_"
                )

            if component is not None:
                if "=" in arg:
                    component_attention_backends[component] = arg.split("=", 1)[1]
                elif i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith(
                    "-"
                ):
                    i += 1
                    component_attention_backends[component] = unknown_args[i]
                else:
                    remaining.append(arg)
                    i += 1
                    continue
            else:
                remaining.append(arg)
            i += 1
        return component_attention_backends, remaining

    @classmethod
    def from_cli_args(
        cls,
        args: argparse.Namespace,
        unknown_args: list[str] | None = None,
        default_args: dict[str, Any] | None = None,
    ) -> "ServerArgs":
        if unknown_args is None:
            unknown_args = []

        # extract dynamic --<component>-path from unknown args
        dynamic_paths, remaining = cls._extract_component_paths(unknown_args)
        dynamic_attention_backends, remaining = (
            cls._extract_component_attention_backends(remaining)
        )
        if remaining:
            raise SystemExit(f"error: unrecognized arguments: {' '.join(remaining)}")

        provided_args = cls.get_provided_args(args, unknown_args)
        explicit_arg_names = set(provided_args)

        # Handle config file
        config_file = provided_args.get("config")
        if config_file:
            config_args = cls.load_config_file(config_file)
            explicit_arg_names.update(config_args)
            provided_args = {**config_args, **provided_args}

        if default_args:
            for key, value in default_args.items():
                provided_args.setdefault(key, value)

        if dynamic_paths:
            existing = dict(provided_args.get("component_paths") or {})
            existing.update(dynamic_paths)
            provided_args["component_paths"] = existing
            explicit_arg_names.add("component_paths")
        if dynamic_attention_backends:
            existing = cls._parse_component_attention_backend_map(
                provided_args.get("component_attention_backends")
            )
            existing.update(dynamic_attention_backends)
            provided_args["component_attention_backends"] = existing
            explicit_arg_names.add("component_attention_backends")

        provided_args["_explicit_arg_names"] = explicit_arg_names
        return cls.from_dict(provided_args)

    @classmethod
    def from_dict(cls, kwargs: dict[str, Any]) -> "ServerArgs":
        """Create a ServerArgs object from a dictionary."""
        attrs = [attr.name for attr in dataclasses.fields(cls) if attr.init]
        server_args_kwargs: dict[str, Any] = {}
        explicit_arg_names = kwargs.get("_explicit_arg_names")
        if explicit_arg_names is None:
            explicit_arg_names = set(kwargs)

        component_paths = dict(kwargs.get("component_paths") or {})
        if component_paths:
            server_args_kwargs["component_paths"] = component_paths
        server_args_kwargs["_explicit_arg_names"] = set(explicit_arg_names)

        for attr in attrs:
            if attr == "_explicit_arg_names":
                continue
            elif attr == "pipeline_config":
                pipeline_config = PipelineConfig.from_kwargs(kwargs)
                logger.debug(f"Using PipelineConfig: {type(pipeline_config)}")
                server_args_kwargs["pipeline_config"] = pipeline_config
            elif attr == "nunchaku_config":
                nunchaku_config = NunchakuSVDQuantArgs.from_dict(kwargs)
                server_args_kwargs["nunchaku_config"] = nunchaku_config
            elif attr in kwargs:
                server_args_kwargs[attr] = kwargs[attr]

        return cls(**server_args_kwargs)

    @staticmethod
    def load_config_file(config_file: str) -> dict[str, Any]:
        """Load a config file."""
        if config_file.endswith(".json"):
            with open(config_file, "r") as f:
                return json.load(f)
        elif config_file.endswith((".yaml", ".yml")):
            try:
                import yaml
            except ImportError:
                raise ImportError(
                    "Please install PyYAML to use YAML config files. "
                    "`pip install pyyaml`"
                )
            with open(config_file, "r") as f:
                return yaml.safe_load(f)
        else:
            raise ValueError(f"Unsupported config file format: {config_file}")

    @classmethod
    def from_kwargs(cls, **kwargs: Any) -> "ServerArgs":
        explicit_arg_names = set(kwargs)

        # Convert backend string to enum if necessary
        if "backend" in kwargs and isinstance(kwargs["backend"], str):
            kwargs["backend"] = Backend.from_string(kwargs["backend"])

        # Convert disagg_role string to enum if necessary
        if "disagg_role" in kwargs and isinstance(kwargs["disagg_role"], str):
            kwargs["disagg_role"] = RoleType.from_string(kwargs["disagg_role"])

        kwargs["pipeline_config"] = PipelineConfig.from_kwargs(kwargs)
        kwargs["_explicit_arg_names"] = explicit_arg_names
        return cls(**kwargs)

    @staticmethod
    def get_provided_args(
        args: argparse.Namespace, unknown_args: list[str]
    ) -> dict[str, Any]:
        """Get the arguments provided by the user."""
        provided_args = {}
        # We need to check against the raw command-line arguments to see what was
        # explicitly provided by the user, vs. what's a default value from argparse.
        raw_argv = sys.argv + unknown_args

        # Create a set of argument names that were present on the command line.
        # This handles both styles: '--arg=value' and '--arg value'.
        provided_arg_names = set(getattr(args, "_sglang_explicit_arg_names", ()))
        for arg in raw_argv:
            if arg.startswith("--"):
                # 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)
        cli_aliases = {
            "cfg_parallel_size": "cfg_parallel_degree",
            "data_parallel_size": "dp_size",
            "dp": "dp_size",
            "layerwise_offload_modules": "layerwise_offload_components",
            "mode": "performance_mode",
        }
        for alias_name, dest_name in cli_aliases.items():
            if alias_name in provided_arg_names:
                provided_arg_names.add(dest_name)

        # Populate provided_args if the argument from the namespace was on the command line.
        for k, v in vars(args).items():
            if k.startswith("_sglang_"):
                continue
            if k in provided_arg_names:
                provided_args[k] = v

        return provided_args

    def _validate_pipeline(self):
        if self.pipeline_config is None:
            raise ValueError("pipeline_config is not set in ServerArgs")

        self.pipeline_config.check_pipeline_config()

    def _validate_offload(self):
        # validate dit_offload_prefetch_size
        if self.dit_offload_prefetch_size > 1 and (
            isinstance(self.dit_offload_prefetch_size, float)
            and not self.dit_offload_prefetch_size.is_integer()
        ):
            self.dit_offload_prefetch_size = int(
                math.floor(self.dit_offload_prefetch_size)
            )
            logger.info(
                f"Invalid --dit-offload-prefetch-size value passed, truncated to: {self.dit_offload_prefetch_size}"
            )

        if 0.5 <= self.dit_offload_prefetch_size < 1.0:
            logger.info(
                "We do not recommend --dit-offload-prefetch-size to be between 0.5 and 1.0"
            )

        # validate layerwise offload conflicts
        if envs.SGLANG_CACHE_DIT_ENABLED and self.use_fsdp_inference:
            if self.is_arg_explicitly_set("use_fsdp_inference"):
                raise ValueError(
                    "FSDP inference cannot be enabled together with cache-dit. "
                    "cache-dit wraps known DiT block structures, while FSDP wraps "
                    "and shards modules before cache-dit can inspect them. "
                    "Please disable --use-fsdp-inference or disable "
                    "SGLANG_CACHE_DIT_ENABLED."
                )
            logger.warning(
                "cache-dit is enabled, automatically disabling use_fsdp_inference."
            )
            self.use_fsdp_inference = False

        if self.layerwise_offload_components:
            if self.dit_offload_prefetch_size < 0.0:
                raise ValueError("dit_offload_prefetch_size must be non-negative")

            is_dit_layerwise_offload_selected = self.is_dit_layerwise_offload_selected
            if self.use_fsdp_inference and is_dit_layerwise_offload_selected:
                logger.warning(
                    "layerwise offload is selected for DiT components, automatically disabling use_fsdp_inference."
                )
                self.use_fsdp_inference = False

            if envs.SGLANG_CACHE_DIT_ENABLED and is_dit_layerwise_offload_selected:
                raise ValueError(
                    "DiT layerwise offload cannot be enabled together with cache-dit. "
                    "cache-dit may reuse skipped blocks whose weights have been released by layerwise offload, "
                    "causing shape mismatch errors. "
                    "Please disable --dit-layerwise-offload, remove DiT from --layerwise-offload-components, "
                    "or disable SGLANG_CACHE_DIT_ENABLED."
                )

            if (
                self.performance_mode == "memory"
                or self.is_arg_explicitly_set("layerwise_offload_components")
                or self.dit_layerwise_offload
            ):
                logger.info_once(
                    "Using layerwise offload components: "
                    f"{', '.join(self.layerwise_offload_components)}. "
                    "This reduces peak GPU memory and can increase latency; use "
                    "--performance-mode speed for GPU-resident defaults when memory allows."
                )

    def _validate_parallelism(self):
        if self.sp_degree > self.num_gpus or self.num_gpus % self.sp_degree != 0:
            raise ValueError(
                f"num_gpus ({self.num_gpus}) must be >= and divisible by sp_degree ({self.sp_degree})"
            )

        if (
            self.hsdp_replicate_dim > self.num_gpus
            or self.num_gpus % self.hsdp_replicate_dim != 0
        ):
            raise ValueError(
                f"num_gpus ({self.num_gpus}) must be >= and divisible by hsdp_replicate_dim ({self.hsdp_replicate_dim})"
            )

        if (
            self.hsdp_shard_dim > self.num_gpus
            or self.num_gpus % self.hsdp_shard_dim != 0
        ):
            raise ValueError(
                f"num_gpus ({self.num_gpus}) must be >= and divisible by hsdp_shard_dim ({self.hsdp_shard_dim})"
            )

        if self.num_gpus % self.dp_size != 0:
            raise ValueError(
                f"num_gpus ({self.num_gpus}) must be divisible by dp_size ({self.dp_size})"
            )

        if self.dp_size < 1:
            raise ValueError("--dp-size must be a natural number")

        if self.dp_size > 1:
            raise ValueError("DP is not yet supported")

        num_gpus_per_group = self.dp_size * self.tp_size
        if self.enable_cfg_parallel:
            num_gpus_per_group *= self.cfg_parallel_degree

        if self.num_gpus % num_gpus_per_group != 0:
            raise ValueError(
                f"num_gpus ({self.num_gpus}) must be divisible by (dp_size * tp_size"
                f"{f' * {self.cfg_parallel_degree}' if self.enable_cfg_parallel else ''}"
                f") = {num_gpus_per_group}"
            )

        if self.sp_degree != self.ring_degree * self.ulysses_degree:
            raise ValueError(
                f"sp_degree ({self.sp_degree}) must equal ring_degree * ulysses_degree "
                f"({self.ring_degree} * {self.ulysses_degree} = {self.ring_degree * self.ulysses_degree})"
            )

        if os.getenv("SGLANG_CACHE_DIT_ENABLED", "").lower() == "true":
            has_sp = self.sp_degree > 1
            has_tp = self.tp_size > 1
            if has_sp and has_tp:
                logger.warning(
                    "cache-dit is enabled with hybrid parallelism (SP + TP). "
                    "Proceeding anyway (SGLang integration may support this mode)."
                )

    def _validate_cfg_parallel(self):
        if self.enable_cfg_parallel and self.num_gpus == 1:
            raise ValueError(
                "CFG Parallelism is enabled via `--enable-cfg-parallel`, but num_gpus == 1"
            )

    def _validate_batching(self):
        if self.batching_mode != "dynamic":
            raise ValueError("batching_mode must be one of: dynamic")
        if self.batching_max_size < 1:
            raise ValueError("batching_max_size must be >= 1")
        if self.batching_delay_ms < 0:
            raise ValueError("batching_delay_ms must be >= 0")

    def _set_default_attention_backend(self) -> None:
        """Configure ROCm defaults when users do not specify an attention backend."""
        if current_platform.is_rocm():
            default_backend = AttentionBackendEnum.AITER.name.lower()
            self.attention_backend = default_backend
            logger.info(
                "Attention backend not specified. Using '%s' by default on ROCm "
                "to match SGLang SRT defaults.",
                default_backend,
            )


@dataclasses.dataclass
class PortArgs:
    # The ipc filename for scheduler (rank 0) to receive inputs from tokenizer (zmq)
    scheduler_input_ipc_name: str

    # The port for nccl initialization (torch.dist)
    nccl_port: int

    # The ipc filename for rpc call between Engine and Scheduler
    rpc_ipc_name: str

    # The ipc filename for Scheduler to send metrics
    metrics_ipc_name: str

    # Master port for distributed inference
    master_port: int | None = None

    @staticmethod
    def from_server_args(
        server_args: ServerArgs, dp_rank: Optional[int] = None
    ) -> "PortArgs":
        if server_args.nccl_port is None:
            nccl_port = server_args.scheduler_port + random.randint(100, 1000)
            while True:
                if is_port_available(nccl_port):
                    break
                if nccl_port < 60000:
                    nccl_port += 42
                else:
                    nccl_port -= 43
        else:
            nccl_port = server_args.nccl_port

        # Normal case, use IPC within a single node
        return PortArgs(
            scheduler_input_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}",
            nccl_port=nccl_port,
            rpc_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}",
            metrics_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}",
            master_port=server_args.master_port,
        )


_global_server_args = None


def prepare_server_args(argv: list[str]) -> ServerArgs:
    """
    Prepare the inference arguments from the command line arguments.
    """
    parser = FlexibleArgumentParser()
    ServerArgs.add_cli_args(parser)
    raw_args, unknown_args = parser.parse_known_args(argv)
    server_args = ServerArgs.from_cli_args(raw_args, unknown_args)
    return server_args


def set_global_server_args(server_args: ServerArgs):
    """
    Set the global sgl_diffusion config for each process
    """
    global _global_server_args
    _global_server_args = server_args


def get_global_server_args() -> ServerArgs:
    if _global_server_args is None:
        # in ci, usually when we test custom ops/modules directly,
        # we don't set the sgl_diffusion config. In that case, we set a default
        # config.
        # TODO(will): may need to handle this for CI.
        raise ValueError("Global sgl_diffusion args is not set.")
    return _global_server_args