vllm-omni/vllm_omni/entrypoints/cli/serve.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""
Omni serve command for vLLM-Omni.

Supports both multi-stage LLM models (e.g., Qwen2.5-Omni) and
diffusion models (e.g., Qwen-Image) through the same CLI interface.
"""

import argparse
import json
import math
import os
import signal
from types import FrameType
from typing import Any, cast

import uvloop
from vllm.entrypoints.cli.types import CLISubcommand
from vllm.entrypoints.launchers.cli_args import make_arg_parser, validate_parsed_serve_args
from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
from vllm.logger import init_logger

from vllm_omni.entrypoints.cli.logo import log_logo
from vllm_omni.entrypoints.openai.api_server import omni_run_server
from vllm_omni.entrypoints.utils import parse_stage_overrides
from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace

logger = init_logger(__name__)

DESCRIPTION = """Launch a local OpenAI-compatible API server to serve Omni models
via HTTP. Supports both multi-stage LLM models and diffusion models.

The server automatically detects the model type:
- LLM models: Served via /v1/chat/completions endpoint
- Diffusion models: Served via /v1/images/generations endpoint

Examples:
  # Start an Omni LLM server
  vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091

  # Start a diffusion model server
  vllm serve Qwen/Qwen-Image --omni --port 8091

Search by using: `--help=<ConfigGroup>` to explore options by section (e.g.,
--help=OmniConfig)
  Use `--help=all` to show all available flags at once.
"""


def _parse_stage_overrides(value: str) -> dict[str, dict[str, Any]]:
    """Adapt shared stage-override validation to argparse's error type."""
    try:
        parsed = parse_stage_overrides(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError(str(exc)) from exc
    if parsed is None:
        raise argparse.ArgumentTypeError("--stage-overrides requires a JSON object")
    return parsed


def _nonneg_finite_float(value: str) -> float:
    """Argparse type for finite, non-negative floats (rejects nan/inf)."""
    try:
        parsed = float(value)
    except ValueError as exc:
        raise argparse.ArgumentTypeError(f"invalid float value: {value!r}") from exc
    if not math.isfinite(parsed) or parsed < 0:
        raise argparse.ArgumentTypeError(f"must be a finite non-negative number, got {value!r}")
    return parsed


def _json_object(value: str) -> dict[str, object]:
    try:
        parsed = json.loads(value)
    except json.JSONDecodeError as exc:
        raise argparse.ArgumentTypeError(f"must be valid JSON: {exc}") from exc
    if not isinstance(parsed, dict):
        raise argparse.ArgumentTypeError("must be a JSON object")
    return parsed


def _ensure_vllm_platform():
    """Ensure vLLM's current_platform is valid before arg parsing.

    Upstream vLLM's argument parser now instantiates DeviceConfig during
    ``make_arg_parser``, which requires a resolved platform with a non-empty
    ``device_type``.  In some environments (e.g. editable installs with
    broken package metadata), vLLM's own platform auto-detection may fail
    and fall back to ``UnspecifiedPlatform``.  When that happens, use the
    Omni platform (which has its own detection logic) as a drop-in
    replacement so that argument parsing succeeds.
    """
    from vllm import platforms as vllm_platforms

    if vllm_platforms.current_platform.is_unspecified():
        from vllm_omni.platforms import current_omni_platform

        if not current_omni_platform.is_unspecified():
            vllm_platforms.current_platform = current_omni_platform
            logger.debug(
                "Replaced vLLM UnspecifiedPlatform with omni platform %s",
                type(current_omni_platform).__name__,
            )
        else:
            from vllm.platforms.cpu import CpuPlatform

            vllm_platforms.current_platform = CpuPlatform()
            logger.debug(
                "Both vLLM and omni platforms are unspecified, falling back to CpuPlatform for arg parsing",
            )


class OmniServeCommand(CLISubcommand):
    """The `serve` subcommand for the vLLM CLI."""

    name = "serve"
    # Parser stashed at subparser_init so ``cmd`` can resolve each user-typed
    # flag to its real ``dest`` via the parser's action table.
    _parser: TrackingArgumentParser

    @staticmethod
    def cmd(args: TrackingNamespace) -> None:
        if not os.environ.get("VLLM_DISABLE_LOG_LOGO"):
            os.environ["VLLM_DISABLE_LOG_LOGO"] = "1"
            log_logo()

        # If model is specified in CLI (as positional arg), it takes precedence
        if hasattr(args, "model_tag") and args.model_tag is not None:
            args.model = args.model_tag

        if getattr(args, "no_guardrails", False):
            existing = getattr(args, "model_config", None)
            model_config = dict(existing) if isinstance(existing, dict) else {}
            model_config["guardrails"] = False
            args.model_config = model_config
            explicit_keys = getattr(args, "explicit_keys", None)
            if explicit_keys is not None:
                args.explicit_keys = explicit_keys | {"model_config"}

        if args.headless:
            run_headless(args)
        else:
            uvloop.run(omni_run_server(args))

    def validate(self, args: argparse.Namespace) -> None:
        if args.stage_id is not None and (args.omni_master_address is None or args.omni_master_port is None):
            raise ValueError("--stage-id requires both --omni-master-address and --omni-master-port to be set")

        # --omni-replica-address is only consulted in run_headless(); reject it
        # on the head so a misconfigured launch fails loudly instead of being
        # silently ignored.
        if getattr(args, "omni_replica_address", None) is not None and not args.headless:
            raise ValueError("--omni-replica-address requires --headless to be set")

        # --omni-dp-size-local is process-local. A value other than 1 only
        # makes sense when this process owns a stage (head or headless).
        omni_dp_size_local = getattr(args, "omni_dp_size_local", None)
        if omni_dp_size_local is not None:
            if omni_dp_size_local < 1:
                raise ValueError(f"--omni-dp-size-local must be >= 1, got {omni_dp_size_local}")
            if omni_dp_size_local != 1 and args.stage_id is None:
                raise ValueError("--omni-dp-size-local != 1 requires --stage-id to be set")

        # vLLM CLI args that omni does not honor: parallelism comes from the
        # per-stage YAML (parallel_config:, enable_expert_parallel:) and the
        # process-local replica count from --omni-dp-size-local. Passing the
        # vLLM equivalents on the command line would silently disagree with
        # those sources of truth, so reject them at parse time.
        if getattr(args, "omni", False):
            explicit_cli_keys: set[str] = getattr(args, "_cli_explicit_keys", set()) or set()
            prohibited_with_omni: dict[str, str] = {
                "data_parallel_size": "--data-parallel-size",
                "data_parallel_size_local": "--data-parallel-size-local",
                "data_parallel_address": "--data-parallel-address",
                "data_parallel_rpc_port": "--data-parallel-rpc-port",
                "data_parallel_start_rank": "--data-parallel-start-rank",
                "data_parallel_backend": "--data-parallel-backend",
                "api_server_count": "--api-server-count",
                "enable_expert_parallel": "--enable-expert-parallel",
            }
            offenders = sorted(flag for dest, flag in prohibited_with_omni.items() if dest in explicit_cli_keys)
            if offenders:
                raise ValueError(
                    "The following CLI args are not supported under --omni: "
                    f"{', '.join(offenders)}. Configure parallelism through the "
                    "per-stage YAML (`--deploy-config`) "
                    "and replica count via the per-stage `num_replicas` config field "
                    "(single-runtime) or `--omni-dp-size-local` (headless / multi-runtime)."
                )

        lora_backend = getattr(args, "lora_backend", None)
        lora_path = getattr(args, "lora_path", None)
        if lora_backend == "distill" and not lora_path:
            raise ValueError("--lora-backend distill requires --lora-path")
        if (lora_backend or "peft") == "peft" and isinstance(lora_path, list) and len(lora_path) > 1:
            raise ValueError("--lora-backend peft accepts only one startup LoRA path")

        # --omni-lb-policy is validated against the LoadBalancingPolicy enum.
        omni_lb_policy = getattr(args, "omni_lb_policy", None)
        if omni_lb_policy is not None:
            from vllm_omni.distributed.omni_coordinator import LoadBalancingPolicy

            try:
                LoadBalancingPolicy(omni_lb_policy)
            except ValueError as exc:
                valid = ", ".join(p.value for p in LoadBalancingPolicy)
                raise ValueError(f"--omni-lb-policy={omni_lb_policy!r} is not one of: {valid}") from exc

        omni_heartbeat_timeout = getattr(args, "omni_heartbeat_timeout", None)
        if omni_heartbeat_timeout is not None and omni_heartbeat_timeout <= 0:
            raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {omni_heartbeat_timeout}")

        # Skip validation for diffusion models as they have different requirements
        from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model

        model = getattr(args, "model_tag", None) or getattr(args, "model", None)
        if model and is_diffusion_model(model):
            logger.info("Detected diffusion model: %s", model)
            return
        validate_parsed_serve_args(args)

    def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgumentParser:
        serve_parser = subparsers.add_parser(
            self.name,
            description=DESCRIPTION,
            usage="vllm serve [model_tag] --omni [options]",
        )

        _ensure_vllm_platform()
        serve_parser = make_arg_parser(serve_parser)
        serve_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name)

        # Create OmniConfig argument group for omni-related parameters
        # This ensures the parameters appear in --help output
        omni_config_group = serve_parser.add_argument_group(
            title="OmniConfig", description="Configuration for vLLM-Omni multi-stage and diffusion models."
        )

        omni_config_group.add_argument(
            "--omni",
            action="store_true",
            help="Enable vLLM-Omni mode for multi-modal and diffusion models",
        )

        try:
            omni_config_group.add_argument(
                "--enable-sleep-mode",
                action="store_true",
                default=False,
                help="Enable GPU memory pool for sleep mode.",
            )
        except argparse.ArgumentError:
            pass

        omni_config_group.add_argument(
            "--task-type",
            type=str,
            default=None,
            help="Model-defined startup task type. The selected model validates "
            "supported values; for example, TTS models accept CustomVoice, "
            "VoiceDesign, or Base, while diffusion models may use it to select "
            "task-specific weights. If omitted, the model default is used.",
        )
        # Forced aligner / word timestamps. Either flag opts in (--forced-aligner-config
        # alone works when its YAML sets the model); heavier knobs live in that YAML.
        omni_config_group.add_argument(
            "--forced-aligner",
            type=str,
            default=None,
            help=(
                "Enable streaming TTS word timestamps via a forced aligner. "
                "Pass the aligner model path/name, e.g. 'Qwen/Qwen3-ForcedAligner-0.6B'. "
                "Disabled when omitted."
            ),
        )
        omni_config_group.add_argument(
            "--forced-aligner-config",
            type=str,
            default=None,
            help=(
                "Optional YAML file for forced aligner settings (model, runner, "
                "gpu_memory_utilization, dtype, max_model_len). The --forced-aligner "
                "flag, when set, overrides the YAML model field."
            ),
        )
        omni_config_group.add_argument(
            "--forced-aligner-device",
            type=str,
            default=None,
            help="Device(s) for the forced-aligner stage (e.g. '2'). Defaults to "
            "sharing an existing stage's GPU when unset.",
        )
        omni_config_group.add_argument(
            "--deploy-config",
            type=str,
            default=None,
            help="Path to a deploy config YAML (new format with stages/engine_args).",
        )
        omni_config_group.add_argument(
            "--strategy-config",
            type=str,
            default=None,
            help="Path to a composable-parallel strategy.yaml. Only applies to "
            "registry-based models (e.g. qwen2_5_omni): its derived parallel sizing "
            "is overlaid onto the registry-merged stages before per-stage engine "
            "args are built.",
        )
        omni_config_group.add_argument(
            "--stage-overrides",
            type=_parse_stage_overrides,
            default=None,
            help="Per-stage JSON overrides. Example: "
            '\'{"0": {"gpu_memory_utilization": 0.8}, "2": {"enforce_eager": true}}\'',
        )
        omni_config_group.add_argument(
            "--async-chunk",
            action=argparse.BooleanOptionalAction,
            default=None,
            help="Override the deploy YAML's ``async_chunk:`` bool. Unset leaves the YAML value in force.",
        )
        omni_config_group.add_argument(
            "--stage-id",
            type=int,
            default=None,
            help="Select and launch a single stage by stage_id.",
        )
        omni_config_group.add_argument(
            "--replica-id",
            type=int,
            default=None,
            help=(
                "Deprecated and ignored — replica ids are auto-assigned by the "
                "master server. Specifying this flag prints a warning and has "
                "no effect."
            ),
        )
        omni_config_group.add_argument(
            "--stage-init-timeout",
            type=int,
            default=300,
            help="The timeout for initializing a single stage in seconds (default: 300)",
        )
        omni_config_group.add_argument(
            "--init-timeout",
            type=int,
            default=600,
            help="The timeout for initializing the stages.",
        )
        omni_config_group.add_argument(
            "--shm-threshold-bytes",
            type=int,
            default=65536,
            help="The threshold for the shared memory size.",
        )
        omni_config_group.add_argument(
            "--log-stats",
            action="store_true",
            help="Enable logging the stats.",
        )
        omni_config_group.add_argument(
            "--log-file",
            type=str,
            default=None,
            help="The path to the log file.",
        )
        omni_config_group.add_argument(
            "--batch-timeout",
            type=int,
            default=10,
            help="The timeout for the batch.",
        )
        omni_config_group.add_argument(
            "--worker-backend",
            type=str,
            default="multi_process",
            choices=["multi_process", "ray"],
            help="The backend to use for stage workers.",
        )
        omni_config_group.add_argument(
            "--ray-address",
            type=str,
            default=None,
            help="The address of the Ray cluster to connect to.",
        )
        omni_config_group.add_argument(
            "--omni-master-address",
            "-oma",
            type=str,
            help="Hostname or IP address of the Omni orchestrator (master).",
        )
        omni_config_group.add_argument(
            "--omni-master-port",
            "-omp",
            type=int,
            help="Port of the Omni orchestrator (master).",
        )
        omni_config_group.add_argument(
            "--omni-replica-address",
            "-ora",
            type=str,
            default=None,
            help=(
                "Local bind address (this host's IP) that the headless stage "
                "advertises to the Omni master for its handshake/input/output "
                "ZMQ sockets. If unset, auto-detected via a UDP-connect "
                "routing probe against --omni-master-address. Override only "
                "when the auto-detected IP is wrong (e.g. multi-NIC host "
                "where the master is reachable on the wrong interface)."
            ),
        )
        omni_config_group.add_argument(
            "--omni-dp-size-local",
            type=int,
            default=1,
            help=(
                "Number of stage replicas this runtime launches locally for its "
                "own --stage-id. Process-local: head and every headless invocation "
                "read their own copy; values may differ across invocations. "
                "Requires --stage-id to be set when not equal to 1."
            ),
        )
        omni_config_group.add_argument(
            "--omni-lb-policy",
            type=str,
            default="random",
            choices=["random", "round-robin", "least-queue-length"],
            help=(
                "Per-stage load-balancing policy used by the head's StagePool to "
                "route requests across UP replicas. Only consulted on the head runtime."
            ),
        )
        omni_config_group.add_argument(
            "--omni-heartbeat-timeout",
            type=float,
            default=30.0,
            help=(
                "Seconds before an unreporting replica is marked ERROR in the "
                "OmniCoordinator. Only consulted on the head runtime."
            ),
        )

        # Diffusion model specific arguments
        omni_config_group.add_argument(
            "--num-gpus",
            type=int,
            default=None,
            help="Number of GPUs to use for diffusion model inference.",
        )
        omni_config_group.add_argument(
            "--model-class-name",
            dest="model_class_name",
            type=str,
            default=None,
            help="Override the diffusion pipeline class name (e.g. LTX2Pipeline).",
        )
        omni_config_group.add_argument(
            "--diffusion-load-format",
            dest="diffusion_load_format",
            type=str,
            default=None,
            choices=["default", "custom_pipeline", "dummy", "diffusers"],
            help=(
                "How to load the diffusion pipeline: native/registry (default), "
                "custom_pipeline, dummy, or diffusers for the HF diffusers adapter."
            ),
        )
        omni_config_group.add_argument(
            "--lora-path",
            type=str,
            nargs="+",
            default=None,
            help=(
                "LoRA checkpoint path(s) loaded when the diffusion server starts. "
                "The distill backend accepts one file per pipeline transformer."
            ),
        )
        omni_config_group.add_argument(
            "--lora-backend",
            type=str,
            choices=["peft", "distill"],
            default=None,
            help=(
                "Diffusion LoRA loading backend. 'distill' fuses checkpoint files "
                "into the base model at server startup; 'peft' uses the adapter manager."
            ),
        )
        omni_config_group.add_argument(
            "--lora-scale",
            type=float,
            default=None,
            help="Scale for a startup PEFT LoRA. Distilled LoRAs are fused at their checkpoint scale.",
        )
        omni_config_group.add_argument(
            "--diffusion-compile-granularity",
            choices=["regional", "full"],
            default=None,
            help=(
                "Compilation scope for the generic diffusion model runner. "
                "'regional' compiles repeated blocks (default); 'full' compiles the whole transformer and is "
                "incompatible with HSDP, sequence parallelism, CPU offload, and layerwise offload."
            ),
        )
        omni_config_group.add_argument(
            "--diffusion-compile-dynamic",
            action=argparse.BooleanOptionalAction,
            default=None,
            help=(
                "Use dynamic shapes for the selected generic diffusion compile scope. "
                "Disable for fixed-shape workloads with --no-diffusion-compile-dynamic."
            ),
        )
        omni_config_group.add_argument(
            "--fa-deterministic",
            dest="fa_deterministic",
            action="store_true",
            default=False,
            help=(
                "Request FlashAttention deterministic=True on the local FLASH_ATTN dense path. "
                "Slower than the library default deterministic=False; intended for accuracy CI. "
                "Serving default remains non-deterministic."
            ),
        )
        omni_config_group.add_argument(
            "--diffusers-load-kwargs",
            dest="diffusers_load_kwargs",
            type=json.loads,
            default="{}",
            help=(
                "JSON object passed to DiffusionPipeline.from_pretrained()."
                "It overrides corresponding parameters in the standard vLLM-Omni interface."
                '(e.g. \'{"use_safetensors": true, "variant": "fp16"}\').'
            ),
        )
        omni_config_group.add_argument(
            "--diffusers-call-kwargs",
            dest="diffusers_call_kwargs",
            type=json.loads,
            default="{}",
            help=(
                "JSON object passed to pipeline.__call__(). "
                "Useful for model-specific sampling parameters not covered by the vLLM-Omni interface."
                "During request time, it is overridden by corresponding parameters in the vLLM-Omni interface."
                '(e.g. \'{"num_inference_steps": 30, "guidance_scale": 7.5}\').'
            ),
        )
        omni_config_group.add_argument(
            "--usp",
            "--ulysses-degree",
            dest="ulysses_degree",
            type=int,
            default=None,
            help="Ulysses Sequence Parallelism degree for diffusion models. "
            "Equivalent to setting DiffusionParallelConfig.ulysses_degree.",
        )
        omni_config_group.add_argument(
            "--ulysses-mode",
            type=str,
            default="strict",
            choices=["strict", "advanced_uaa"],
            help="Ulysses sequence-parallel mode for diffusion models. "
            "'strict' keeps the original divisibility requirements; "
            "'advanced_uaa' enables the experimental UAA path for uneven sequence/head shapes.",
        )
        omni_config_group.add_argument(
            "--ulysses-a2a-permute",
            action=argparse.BooleanOptionalAction,
            default=None,
            help=(
                "Enable fused permute-free Ulysses all-to-all over NCCL symmetric memory. "
                "Only strict Ulysses layouts are eligible. Defaults to disabled."
            ),
        )
        omni_config_group.add_argument(
            "--ring",
            "--ring-degree",
            dest="ring_degree",
            type=int,
            default=None,
            help="Ring Sequence Parallelism degree for diffusion models. "
            "Equivalent to setting DiffusionParallelConfig.ring_degree.",
        )
        omni_config_group.add_argument(
            "--allgather-degree",
            dest="allgather_degree",
            type=int,
            default=None,
            help="AllGather-KV Sequence Parallelism degree for non-causal diffusion attention. "
            "Equivalent to setting DiffusionParallelConfig.allgather_degree.",
        )
        omni_config_group.add_argument(
            "--diffusion-quantization-config",
            type=json.loads,
            default=None,
            help=(
                "JSON string for diffusion quantization_config. "
                'Example: \'{"method":"fp8","activation_scheme":"dynamic"}\'.'
            ),
        )
        omni_config_group.add_argument(
            "--force-cutlass-fp8",
            action="store_true",
            default=None,
            help=(
                "Diffusion-only runtime override for ModelOpt FP8 checkpoints: "
                "force CUTLASS FP8 linear kernels on CUDA SM89+ devices. "
                "Ignored for BF16, non-ModelOpt FP8, ROCm, and older CUDA GPUs."
            ),
        )

        # HSDP (Hybrid Sharded Data Parallel) parameters
        omni_config_group.add_argument(
            "--use-hsdp",
            dest="use_hsdp",
            action="store_true",
            help="Enable HSDP (Hybrid Sharded Data Parallel) for diffusion models. "
            "Shards model weights across GPUs to reduce per-GPU memory usage.",
        )
        omni_config_group.add_argument(
            "--hsdp-shard-size",
            type=int,
            default=-1,
            help="Number of GPUs to shard weights across. -1 = auto (world_size / replicate_size).",
        )
        omni_config_group.add_argument(
            "--hsdp-replicate-size",
            type=int,
            default=1,
            help="Number of replica groups for HSDP. Each group holds a full sharded copy.",
        )

        # Attention backend configuration
        omni_config_group.add_argument(
            "--diffusion-attention-backend",
            dest="diffusion_attention_backend",
            type=str,
            default=None,
            help="Diffusion attention backend (shorthand). "
            "Sets the default backend for all diffusion attention roles, e.g. 'FLASH_ATTN'. "
            "May be combined with --diffusion-attention-config.per_role.* overrides, "
            "but mutually exclusive with --diffusion-attention-config.default.backend.",
        )
        omni_config_group.add_argument(
            "--fastvideo-vsa-topk",
            type=int,
            default=None,
            help="Number of key/value blocks selected per query block by FASTVIDEO_VSA.",
        )
        omni_config_group.add_argument(
            "--diffusion-attention-config",
            "-dac",
            dest="diffusion_attention_config",
            type=json.loads,
            default=None,
            help="Diffusion attention config. Accepts JSON or vLLM-style dotted flags. "
            "Examples: "
            "--diffusion-attention-config.default.backend FLASH_ATTN, "
            "--diffusion-attention-config.default.backend TRTLLM_ATTN "
            "--diffusion-attention-config.default.skip_softmax.target_sparsity 0.5, "
            "--diffusion-attention-config.per_role.cross.backend SAGE_ATTN, "
            '--diffusion-attention-config \'{"default": {"backend": "FLASH_ATTN"}, '
            '"per_role": {"cross": {"backend": "SAGE_ATTN"}}}\'.',
        )

        # Cache optimization parameters
        omni_config_group.add_argument(
            "--cache-backend",
            type=str,
            default="none",
            help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache', 'step_cache'",
        )
        omni_config_group.add_argument(
            "--cache-config",
            type=str,
            default=None,
            help="JSON string of cache configuration. "
            "TeaCache: '{\"rel_l1_thresh\": 0.2}'. "
            'MagCache: \'{"mag_threshold": 0.24, "mag_max_skip_steps": 5, "mag_retention_ratio": 0.1}\'. '
            "Calibration mode: add '\"mag_calibrate\": true'",
        )
        omni_config_group.add_argument(
            "--video-output-transport",
            type=_json_object,
            default=None,
            help=(
                "JSON object configuring video output preparation, for example '{\"enable_device_postprocess\": true}'."
            ),
        )
        omni_config_group.add_argument(
            "--enable-cache-dit-summary",
            action="store_true",
            help="Enable cache-dit summary logging after diffusion forward passes.",
        )
        omni_config_group.add_argument(
            "--step-execution",
            action="store_true",
            help="Enable per-step diffusion execution so running requests can be aborted between denoise steps.",
        )
        omni_config_group.add_argument(
            "--request-batch-max-wait-ms",
            type=_nonneg_finite_float,
            default=0.0,
            help="Request-mode batch admission: max milliseconds to wait for compatible "
            "requests to accumulate before scheduling a fused forward wave. "
            "0 disables admission (default).",
        )

        # VAE memory optimization parameters
        omni_config_group.add_argument(
            "--vae-use-slicing",
            action="store_true",
            help="Enable VAE slicing for memory optimization (useful for mitigating OOM issues).",
        )
        omni_config_group.add_argument(
            "--vae-use-tiling",
            action="store_true",
            help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).",
        )

        # Parallel weight loading (faster diffusion startup)
        omni_config_group.add_argument(
            "--disable-multithread-weight-load",
            action="store_false",
            dest="enable_multithread_weight_load",
            default=True,
            help="Disable multi-threaded safetensors loading (default: enabled with 4 threads).",
        )
        omni_config_group.add_argument(
            "--num-weight-load-threads",
            type=int,
            default=4,
            help="Number of threads for parallel weight loading (default: 4).",
        )

        # diffusion model offload parameters
        omni_config_group.add_argument(
            "--diffusion-offload-config",
            type=json.loads,
            default=None,
            help="Diffusion CPU-offload config as JSON. "
            "Set mode to module or layer, list dit and/or text_encoder in "
            "components, and put layer-only tuning under layer_options. "
            "Layer settings are weight_transfer (rank-local or allgather) and "
            "resident_layers (DiT only).",
        )
        omni_config_group.add_argument(
            "--enable-cpu-offload",
            action="store_true",
            help="Compatibility alias for model-level CPU offload. New integrations should use "
            "--diffusion-offload-config with mode=module and explicit components.",
        )
        omni_config_group.add_argument(
            "--enable-layerwise-offload",
            action="store_true",
            help="Compatibility alias for layerwise CPU offload. New integrations should use "
            "--diffusion-offload-config with mode=layer and explicit components.",
        )
        omni_config_group.add_argument(
            "--enable-distributed-layerwise-offload",
            action="store_true",
            help="Compatibility alias for distributed layerwise CPU offload. "
            "New integrations should use mode=layer and configure weight transfer per component.",
        )
        omni_config_group.add_argument(
            "--dlo-use-allgather",
            dest="dlo_use_allgather",
            action="store_true",
            default=True,
            help="Compatibility option; use component weight_transfer=allgather in new configurations. "
            "Use shard + AllGather for weight reconstruction (default: True). "
            "When disabled (--dlo-no-use-allgather), each rank streams the "
            "standard loader's rank-local tensors via H2D only — no additional "
            "DP sharding, no AllGather, and no concurrent-request requirement.",
        )
        omni_config_group.add_argument(
            "--dlo-no-use-allgather",
            dest="dlo_use_allgather",
            action="store_false",
            help=(
                "Compatibility option; use component weight_transfer=rank-local in new configurations. "
                "Disable AllGather and stream standard-loader rank-local weights "
                "independently (including existing TP shards)."
            ),
        )
        omni_config_group.add_argument(
            "--dlo-resident-layers",
            type=int,
            default=0,
            help="Compatibility option; use layer_options.dit.resident_layers in new configurations.",
        )
        omni_config_group.add_argument(
            "--host-weight-runtime-mode",
            choices=("disabled", "preferred", "required"),
            default="disabled",
            help=(
                "Host Weight Runtime policy for eligible no-AllGather DLO: "
                "disabled does not consult HWR; preferred restores an exact hit "
                "or canonically loads and publishes on a miss; required restores "
                "an exact hit or fails startup. Populate a required store with "
                "preferred first."
            ),
        )
        omni_config_group.add_argument(
            "--host-weight-runtime-root",
            type=str,
            default=None,
            help=(
                "Writable node-local Host Weight Runtime store shared by workers "
                "in one storage domain. Required for preferred and required; use "
                "the same persistent path for population and serving."
            ),
        )
        omni_config_group.add_argument(
            "--dlo-host-registration-limit-gib",
            type=float,
            default=0.0,
            help=(
                "Optional per-worker GiB ceiling for registering an HWR mmap for direct H2D. "
                "Zero applies no additional ceiling. Eligible no-AllGather HWR hits attempt registration "
                "under the existing pinned-memory policy and fall back to bounded staging when unavailable."
            ),
        )
        # Video model parameters (e.g., Wan2.2) - engine-level
        omni_config_group.add_argument(
            "--boundary-ratio",
            type=float,
            default=None,
            help="Boundary split ratio for low/high DiT in video models (e.g., 0.875 for Wan2.2).",
        )
        omni_config_group.add_argument(
            "--flow-shift",
            type=float,
            default=None,
            help="Scheduler flow_shift for video models (e.g., 5.0 for 720p, 12.0 for 480p).",
        )
        # Diffusion KV-cache quantization uses dedicated flags so we do not reuse
        # vLLM's --kv-cache-dtype (AR cache dtype, default "auto").
        omni_config_group.add_argument(
            "--diffusion-kv-cache-dtype",
            type=str,
            default=None,
            help="Diffusion attention KV cache dtype (e.g. fp8). Separate from vLLM --kv-cache-dtype.",
        )
        omni_config_group.add_argument(
            "--diffusion-kv-cache-skip-steps",
            type=str,
            default=None,
            help="Diffusion KV-cache quantization skip-step selector, e.g. '0-9,20,25-30'.",
        )
        omni_config_group.add_argument(
            "--diffusion-kv-cache-skip-layers",
            type=str,
            default=None,
            help="Diffusion KV-cache quantization skip-layer selector, e.g. '0,1,4-8'.",
        )
        omni_config_group.add_argument(
            "--cfg-parallel-size",
            type=int,
            default=1,
            help="Number of devices used to execute diffusion guidance passes in parallel. "
            "Equivalent to setting DiffusionParallelConfig.cfg_parallel_size.",
        )
        omni_config_group.add_argument(
            "--vae-patch-parallel-size",
            type=int,
            default=1,
            help="VAE Patch Parallelism degree for diffusion models. "
            "Distributes VAE decode workload across multiple ranks by splitting the latent spatially. "
            "Equivalent to setting DiffusionParallelConfig.vae_patch_parallel_size.",
        )
        omni_config_group.add_argument(
            "--text-encoder-tp-size",
            type=int,
            default=None,
            help="Tensor-parallel degree for the diffusion text encoder. "
            "Shards the encoder across the first N DiT ranks. "
            "Equivalent to setting DiffusionParallelConfig.text_encoder_tp_size.",
        )
        omni_config_group.add_argument(
            "--vae-parallel-mode",
            type=str,
            default="tile",
            choices=["tile", "spatial_shard_height", "spatial_shard_width"],
            help="VAE parallel decode strategy for diffusion models. "
            "'tile' (default) uses patch/tile parallel decode; "
            "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits "
            "decoder feature maps along height/width and exchanges halo regions. The "
            "'spatial_shard_*' modes require vae_patch_parallel_size to match the DiT group size. "
            "Equivalent to setting DiffusionParallelConfig.vae_parallel_mode.",
        )

        # Default sampling parameters
        omni_config_group.add_argument(
            "--default-sampling-params",
            type=str,
            help="Json str for Default sampling parameters, \n"
            'Structure: {"<stage_id>": {<sampling_param>: value, ...}, ...}\n'
            'e.g., \'{"0": {"num_inference_steps":50, "guidance_scale":1}}\'. '
            "Currently only supports diffusion models.",
        )
        # Diffusion model mixed precision
        omni_config_group.add_argument(
            "--max-generated-image-size",
            default=7680 * 4320,  # 8K resolution
            type=int,
            help="Maximum generated image size in pixels (height * width).",
        )
        # Diffusion model (mainly video generation models) streaming output mode
        omni_config_group.add_argument(
            "--diffusion-streaming-output",
            dest="diffusion_streaming_output",
            action="store_true",
            default=False,
            help="Enable chunked streaming output for diffusion (mainly video generation) models that support it.",
        )
        # TTS-specific parameters
        omni_config_group.add_argument(
            "--tts-max-instructions-length",
            type=int,
            default=None,
            help="Maximum length for TTS voice style instructions (overrides the pipeline default, default: 500).",
        )

        # Disable safety guardrails for this server (currently only applicable for Cosmos3)
        # TODO: drop once --model-config-override lands (3/N config refactor)
        omni_config_group.add_argument(
            "--no-guardrails",
            dest="no_guardrails",
            action="store_true",
            help="Disable Cosmos3 text/video safety guardrails for this server.",
        )

        # Enable diffusion pipeline profiling
        omni_config_group.add_argument(
            "--enable-diffusion-pipeline-profiler",
            action="store_true",
            help="Enable diffusion pipeline profiler to display stage durations.",
        )
        omni_config_group.add_argument(
            "--enable-ar-profiler",
            action="store_true",
            help="Enable AR stage profiler to include AR stage timing in stage_durations.",
        )
        omni_config_group.add_argument(
            "--enable-orch-monitor",
            action="store_true",
            help="Enable orchestrator window monitor and write a JSON log at shutdown.",
        )

        # Supplementary auxiliary text encoder parameters
        # (e.g., the meta llama/meta llama-3.1-8b-instrument used by hidream)
        omni_config_group.add_argument(
            "--auxiliary-text-encoder",
            type=str,
            default=None,
            help="Auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
        )

        # Stash via type(self) so the docs hook (which execs this function in a
        # sandboxed globals dict via ``DummySelf``) doesn't fail on a NameError.
        type(self)._parser = serve_parser

        return serve_parser


def run_headless(args: TrackingNamespace) -> None:
    """Run a single stage in headless mode.

    Honors ``--omni-dp-size-local``: launches that many replicas locally for
    ``--stage-id``. Each replica registers with the head's OmniMasterServer
    (auto-assigned replica id when ``--omni-dp-size-local > 1`` so multiple
    headless invocations can coexist) and reports heartbeats to the head's
    OmniCoordinator.
    """
    from vllm.v1.executor.multiproc_executor import MultiprocExecutor
    from vllm.version import __version__ as VLLM_VERSION

    from vllm_omni.config.resolver import resolve_omni_config
    from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_omni_kv_config_for_stage
    from vllm_omni.engine.stage_engine_startup import (
        get_headless_replica_devices,
        launch_headless_diffusion_replicas,
        launch_headless_llm_replicas,
    )
    from vllm_omni.engine.stage_init_utils import (
        build_engine_args_dict,
        build_vllm_config,
        get_stage_connector_spec,
        inject_omni_kv_connector_config,
        load_omni_transfer_config_for_model,
        prepare_engine_environment,
    )

    model = args.model
    stage_id: int | None = args.stage_id
    omni_master_address: str | None = args.omni_master_address
    omni_master_port: int | None = args.omni_master_port
    worker_backend: str | None = args.worker_backend
    omni_replica_address: str | None = getattr(args, "omni_replica_address", None)
    omni_dp_size_local: int = max(1, int(getattr(args, "omni_dp_size_local", 1) or 1))

    if not model:
        raise ValueError("Failed to pass model from kwargs")
    if stage_id is None:
        raise ValueError("--stage-id is required in headless mode")
    if omni_master_address is None or omni_master_port is None:
        raise ValueError("--omni-master-address and --omni-master-port are required in headless mode")
    if worker_backend != "multi_process":
        raise ValueError("headless mode requires worker_backend=multi_process")

    # Filter down to a dict of things explicitly requested by the user
    args_dict = args.get_explicit_kwargs_dict()
    # This CLI-only negative alias is consumed below when selecting the
    # launcher log_stats value; it is not a per-stage config override.
    args_dict.pop("disable_log_stats", None)

    deploy_config_path = args_dict.pop("deploy_config", None)
    strategy_config_path = args_dict.pop("strategy_config", None)
    stage_overrides = args_dict.pop("stage_overrides", None)

    # ``--replica-id`` is deprecated and ignored — replica ids are
    # auto-assigned by ``OmniMasterServer`` so headless processes carry
    # no knowledge of their per-replica id at launch time. Warn (don't
    # error) when the operator still supplies it so existing launchers
    # keep working with a single log line.
    if "replica_id" in args_dict:
        logger.warning(
            "[Headless] --replica-id is deprecated and ignored "
            "(supplied value: %s). Replica ids are auto-assigned by the "
            "master server.",
            args.replica_id,
        )
        args_dict.pop("replica_id")

    resolved = resolve_omni_config(
        model,
        # store_true cannot express an explicit False: absent maps to None
        # ("not specified") so the deploy yaml's per-stage value applies.
        trust_remote_code=getattr(args, "trust_remote_code", None) or None,
        cli_overrides=args_dict,
        deploy_config_path=deploy_config_path,
        stage_overrides=stage_overrides,
        strategy_config_path=strategy_config_path,
    )
    config_path = resolved.config_path
    stage_configs = list(resolved.stage_configs)

    try:
        stage_cfg = resolved.stage_by_id(stage_id)
    except KeyError:
        raise ValueError(
            f"No stage config found for stage_id={stage_id}. Available stage ids: {[c.stage_id for c in stage_configs]}"
        ) from None

    prepare_engine_environment()
    per_replica_devices = get_headless_replica_devices(stage_cfg, stage_id, omni_dp_size_local)

    if stage_cfg.stage_type == "diffusion":
        launch_headless_diffusion_replicas(
            model=model,
            stage_cfg=stage_cfg,
            stage_configs=stage_configs,
            stage_id=stage_id,
            omni_master_address=omni_master_address,
            omni_master_port=omni_master_port,
            omni_dp_size_local=omni_dp_size_local,
            per_replica_devices=per_replica_devices,
            config_path=cast(str, config_path),
            replica_bind_address=omni_replica_address,
        )
        return

    omni_transfer_config = load_omni_transfer_config_for_model(model, config_path)
    omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id)
    stage_connector_spec = get_stage_connector_spec(
        omni_transfer_config=omni_transfer_config,
        stage_id=stage_id,
        async_chunk=bool(stage_cfg.engine_args.get("async_chunk", False)),
    )

    # ``runtime_cfg`` is mostly inherited from the parent's
    # CUDA_VISIBLE_DEVICES; when ``--omni-dp-size-local > 1`` we additionally
    # bracket each replica's spawn below with setup_stage_devices so they
    # don't all stack on cuda:0 (see ``per_replica_devices`` above).
    # Headless startup still supplies the legacy OmegaConf stage shape through
    # the stable adapter entry point. The implementation switches only when
    # RFC #4021 threads structured stage configs through the launch plan.
    engine_args_dict = build_engine_args_dict(
        stage_cfg,
        model,
        stage_connector_spec=stage_connector_spec,
        cli_tokenizer=getattr(args, "tokenizer", None),
    )

    inject_omni_kv_connector_config(engine_args_dict, omni_kv_connector, stage_id)

    vllm_config, executor_class = build_vllm_config(
        stage_cfg,
        model,
        stage_connector_spec=stage_connector_spec,
        engine_args_dict=engine_args_dict,
        headless=True,
    )
    parallel_config = vllm_config.parallel_config

    shutdown_requested = False

    def signal_handler(signum: int, frame: FrameType | None) -> None:
        nonlocal shutdown_requested
        logger.debug("Received %d signal.", signum)
        if not shutdown_requested:
            shutdown_requested = True
            raise SystemExit

    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)

    if parallel_config.node_rank_within_dp > 0:
        head_node_address = f"{parallel_config.master_addr}:{parallel_config.master_port}"
        logger.info(
            "Launching vLLM-Omni (v%s) headless multiproc executor, "
            "with head node address %s for torch.distributed process group.",
            VLLM_VERSION,
            head_node_address,
        )

        executor = MultiprocExecutor(vllm_config, monitor_workers=False)
        executor.start_worker_monitor(inline=True)
        return

    log_stats = bool(args.log_stats)
    if args.disable_log_stats:
        log_stats = False

    launch_headless_llm_replicas(
        vllm_config=vllm_config,
        executor_class=executor_class,
        log_stats=log_stats,
        omni_master_address=omni_master_address,
        omni_master_port=omni_master_port,
        stage_id=stage_id,
        stage_config=stage_cfg,
        omni_dp_size_local=omni_dp_size_local,
        per_replica_devices=per_replica_devices,
        replica_bind_address=omni_replica_address,
    )


def cmd_init() -> list[CLISubcommand]:
    return [OmniServeCommand()]