vllm-omni/vllm_omni/diffusion/data.py

# adapted from sglang and fastvideo
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import copy
import json
import math
import os
import random
import warnings
from collections.abc import Callable, Mapping
from dataclasses import dataclass, field, fields
from enum import Enum
from typing import TYPE_CHECKING, Any

import diffusers
import huggingface_hub
import torch
from PIL import Image
from pydantic import Field, model_validator
from typing_extensions import Self
from vllm.config.utils import config
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig,
)
from vllm.transformers_utils.repo_utils import get_model_path

from vllm_omni.diffusion.diffusion_kv.config import (
    DiffusionKVCacheMode,
    parse_diffusion_kv_cache_mode,
)
from vllm_omni.diffusion.lora.manager import LoRABackend
from vllm_omni.diffusion.media import DiffusionMediaOutput
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
from vllm_omni.diffusion.utils.network_utils import is_port_available
from vllm_omni.errors import client_error_metadata
from vllm_omni.quantization import build_quant_config

if TYPE_CHECKING:
    from vllm.config import KVTransferConfig, ProfilerConfig

# Import after TYPE_CHECKING to avoid circular imports at runtime
# The actual import is deferred to __post_init__ to avoid import order issues

logger = init_logger(__name__)

# Accepted values for ``OmniDiffusionConfig.vae_fast_path``.
VAE_FAST_PATH_LEVELS: tuple[str, ...] = ("off", "lossless", "channels_last")


def _move_diffusion_alias(
    normalized: dict[str, Any],
    legacy_name: str,
    canonical_name: str,
) -> None:
    legacy_value = normalized.pop(legacy_name, None)
    if legacy_value is None:
        return
    if normalized.get(canonical_name) is not None:
        raise ValueError(f"Diffusion config fields {legacy_name!r} and {canonical_name!r} cannot both be provided.")
    warnings.warn(
        f"Diffusion config field {legacy_name!r} is deprecated; use {canonical_name!r}.",
        FutureWarning,
        stacklevel=3,
    )
    normalized[canonical_name] = legacy_value


def normalize_omni_diffusion_kwargs(
    raw_kwargs: Mapping[str, Any],
    *,
    apply_defaults: bool = True,
) -> dict[str, Any]:
    """Normalize diffusion kwargs, deferring defaults until sources are merged."""
    config_kwargs = dict(raw_kwargs)

    dtype = config_kwargs.get("dtype")
    if dtype is None:
        if apply_defaults:
            config_kwargs["dtype"] = "auto"
    elif isinstance(dtype, torch.dtype):
        config_kwargs["dtype"] = str(dtype).removeprefix("torch.")
    elif not isinstance(dtype, str):
        raise TypeError(f"Provided dtype must be a string or torch.dtype, got {type(dtype).__name__}")

    for legacy, canonical in (
        ("static_lora_scale", "lora_scale"),
        ("quantization", "quantization_config"),
        ("diffusion_quantization_config", "quantization_config"),
        ("max_batch_size", "max_num_seqs"),
        ("kv_cache_skip_steps", "diffusion_kv_cache_skip_steps"),
        ("kv_cache_skip_layers", "diffusion_kv_cache_skip_layers"),
    ):
        _move_diffusion_alias(config_kwargs, legacy, canonical)
    _move_diffusion_alias(config_kwargs, "kv_cache_dtype", "diffusion_kv_cache_dtype")

    # Handle "diffusion_attention_backend" shorthand: merge into
    # diffusion_attention_config before field filtering.
    diffusion_attn_backend = config_kwargs.pop("diffusion_attention_backend", None)
    fastvideo_vsa_topk = config_kwargs.pop("fastvideo_vsa_topk", None)
    if diffusion_attn_backend is not None or fastvideo_vsa_topk is not None:
        if diffusion_attn_backend is not None:
            warnings.warn(
                "Diffusion config field 'diffusion_attention_backend' is deprecated; use 'diffusion_attention_config'.",
                FutureWarning,
                stacklevel=2,
            )
        existing = config_kwargs.get("diffusion_attention_config")
        config_kwargs["diffusion_attention_config"] = parse_attention_config(
            existing,
            attention_backend=diffusion_attn_backend,
            fastvideo_vsa_topk=fastvideo_vsa_topk,
        )

    auxiliary_text_encoder = config_kwargs.pop("auxiliary_text_encoder", None)
    if auxiliary_text_encoder is not None:
        extras = dict(config_kwargs.get("extras") or {})
        if extras.get("auxiliary_text_encoder") is not None:
            raise ValueError(
                "Diffusion engine field 'auxiliary_text_encoder' cannot be provided both at the top level and in "
                "'extras'."
            )
        extras["auxiliary_text_encoder"] = auxiliary_text_encoder
        config_kwargs["extras"] = extras

    # Check environment variable as fallback for cache_backend.
    # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND.
    if "cache_backend" not in config_kwargs and apply_defaults:
        cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
        config_kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none"
    elif "cache_backend" in config_kwargs and config_kwargs["cache_backend"] is None and apply_defaults:
        # Callers (e.g. example CLIs with `default=None`) pass an explicit
        # None for "no cache"; canonicalize it so every consumer sees the
        # declared `str` value instead of relying on per-model None handling.
        config_kwargs["cache_backend"] = "none"

    cache_config = config_kwargs.get("cache_config")
    if isinstance(cache_config, str):
        try:
            config_kwargs["cache_config"] = json.loads(cache_config)
        except json.JSONDecodeError:
            logger.warning("Invalid cache_config JSON, using backend defaults.")
            config_kwargs.pop("cache_config", None)

    if config_kwargs.get("streaming_output") is None and config_kwargs.get("diffusion_streaming_output") is not None:
        config_kwargs["streaming_output"] = config_kwargs["diffusion_streaming_output"]
    config_kwargs.pop("diffusion_streaming_output", None)

    # Convert optional YAML null values to empty containers.
    for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"):
        if key in config_kwargs and config_kwargs[key] is None:
            config_kwargs[key] = {}

    return config_kwargs


def validate_host_weight_runtime_options(*, mode: object, root: object) -> None:
    """Validate HWR policy without touching the configured storage domain.

    Filesystem locality and store construction belong to the eligible loader
    path. Keeping this check purely structural is what lets disabled and
    AllGather DLO configurations retain zero HWR interaction.
    """
    if mode not in {"disabled", "preferred", "required"}:
        raise ValueError("host_weight_runtime_mode must be disabled, preferred, or required")
    if mode != "disabled" and (not isinstance(root, str) or not root.strip()):
        raise ValueError("enabled Host Weight Runtime requires host_weight_runtime_root")


def validate_dlo_host_registration_options(
    *,
    limit_gib: object,
    enable_dlo: bool,
    use_allgather: bool,
    hwr_mode: object,
) -> float:
    """Validate the optional transport budget without probing CUDA or HWR."""
    if not isinstance(limit_gib, (int, float, str)):
        raise TypeError(f"dlo_host_registration_limit_gib must be a number; got {type(limit_gib).__name__}")
    value = float(limit_gib)
    if not math.isfinite(value) or value < 0:
        raise ValueError("dlo_host_registration_limit_gib must be finite and >= 0")
    if value and (not enable_dlo or use_allgather or hwr_mode == "disabled"):
        raise ValueError("dlo_host_registration_limit_gib requires enabled no-AllGather DLO and Host Weight Runtime")
    return value


def validate_omni_diffusion_kwargs(
    kwargs: Mapping[str, Any],
    allowed_fields: set[str] | frozenset[str],
    *,
    stage_id: int | str | None = None,
) -> None:
    """Reject every field without an owner."""
    unknown = set(kwargs) - allowed_fields
    if unknown:
        names = ", ".join(repr(name) for name in sorted(unknown))
        suffix = "" if stage_id is None else f" for stage {stage_id}"
        raise ValueError(f"Unknown diffusion config field(s){suffix}: {names}")


def parse_kv_cache_skip_selector(
    selector: str | list[int] | tuple[int, ...] | set[int] | None,
) -> set[int] | None:
    """Parse a non-negative index selector such as "0-9,20,25-30"."""
    if selector is None:
        return None
    values: set[int]
    if isinstance(selector, set):
        values = selector
    elif isinstance(selector, (list, tuple)):
        values = set(selector)
    elif isinstance(selector, str):
        text = selector.strip()
        if not text:
            return None
        values = set()
        for chunk in text.split(","):
            token = chunk.strip()
            if not token:
                continue
            if "-" in token:
                start_str, end_str = token.split("-", 1)
                try:
                    start = int(start_str.strip())
                    end = int(end_str.strip())
                except ValueError as exc:
                    raise ValueError(f"Invalid range token '{token}' in selector '{selector}'.") from exc
                if start < 0 or end < 0 or start > end:
                    raise ValueError(f"Invalid range token '{token}' in selector '{selector}'.")
                values.update(range(start, end + 1))
            else:
                try:
                    index = int(token)
                except ValueError as exc:
                    raise ValueError(f"Invalid index token '{token}' in selector '{selector}'.") from exc
                if index < 0:
                    raise ValueError(f"Negative index '{index}' is not allowed in selector '{selector}'.")
                values.add(index)
    else:
        raise TypeError(f"Unsupported selector type: {type(selector)!r}")

    for idx in values:
        if not isinstance(idx, int):
            raise TypeError(f"Selector index must be int, got {type(idx)!r}")
        if idx < 0:
            raise ValueError("Selector indices must be non-negative.")
    return values


@config
@dataclass
class DiffusionParallelConfig:
    """Configuration for diffusion model distributed execution."""

    pipeline_parallel_size: int = 1
    """Number of pipeline parallel stages."""

    data_parallel_size: int | None = None
    """Number of data parallel groups. Inferred from WORLD when omitted."""

    tensor_parallel_size: int = 1
    """Number of tensor parallel groups."""

    enable_expert_parallel: bool = False
    """Enable expert parallelism for MoE layers (TP is still used for non-MoE layers)."""

    sequence_parallel_size: int | None = None
    """Number of sequence parallel groups.
    sequence_parallel_size = ulysses_degree * ring_degree, or allgather_degree
    when AllGather-KV is enabled."""

    ulysses_degree: int = 1
    """Number of GPUs used for ulysses sequence parallelism."""

    ring_degree: int = 1
    """Number of GPUs used for ring sequence parallelism."""

    allgather_degree: int = 1
    """Number of GPUs used for AllGather-KV sequence parallelism (causal=False only)."""

    ulysses_mode: str = "strict"
    """Ulysses sequence-parallel mode.

    - "strict": Require divisibility constraints (fastest, default).
    - "advanced_uaa": Enable UAA ("Ulysses Anything Attention") to support
      uneven sequence lengths and non-divisible head counts.

    Note:
    - Ring attention does not support `attention_mask`, so models that rely on
      mask-based auto-padding are still incompatible with Ring.
    - When used in hybrid Ulysses+Ring, Ring requires consistent per-rank
      sequence shapes across the ring group.
    """

    ulysses_a2a_permute: bool = False
    """Use fused permute-free all-to-all for eligible strict Ulysses exchanges."""

    cfg_parallel_size: int = 1
    """Number of ranks used to execute guidance passes in parallel."""

    vae_patch_parallel_size: int = 1
    """Number of ranks used for VAE patch/tile parallelism (decode/encode)."""

    text_encoder_tp_size: int = 1
    """Number of ranks used to tensor-parallel shard the diffusion text encoder.

    Ranks are the first ``text_encoder_tp_size`` DiT ranks.  Defaults to 1,
    which keeps the encoder fully resident on the DiT main rank (historical
    behavior).  Values > 1 shard the Qwen3-VL encoder across the encoder TP
    ranks and run the encode with distributed collectives over the encoder
    process group.
    """

    vae_parallel_mode: str = "tile"
    """VAE parallel decode strategy.

    - "tile": Patch/tile parallel decode (default). Each rank decodes a subset
      of spatial tiles and the results are stitched on rank 0.
    - "batch": Decode complete images on different ranks, preserving their
      batch order. Requires AutoencoderKL/Flux2 and DP/PP/CFG sizes of 1.
    - "spatial_shard_height": Spatially-sharded decode that splits decoder
      feature maps along height and exchanges halo rows around spatial
      convolutions.
    - "spatial_shard_width": Same as "spatial_shard_height" but sharded along width.

    The "spatial_shard_*" modes are decode-only and currently require
    ``vae_patch_parallel_size`` to match the DiT group size; otherwise the VAE
    falls back to tile-parallel decode at runtime.
    """

    use_hsdp: bool = False
    """Enable Hybrid Sharded Data Parallel (HSDP) for model weight sharding."""

    mask_sp_padding: bool = False
    """If True, generate a boolean attention mask for zero-padded SP tokens
    when sequence length is not divisible by the SP world size. The mask
    routes attention through the varlen path (unpad→kernel→repad), which is
    correct but carries additional overhead. When False (default), padding
    tokens are left unmasked; since _shard_with_auto_pad always pads with
    zeros, their contribution to attention output is negligible."""

    hsdp_shard_size: int = -1
    """Number of GPUs to shard weights across within each replica group. -1 means auto-calculate."""

    hsdp_replicate_size: int = 1
    """Number of replica groups for HSDP. Each replica holds a full sharded copy."""

    @model_validator(mode="after")
    def _validate_parallel_config(self) -> Self:
        """Validates the config relationships among the parallel strategies."""
        assert self.pipeline_parallel_size > 0, "Pipeline parallel size must be > 0"
        assert self.data_parallel_size is None or self.data_parallel_size > 0, "Data parallel size must be > 0"
        assert self.tensor_parallel_size > 0, "Tensor parallel size must be > 0"
        assert self.sequence_parallel_size is not None
        assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0"
        assert self.ulysses_degree > 0, "Ulysses degree must be > 0"
        assert self.ring_degree > 0, "Ring degree must be > 0"
        assert self.allgather_degree > 0, "AllGather degree must be > 0"
        assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0"
        assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0"
        assert self.vae_parallel_mode in {"tile", "batch", "spatial_shard_height", "spatial_shard_width"}, (
            "vae_parallel_mode must be one of {'tile', 'batch', 'spatial_shard_height', 'spatial_shard_width'}, "
            f"but got {self.vae_parallel_mode!r}."
        )
        if self.vae_parallel_mode == "batch" and (
            self.data_parallel_size not in (None, 1) or self.pipeline_parallel_size != 1 or self.cfg_parallel_size != 1
        ):
            raise ValueError("VAE batch parallel decode requires DP, PP, and CFG parallel sizes to be 1")
        if self.allgather_degree > 1:
            assert self.ulysses_degree == 1 and self.ring_degree == 1, (
                "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. "
                f"Got ulysses_degree={self.ulysses_degree}, ring_degree={self.ring_degree}, "
                f"allgather_degree={self.allgather_degree}."
            )
        expected_sp_size = (
            self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
        )
        assert self.sequence_parallel_size == expected_sp_size, (
            f"Sequence parallel size must be {expected_sp_size}, but got {self.sequence_parallel_size}"
        )
        assert self.ulysses_mode in {"strict", "advanced_uaa"}, (
            f"ulysses_mode must be one of {{'strict','advanced_uaa'}}, but got {self.ulysses_mode!r}."
        )

        # Validate HSDP configuration
        if self.use_hsdp:
            assert self.hsdp_replicate_size > 0, "HSDP replicate size must be > 0"
            assert self.hsdp_shard_size > 0, "HSDP shard size must be > 0 (should be set in __post_init__)"
        return self

    def __post_init__(self) -> None:
        if self.sequence_parallel_size is None:
            self.sequence_parallel_size = (
                self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
            )

        # Until the runtime WORLD size is known, an omitted DP dimension means
        # one replica. OmniDiffusionConfig resolves it against num_gpus below.
        data_parallel_size = self.data_parallel_size or 1
        other_parallel_world_size = (
            self.pipeline_parallel_size
            * data_parallel_size
            * self.tensor_parallel_size
            * self.sequence_parallel_size
            * self.cfg_parallel_size
        )

        # Handle HSDP configuration
        # HSDP can work in two modes:
        # 1. Standalone: when other parallelism is all 1, HSDP determines world_size
        # 2. Combined: HSDP overlays on top of other parallelism
        if self.use_hsdp:
            incompatible = []
            if self.tensor_parallel_size > 1:
                incompatible.append("TP")
            if data_parallel_size > 1:
                incompatible.append("DP")
            if self.pipeline_parallel_size > 1:
                incompatible.append("PP")
            if self.enable_expert_parallel:
                incompatible.append("EP")
            if incompatible:
                raise ValueError("HSDP (FSDP2) is not compatible with " + ", ".join(incompatible))
            if self.hsdp_shard_size == -1:
                # Auto-calculate: use other_parallel_world_size as shard_size
                if self.hsdp_replicate_size <= 0:
                    raise ValueError("hsdp_replicate_size must be > 0")
                if other_parallel_world_size == 1:
                    raise ValueError(
                        "Cannot auto-calculate hsdp_shard_size when other parallelism is all 1. "
                        "Please specify hsdp_shard_size explicitly for standalone HSDP."
                    )
                if other_parallel_world_size % self.hsdp_replicate_size != 0:
                    raise ValueError(
                        f"Invalid HSDP configuration: replicate_size ({self.hsdp_replicate_size}) "
                        f"must evenly divide world_size ({other_parallel_world_size}) when shard_size is -1."
                    )
                self.hsdp_shard_size = other_parallel_world_size // self.hsdp_replicate_size
                self.world_size = other_parallel_world_size
            else:
                # Explicit shard_size: HSDP can work standalone or combined
                hsdp_world_size = self.hsdp_replicate_size * self.hsdp_shard_size
                if other_parallel_world_size == 1:
                    # Standalone HSDP: world_size is determined by HSDP
                    self.world_size = hsdp_world_size
                else:
                    # Combined: HSDP must match other parallelism world_size
                    if hsdp_world_size != other_parallel_world_size:
                        raise ValueError(
                            f"HSDP dimensions "
                            f"({self.hsdp_replicate_size} × {self.hsdp_shard_size} = {hsdp_world_size}) "
                            f"must equal world_size from other parallelism ({other_parallel_world_size})"
                        )
                    self.world_size = other_parallel_world_size
        else:
            self.world_size = other_parallel_world_size

    def resolve_data_parallel_size(self, world_size: int) -> int:
        """Resolve or validate diffusion DP against the actual WORLD size."""
        if world_size <= 0:
            raise ValueError(f"WORLD size must be > 0, but got {world_size}")

        if self.use_hsdp:
            if self.data_parallel_size not in (None, 1):
                raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1")
            expected_world_size = self.hsdp_replicate_size * self.hsdp_shard_size
            if world_size != expected_world_size:
                raise ValueError(
                    f"WORLD size ({world_size}) must equal HSDP size "
                    f"({self.hsdp_replicate_size} x {self.hsdp_shard_size} = {expected_world_size})"
                )
            self.data_parallel_size = 1
            self.world_size = world_size
            return 1

        assert self.sequence_parallel_size is not None
        non_dp_size = (
            self.pipeline_parallel_size
            * self.tensor_parallel_size
            * self.sequence_parallel_size
            * self.cfg_parallel_size
        )
        if world_size % non_dp_size != 0:
            raise ValueError(f"WORLD size ({world_size}) must be divisible by non-DP parallel size ({non_dp_size})")
        inferred_data_parallel_size = world_size // non_dp_size
        if self.vae_parallel_mode == "batch" and inferred_data_parallel_size != 1:
            raise ValueError("VAE batch parallel decode requires DP, PP, and CFG parallel sizes to be 1")
        if self.data_parallel_size is not None and self.data_parallel_size != inferred_data_parallel_size:
            raise ValueError(
                f"data_parallel_size ({self.data_parallel_size}) does not match WORLD-derived value "
                f"({inferred_data_parallel_size})"
            )
        self.data_parallel_size = inferred_data_parallel_size
        self.world_size = world_size
        return inferred_data_parallel_size

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "DiffusionParallelConfig":
        """
        Create DiffusionParallelConfig from a dictionary.

        Args:
            data: Dictionary containing parallel configuration parameters

        Returns:
            DiffusionParallelConfig instance with parameters set from dict
        """
        if not isinstance(data, dict):
            raise TypeError(f"Expected parallel config dict, got {type(data)!r}")
        return cls(**data)

    @classmethod
    def from_stage_overrides(cls, kwargs: Mapping[str, Any]) -> "DiffusionParallelConfig":
        """Resolve nested or flat stage inputs through the config's own defaults."""
        parallel_config = kwargs.get("parallel_config")
        if isinstance(parallel_config, cls):
            return parallel_config

        valid_fields = {item.name for item in fields(cls) if item.init}
        if isinstance(parallel_config, Mapping):
            return cls.from_dict(dict(parallel_config))
        elif parallel_config is not None and hasattr(parallel_config, "__dict__"):
            return cls.from_dict(dict(vars(parallel_config)))
        elif parallel_config is not None:
            raise TypeError(
                f"parallel_config must be a mapping or DiffusionParallelConfig, got {type(parallel_config).__name__}"
            )
        else:
            values = {key: value for key, value in kwargs.items() if key in valid_fields}

        return cls(**{key: value for key, value in values.items() if value is not None})


@dataclass
class TransformerConfig:
    """Container for raw transformer configuration dictionaries."""

    params: dict[str, Any] = field(default_factory=dict)
    quant_method: str | None = None
    quant_config: "QuantizationConfig | None" = None

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "TransformerConfig":
        if not isinstance(data, dict):
            raise TypeError(f"Expected transformer config dict, got {type(data)!r}")
        params = dict(data)  # copy to avoid mutating caller's dict

        quant_method: str | None = None
        quant_config: QuantizationConfig | None = None
        disk_qc = params.get("quantization_config")
        if isinstance(disk_qc, dict):
            raw_quant_method = disk_qc.get("quant_method", disk_qc.get("method"))
            quant_config = build_quant_config(disk_qc)
            if quant_config is not None:
                quant_method = raw_quant_method if raw_quant_method is not None else quant_config.get_name()

        return cls(params=params, quant_method=quant_method, quant_config=quant_config)

    def to_dict(self) -> dict[str, Any]:
        return dict(self.params)

    def get(self, key: str, default: Any | None = None) -> Any:
        return self.params.get(key, default)

    def __getattr__(self, item: str) -> Any:
        params = object.__getattribute__(self, "params")
        try:
            return params[item]
        except KeyError as exc:
            raise AttributeError(item) from exc


@dataclass
class DiffusionCacheConfig:
    """
    Configuration for cache adapters (TeaCache, cache-dit, MagCache, etc.).

    This dataclass provides a unified interface for cache configuration parameters.
    It can be initialized from a dictionary and accessed via attributes.

    Common parameters:
        - TeaCache: rel_l1_thresh, coefficients (optional)
        - cache-dit: Fn_compute_blocks, Bn_compute_blocks, max_warmup_steps,
                    residual_diff_threshold, enable_taylorseer, taylorseer_order,
                    scm_steps_mask_policy, scm_steps_policy
        - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio,
                    mag_ratios, mag_calibrate
        - SeaCache: sea_threshold, sea_residual_order,
                    sea_max_consecutive_cached, sea_power_exp
        - step_cache: step_cache_dit_enabled, velocity_sim_thresholds,
                          velocity_skip_countdowns, step_cache_dit_min_history

    Example:
        >>> # From dict (user-facing API) - partial config uses defaults for missing keys
        >>> config = DiffusionCacheConfig.from_dict({"rel_l1_thresh": 0.3})
        >>> # Access via attribute
        >>> print(config.rel_l1_thresh)  # 0.3 (from dict)
        >>> print(config.Fn_compute_blocks)  # 8 (default)
        >>> # Empty dict defers model-specific defaults to the TeaCache backend
        >>> default_config = DiffusionCacheConfig.from_dict({})
        >>> print(default_config.rel_l1_thresh)  # None
    """

    # TeaCache parameters [tea_cache only]
    # None defers to the model-specific TeaCache default (0.2 fallback).
    rel_l1_thresh: float | None = None
    coefficients: list[float] | None = None  # Uses model-specific defaults if None

    # SeaCache parameters [sea_cache only]
    sea_threshold: float = 0.25
    sea_residual_order: int = 1
    sea_max_consecutive_cached: int = 2
    sea_power_exp: float = 3.0

    # MagCache parameters [mag_cache only]
    # Default: 0.24 threshold for accumulated magnitude error
    mag_threshold: float = 0.24
    # Default: 5 maximum consecutive skip steps (K)
    mag_max_skip_steps: int = 5
    # Default: 0.1 fraction of initial steps where skipping is disabled (stability)
    mag_retention_ratio: float = 0.1
    # Default: None magnitude ratios (model-specific, required for inference)
    mag_ratios: list[float] | None = None
    # Default: False calibration mode (computes mag_ratios on first run)
    mag_calibrate: bool = False

    # cache-dit parameters [cache-dit only]
    # Default: 1 forward compute block (optimized for single-transformer models)
    # Use 1 as default instead of cache-dit's 8, optimized for single-transformer models
    # This provides better performance while maintaining quality for most use cases
    Fn_compute_blocks: int = 1
    # Default: 0 backward compute blocks (no fusion by default)
    Bn_compute_blocks: int = 0
    # Default: 4 warmup steps (optimized for few-step distilled models like Z-Image with 8 steps)
    # Use 4 as default warmup steps instead of 8 in cache-dit, making DBCache work
    # for few-step distilled models (e.g., Z-Image with 8 steps)
    max_warmup_steps: int = 4
    # Default: -1 (unlimited cached steps) - DBCache disables caching when previous cached steps exceed this value
    # to prevent precision degradation. Set to -1 for unlimited caching (cache-dit default).
    max_cached_steps: int = -1
    # Default: 0.24 residual difference threshold (higher for more aggressive caching)
    # Use a relatively higher residual diff threshold (0.24) as default to allow more
    # aggressive caching. This is safe because we have max_continuous_cached_steps limit.
    # Without this limit, a lower threshold like 0.12 would be needed.
    residual_diff_threshold: float = 0.24
    # Default: Limit consecutive cached steps to 3 to prevent precision degradation
    # This allows us to use a higher residual_diff_threshold for more aggressive caching
    max_continuous_cached_steps: int = 3
    # Default: Disable TaylorSeer (not suitable for few-step distilled models)
    # TaylorSeer is not suitable for few-step distilled models, so we disable it by default.
    # References:
    # - From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers
    # - Forecast then Calibrate: Feature Caching as ODE for Efficient Diffusion Transformers
    enable_taylorseer: bool = False
    # Default: 1st order TaylorSeer polynomial
    taylorseer_order: int = 1
    # Default: None SCM mask policy (disabled by default)
    scm_steps_mask_policy: str | None = None
    # Default: "dynamic" steps policy for adaptive caching
    scm_steps_policy: str = "dynamic"
    # Used by cache-dit for scm mask generation. If this value changes during inference,
    # we will re-generate the scm mask and refresh the cache context.
    num_inference_steps: int | None = None
    # Force refresh the cache at a specific step index hint, useful for models like
    # GLM-Image (image preprocessing step in editing mode).
    force_refresh_step_hint: int | None = None
    # Policy for force refresh: "once" refreshes only at the hint step,
    # "repeat" refreshes every force_refresh_step_hint steps.
    force_refresh_step_policy: str = "once"

    # step_cache parameters [step_cache only] — DreamZero velocity schedule
    step_cache_dit_enabled: bool = True
    velocity_sim_thresholds: list[float] = field(default_factory=lambda: [0.95, 0.93])
    velocity_skip_countdowns: list[int] = field(default_factory=lambda: [4, 2])
    step_cache_dit_min_history: int = 2
    step_cache_dit_max_history: int = 2

    # Additional parameters that may be passed but not explicitly defined
    _extra_params: dict[str, Any] = field(default_factory=dict, repr=False)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "DiffusionCacheConfig":
        """
        Create DiffusionCacheConfig from a dictionary.

        Args:
            data: Dictionary containing cache configuration parameters

        Returns:
            DiffusionCacheConfig instance with parameters set from dict
        """
        if not isinstance(data, dict):
            raise TypeError(f"Expected cache config dict, got {type(data)!r}")

        # Get all dataclass field names automatically
        field_names = {f.name for f in fields(cls)}

        # Extract parameters that match dataclass fields (excluding private fields)
        known_params = {k: v for k, v in data.items() if k in field_names and not k.startswith("_")}

        # Store extra parameters
        extra_params = {k: v for k, v in data.items() if k not in field_names}

        # Create instance with known params (missing ones will use defaults)
        # Then update _extra_params after creation since it's a private field
        instance = cls(**known_params, _extra_params=extra_params)
        return instance

    def __getattr__(self, item: str) -> Any:
        """
        Allow access to extra parameters via attribute access.

        This enables accessing parameters that weren't explicitly defined
        in the dataclass fields but were passed in the dict.
        """
        if item == "_extra_params" or item.startswith("_"):
            return object.__getattribute__(self, item)

        extra = object.__getattribute__(self, "_extra_params")
        if item in extra:
            return extra[item]

        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'")


def resolve_model_class_name(
    model: str | None,
    diffusion_load_format: str = "default",
    revision: str | None = None,
) -> str | None:
    """Resolve the diffusion pipeline class name from the model config.

    Read-only counterpart of ``OmniDiffusionConfig.enrich_config``, safe to call
    client-side. Returns ``None`` if the pipeline can't be determined.
    """
    from vllm.transformers_utils.config import get_hf_file_to_dict

    from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index, resolve_native_diffusion_model_class

    if not model:
        return None
    native_model_class = resolve_native_diffusion_model_class(model)
    if native_model_class is not None:
        return native_model_class

    is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"}

    # Diffusers models: read _class_name from the pipeline index. Missing
    # local paths can otherwise be interpreted as invalid Hub repo IDs.
    try:
        model_index = get_diffusion_model_index(model, revision=revision)
    except Exception:
        model_index = None
    if model_index is not None:
        return model_index.get("_class_name")
    if diffusion_load_format == "diffusers":
        return "DiffusersAdapterPipeline"

    # Other models: map model_type / architecture from config.json.
    try:
        cfg = get_hf_file_to_dict("config.json", model, revision=revision) or {}
    except Exception:
        cfg = {}
    model_type = cfg.get("model_type")
    architectures = cfg.get("architectures") or []

    from vllm_omni.diffusion.utils.hf_utils import _looks_like_hidream_o1

    if _looks_like_hidream_o1(model, cfg):
        return "HiDreamO1ImagePipeline"
    if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
        return "BagelPipeline"
    if (
        model_type == "lance"
        or "LancePipeline" in architectures
        or cfg.get("model_name") == "Lance"
        or is_lance_subfolder
    ):
        return "LancePipeline"
    if model_type == "neo_chat":
        return "SenseNovaU1Pipeline"
    if "BailingMM2NativeForConditionalGeneration" in architectures or model_type in (
        "bailingmm_moe_v2_lite",
        "ming_flash_omni",
        "ming_flash_omni_thinker",
    ):
        return "MingImagePipeline"
    if model_type == "nextstep":
        return "NextStep11Pipeline"
    if model_type == "s2v":
        return "WanS2VPipeline"
    if model_type == "vla":
        from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero

        return "DreamZeroPipeline" if _looks_like_dreamzero(model) else None
    if len(architectures) == 1:
        return architectures[0]
    return None


def uses_diffusers_adapter(od_config: object) -> bool:
    """Return whether execution uses the Diffusers adapter backend.

    A custom pipeline takes precedence over ``diffusion_load_format`` in the
    worker, so adapter-specific behavior must only apply when no custom
    pipeline override is configured.
    """
    return (
        getattr(od_config, "custom_pipeline_args", None) is None
        and getattr(od_config, "diffusion_load_format", "default") == "diffusers"
    )


@dataclass
class VideoOutputTransportConfig:
    enable_device_postprocess: bool = False
    enable_registered_shm: bool = False
    enable_borrowed_frames: bool = False

    def __post_init__(self) -> None:
        if not isinstance(self.enable_device_postprocess, bool):
            raise TypeError("enable_device_postprocess must be a bool")
        if not isinstance(self.enable_registered_shm, bool):
            raise TypeError("enable_registered_shm must be a bool")
        if not isinstance(self.enable_borrowed_frames, bool):
            raise TypeError("enable_borrowed_frames must be a bool")


@dataclass
class OmniDiffusionConfig:
    # Model and path configuration (for convenience)
    stage_id: int = 0

    model: str | None = None

    model_class_name: str | None = None

    # Optional model-defined startup task. Pipelines may use this to select
    # task-specific components or weights before serving requests.
    task_type: str | None = None

    dtype: torch.dtype = torch.bfloat16

    model_config: dict[str, Any] = field(default_factory=dict)
    tf_model_config: TransformerConfig = field(default_factory=TransformerConfig)

    # Attention
    diffusion_attention_config: "AttentionConfig" = field(default_factory=lambda: AttentionConfig())
    fa_deterministic: bool = False

    # Running mode
    # mode: ExecutionMode = ExecutionMode.INFERENCE

    # Workload type
    # workload_type: WorkloadType = WorkloadType.T2V

    # Cache strategy (legacy)
    cache_strategy: str = "none"
    parallel_config: DiffusionParallelConfig = field(default_factory=DiffusionParallelConfig)

    # Cache backend configuration (NEW)
    cache_backend: str | None = "none"  # "tea_cache", "deep_cache", etc.
    cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict)
    video_output_transport: VideoOutputTransportConfig = field(default_factory=VideoOutputTransportConfig)
    enable_cache_dit_summary: bool = False

    # Prompt-embedding cache. When enabled, ``DiffusionModelRunner`` wraps the
    # pipeline's ``encode_prompt`` so repeated calls with identical prompt
    # arguments (e.g. GRPO rollouts that sample the same prompt many times
    # with different seeds) reuse the text-encoder output instead of re-running
    # it. Safe against inputs that cannot be hashed (tensors, PIL images):
    # those calls transparently bypass the cache.
    enable_prompt_embed_cache: bool = False
    prompt_embed_cache_size: int = 32

    # Opt-in: route per-session world-model state through the shared
    # SessionStateManager (RFC #4480) instead of the model's bespoke cache.
    # Default off; the bespoke path remains the default. A *set*
    # OMNI_DIFFUSION_SESSION_STATE_MANAGER environment variable overrides this
    # in both directions. See docs/features/session_state_manager.md.
    enable_session_state_manager: bool = False

    # Distributed executor backend. ``None`` means auto: ``"uni"`` at
    # ``num_gpus == 1``, ``"mp"`` otherwise. Explicit ``"mp"`` keeps a
    # worker subprocess (process isolation and RPC timeouts).
    distributed_executor_backend: str | None = None
    nccl_port: int | None = None

    # Engine backend selection, resolved by ``DiffusionEngine.resolve_engine_class``
    # (mirrors ``DiffusionExecutor.get_class``). Config files use a string:
    # "default" -> DiffusionEngine, or an import path (set e.g. by a deploy
    # config). Programmatic callers may pass a DiffusionEngine subclass
    # directly — hence ``str | type`` (structured-config mirrors should expose
    # the string form only).
    engine_backend: str | type = "default"

    # Local Diffusion KV ownership and cache-layout mode.
    diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
    # Reuse block-aligned, immutable diffusion prefixes across requests when
    # Scheduler-owned paged KV is active.
    enable_prefix_caching: bool = False
    # Maximum number of native BlockTable rows one public request can own
    # (sequences plus independent contexts). The model adapter defines it.
    diffusion_kv_max_rows_per_request: int | None = None

    # Optional override for the diffusion model runner class (import path).
    # Precedence in the worker: this override > the runner declared by the
    # selected engine class (``default_diffusion_model_runner_cls``) > the
    # platform default. Never mutated by engines.
    diffusion_model_runner_cls: str | None = None

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

    num_gpus: int | None = None

    dist_timeout: int | None = None  # timeout for torch.distributed

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

    # LoRA parameters
    lora_path: str | list[str] | None = None
    lora_scale: float | list[float] = 1.0
    lora_backend: LoRABackend = LoRABackend.PEFT  # available choices: ["peft", "distill"]
    max_cpu_loras: int | None = None

    output_type: str = "pil"

    # CPU offload parameters. Keep the public mapping raw so stage configs can
    # serialize it across processes; __post_init__ validates it once and caches
    # the internal typed resolution used at runtime.
    diffusion_offload_config: dict[str, Any] | None = None
    # Compatibility aliases. Some model-specific legacy stage lifecycles are
    # intentionally broader than the compact dit/text_encoder selector.
    # When enabled, DiT and encoders swap GPU access (mutual exclusion):
    # - Text encoders run on GPU while DiT is on CPU
    # - DiT runs on GPU while encoders are on CPU
    enable_cpu_offload: bool = False
    # Layer-wise offloading (block-level offloading) parameters
    enable_layerwise_offload: bool = False
    # Distributed layer-wise offloading with H2D + AllGather overlap (RFC-1)
    enable_distributed_layerwise_offload: bool = False
    # If True: shard weights 1/dp_size + AllGather (saves CPU memory, requires
    # concurrent requests in DP mode). If False: each rank streams the standard
    # loader's rank-local tensors (including TP-local shards) via H2D only.
    # This avoids AllGather synchronization, while host memory follows the
    # loader's existing rank-local layout instead of adding a second DP shard.
    dlo_use_allgather: bool = True
    # Leading main-DiT blocks kept resident by distributed layerwise offload.
    dlo_resident_layers: int = 0
    # Final-layout Host Weight Runtime policy. The loader only activates this
    # for eligible no-AllGather DLO; all other configurations preserve their
    # existing loader/storage path.
    host_weight_runtime_mode: str = "disabled"
    host_weight_runtime_root: str | None = None
    # Optional per-worker ceiling for registering final-layout HWR mappings.
    # Zero adds no ceiling; pin_cpu_memory controls whether registration is tried.
    dlo_host_registration_limit_gib: float = 0.0

    pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading

    # VAE memory optimization parameters
    vae_use_slicing: bool = False
    vae_use_tiling: bool = False
    # Wan VAE decoder fast path. ``"lossless"`` installs the bit-exact fused
    # kernels on every diffusers Wan VAE, ``"channels_last"`` additionally
    # converts decoder convolution weights to channels-last memory format and
    # fuses RMSNorm+SiLU (faster, not bit-exact), ``"off"`` keeps the reference
    # diffusers implementation.
    vae_fast_path: str = "lossless"

    # STA (Sliding Tile Attention) parameters
    mask_strategy_file_path: str | None = None
    # STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
    skip_time_steps: int = 15

    # MoE kernel backend selection
    moe_backend: str = "auto"

    # Compilation
    enforce_eager: bool = False
    # Controls the generic compilation path used when a pipeline does not
    # provide its own setup_compile() implementation.
    diffusion_compile_granularity: str = "regional"
    diffusion_compile_dynamic: bool = True

    # Parallel weight loading (for faster diffusion model startup)
    enable_multithread_weight_load: bool = True
    enable_broadcast_weight_load: bool = False
    num_weight_load_threads: int = 4

    # Shard meta parameters before reading rank-local HF safetensors slices.
    hsdp_weight_load_strategy: str = "full"

    # Enable sleep mode
    enable_sleep_mode: bool = False

    disable_autocast: bool = False

    # VSA parameters
    VSA_sparsity: float = 0.0  # inference/validation sparsity

    # V-MoBA parameters
    moba_config_path: str | None = None
    # moba_config: dict[str, Any] = field(default_factory=dict)

    # Master port for distributed inference
    # TODO: do not hard code
    master_port: int | None = None

    # Worker extension class for custom functionality
    worker_extension_cls: str | None = None

    # Internal transport of explicit stage runtime.env to remote Ray actors.
    ray_worker_env: dict[str, str] = field(default_factory=dict, init=False, repr=False)

    # Custom pipeline arguments for custom pipelines
    custom_pipeline_args: dict[str, Any] | None = None

    # Diffusion model loading format
    # "default", "custom_pipeline", "dummy", "diffusers" (HF diffusers adapter)
    diffusion_load_format: str = "default"

    # Diffusers adapter kwargs
    # kwargs forwarded to DiffusionPipeline.from_pretrained()
    diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict)
    # kwargs forwarded to pipeline.__call__()
    diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict)
    # Actual diffusers pipeline object (to determine inputs of the dummy run)
    diffusers_pipeline_cls: type[diffusers.DiffusionPipeline] | None = None  # pyright: ignore[reportPrivateImportUsage]

    # http server endpoint config, would be ignored in local mode
    host: str | None = None
    port: int | None = None

    scheduler_port: int = 5555

    # Stage verification
    enable_stage_verification: bool = True

    # 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,
            "text_encoder": True,
            "vae_encoder": True,
        }
    )
    override_transformer_cls_name: str | None = None

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

    # MoE parameters used by Wan2.2
    boundary_ratio: float | None = None
    # Scheduler flow_shift for Wan2.2 (12.0 for 480p, 5.0 for 720p)
    flow_shift: float | None = None

    # Support multi-image inputs and expose any model-specific request limit
    # through a generic config field so serving code stays model-agnostic.
    supports_multimodal_inputs: bool = False
    max_multimodal_image_inputs: int | None = None
    supports_mixed_reference_inputs: bool = False

    log_level: str = "info"

    # Omni configuration (injected from stage config)
    omni_kv_config: dict[str, Any] = field(default_factory=dict)
    # Native vLLM KV connector configuration. PR0 only assembles the connector;
    # the native page data path is owned by the follow-up landing PR.
    kv_transfer_config: "KVTransferConfig | None" = None
    additional_config: dict[str, Any] = field(default_factory=dict)

    profiler_config: "ProfilerConfig | dict[str, Any] | None" = None

    # Model-specific function for collecting CFG KV caches (set at runtime)
    cfg_kv_collect_func: Any | None = None

    # Conditioning keys fetched from the upstream stage over the omni connector
    # rather than carried inline through the orchestrator. Empty disables the
    # worker-side connector receive path.
    stage_input_payload_keys: tuple[str, ...] = ()

    # Keys handed to the next stage over the omni connector. Empty disables the
    # worker-side connector send path.
    stage_output_payload_keys: tuple[str, ...] = ()

    # Quantization: str method name, dict config, QuantizationConfig, or None.
    # str is resolved to {"method": <str>} internally.
    # Per-component: {"transformer": {"method": "fp8"}, "vae": None}
    quantization_config: str | QuantizationConfig | dict[str, Any] | None = None
    # Internal provenance, retained across config projection and worker transport.
    quantization_config_is_auto_detected: bool = False
    # Explicit runtime override for ModelOpt FP8 diffusion checkpoints. This
    # does not enable FP8 by itself; it only selects CUTLASS once the checkpoint
    # has already resolved to vLLM's ModelOpt FP8 linear method.
    force_cutlass_fp8: bool = False

    # Runtime diffusion attention method (not vLLM's --kv-cache-dtype for AR models).
    # None/"auto" keeps native dtype; other values are validated by the selected backend.
    diffusion_kv_cache_dtype: str | None = None
    # Optional skip selectors for KV-cache quantization. Format: "0-9,20,25-30".
    # Listed steps/layers skip quantization; others keep quantized execution.
    diffusion_kv_cache_skip_steps: str | None = None
    diffusion_kv_cache_skip_layers: str | None = None
    diffusion_kv_cache_skip_step_indices: set[int] | None = None
    diffusion_kv_cache_skip_layer_indices: set[int] | None = None

    # Diffusion pipeline Profiling config
    enable_diffusion_pipeline_profiler: bool = False

    # Step mode settings
    step_execution: bool = False

    # Streaming mode settings
    streaming_output: bool = False  # Start (video) generation with initial prompt, but streaming output in chunks

    # Maximum number of sequences to generate in a batch
    max_num_seqs: int = 1

    # Native vLLM KV-cache sizing inputs used only by paged_scheduler mode.
    # These mirror the structured stage cache/scheduler config until all
    # diffusion runtime construction consumes VllmOmniConfig directly.
    kv_cache_memory_bytes: int | None = None
    gpu_memory_utilization: float = 0.9
    max_num_batched_tokens: int | None = None
    max_model_len: int | None = None

    # Request-mode batch admission: wait briefly for compatible requests to
    # accumulate in the scheduler waiting queue before the first schedule() of
    # a wave.  Improves fused forward batch sizes under bursty HTTP ingress.
    # 0 disables admission (default; no added latency).
    request_batch_max_wait_ms: float = 0.0

    # Supplementary model specific parameters
    extras: dict[str, Any] = Field(default_factory=dict)

    @property
    def is_moe(self) -> bool:
        num_experts = self.tf_model_config.get("num_experts", None)
        if not isinstance(num_experts, (list, tuple, int)):
            return False
        if isinstance(num_experts, int):
            return num_experts > 0

        if isinstance(num_experts, (list, tuple)):
            return any(isinstance(n, int) and n > 0 for n in num_experts)

        return False

    def _resolve_master_port(self) -> int:
        """Resolve torch.distributed master port without unnecessary random jitter.

        Precedence:
        1. ``MASTER_PORT`` environment variable (set by orchestrators for multi-replica launch).
        2. Explicit ``master_port`` passed at construction time.
        3. An OS-assigned ephemeral port when neither is provided.
        """
        from vllm.utils.network_utils import get_open_port

        from vllm_omni.diffusion import envs

        env_port = envs.MASTER_PORT
        if env_port is not None:
            return self.settle_port(env_port, port_inc=37)
        if self.master_port is not None:
            return self.settle_port(self.master_port, port_inc=37)
        return self.settle_port(get_open_port(), port_inc=37)

    def settle_port(self, port: int, port_inc: int = 42, max_attempts: int = 100) -> int:
        """
        Find an available port with retry logic.

        Args:
            port: Initial port to check
            port_inc: Port increment for each attempt
            max_attempts: Maximum number of attempts to find an available port

        Returns:
            An available port number

        Raises:
            RuntimeError: If no available port is found after max_attempts
        """
        attempts = 0
        original_port = port

        while attempts < max_attempts:
            if 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 (started from port {original_port})"
        )

    def __post_init__(self):
        if self.hsdp_weight_load_strategy not in {"full", "pre_sharded"}:
            raise ValueError(
                f"hsdp_weight_load_strategy must be 'full' or 'pre_sharded', got {self.hsdp_weight_load_strategy!r}"
            )
        from vllm_omni.diffusion.offloader.config import (
            OffloadStrategy,
            materialize_legacy_offload_flags,
        )

        self.stage_input_payload_keys = tuple(self.stage_input_payload_keys)
        self.stage_output_payload_keys = tuple(self.stage_output_payload_keys)
        if self.vae_fast_path not in VAE_FAST_PATH_LEVELS:
            raise ValueError(f"vae_fast_path must be one of {list(VAE_FAST_PATH_LEVELS)}, got {self.vae_fast_path!r}")
        if self.diffusion_compile_granularity not in {"regional", "full"}:
            raise ValueError(
                "diffusion_compile_granularity must be 'regional' or 'full', "
                f"got {self.diffusion_compile_granularity!r}"
            )
        if not isinstance(self.diffusion_compile_dynamic, bool):
            raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
        self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
        if not isinstance(self.enable_prefix_caching, bool):
            raise TypeError("enable_prefix_caching must be a bool")
        if self.enable_prefix_caching and self.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
            raise ValueError(
                "enable_prefix_caching=True requires diffusion_kv_mode='paged_scheduler'; "
                "set diffusion_kv_mode='paged_scheduler' or disable enable_prefix_caching"
            )
        if self.enable_prefix_caching and self.enable_sleep_mode:
            raise ValueError(
                "Diffusion prefix caching cannot be combined with sleep mode: "
                "sleep discards KV pages without invalidating cached prefixes; "
                "disable enable_prefix_caching or enable_sleep_mode"
            )
        if self.diffusion_kv_max_rows_per_request is not None and (
            type(self.diffusion_kv_max_rows_per_request) is not int or self.diffusion_kv_max_rows_per_request <= 0
        ):
            raise ValueError("diffusion_kv_max_rows_per_request must be a positive integer when set")
        if (
            self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
            and self.diffusion_kv_max_rows_per_request is None
        ):
            raise ValueError("paged_scheduler requires diffusion_kv_max_rows_per_request to be set")
        if self.kv_cache_memory_bytes is not None and self.kv_cache_memory_bytes < 0:
            raise ValueError("kv_cache_memory_bytes must be non-negative")
        if not 0.0 < self.gpu_memory_utilization <= 1.0:
            raise ValueError("gpu_memory_utilization must be in (0, 1]")
        if self.max_num_batched_tokens is not None and self.max_num_batched_tokens <= 0:
            raise ValueError("max_num_batched_tokens must be positive")
        if self.max_model_len is not None and self.max_model_len != -1 and self.max_model_len <= 0:
            raise ValueError("max_model_len must be positive or -1")

        if self.omni_kv_config is None:
            self.omni_kv_config = {}
        elif isinstance(self.omni_kv_config, Mapping):
            self.omni_kv_config = dict(self.omni_kv_config)
        else:
            raise TypeError("omni_kv_config must be a mapping")
        if self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER and self.omni_kv_config.get(
            "need_recv_cache", False
        ):
            raise ValueError(
                "paged_scheduler Diffusion KV does not support imported AR KV; "
                "disable need_recv_cache until connector-aware import is implemented"
            )

        if self.kv_transfer_config is not None:
            from vllm_omni.diffusion.diffusion_kv.kv_connector import parse_kv_transfer_config

            if any(self.omni_kv_config.get(name, False) for name in ("need_send_cache", "need_recv_cache")):
                raise ValueError(
                    "native kv_transfer_config cannot be combined with legacy omni_kv_config transfer; "
                    "configure exactly one KV transfer path"
                )
            if self.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
                raise ValueError("native kv_transfer_config requires diffusion_kv_mode='paged_scheduler'")
            self.kv_transfer_config = parse_kv_transfer_config(self.kv_transfer_config)
            if self.enable_prefix_caching and self.kv_transfer_config is not None:
                raise ValueError(
                    "Diffusion prefix caching cannot be combined with native kv_transfer_config; "
                    "disable enable_prefix_caching for AR KV import or remove kv_transfer_config for local DiT reuse"
                )
            if self.enable_sleep_mode:
                raise ValueError("Native KV transfer does not support sleep mode: registered pages must remain mapped")

        self.master_port = self._resolve_master_port()
        self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0)
        if not math.isfinite(self.request_batch_max_wait_ms) or self.request_batch_max_wait_ms < 0:
            raise ValueError(
                f"request_batch_max_wait_ms must be a finite non-negative number, got {self.request_batch_max_wait_ms}."
            )

        if isinstance(self.profiler_config, dict):
            from vllm.config import ProfilerConfig

            self.profiler_config = ProfilerConfig(**self.profiler_config)

        if self.additional_config is None:
            self.additional_config = {}
        elif isinstance(self.additional_config, Mapping):
            self.additional_config = dict(self.additional_config)
        else:
            raise TypeError(f"additional_config must be a mapping or None, got {type(self.additional_config)!r}")

        # Convert parallel_config dict/DictConfig to DiffusionParallelConfig
        # Use Mapping to handle both plain dicts and OmegaConf DictConfig
        if isinstance(self.parallel_config, Mapping):
            self.parallel_config = DiffusionParallelConfig.from_dict(dict(self.parallel_config))
        elif not isinstance(self.parallel_config, DiffusionParallelConfig):
            self.parallel_config = DiffusionParallelConfig()

        if self.num_gpus is None:
            if self.parallel_config is not None:
                self.num_gpus = self.parallel_config.world_size
            else:
                self.num_gpus = 1

        self.parallel_config.resolve_data_parallel_size(self.num_gpus)
        # Resolve offload only after DP/SP normalization so cached policy
        # validation observes the actual execution topology.
        offload_strategy = materialize_legacy_offload_flags(self)

        if self.diffusion_compile_granularity == "full":
            incompatible_features = []
            assert self.parallel_config.sequence_parallel_size is not None
            if self.parallel_config.use_hsdp:
                incompatible_features.append("HSDP")
            if self.parallel_config.sequence_parallel_size > 1:
                incompatible_features.append("sequence parallelism")
            if offload_strategy is not OffloadStrategy.NONE:
                incompatible_features.append(
                    {
                        OffloadStrategy.MODEL_LEVEL: "CPU offload",
                        OffloadStrategy.LAYER_WISE: "layerwise offload",
                        OffloadStrategy.DISTRIBUTED_LAYER_WISE: "distributed layerwise offload",
                    }[offload_strategy]
                )
            if incompatible_features:
                features = ", ".join(incompatible_features)
                raise ValueError(
                    "diffusion_compile_granularity='full' is incompatible with "
                    f"{features}; use 'regional' compilation instead"
                )

        # Convert string dtype to torch.dtype if needed
        if isinstance(self.dtype, str):
            dtype_map = {
                "auto": torch.bfloat16,
                "bfloat16": torch.bfloat16,
                "bf16": torch.bfloat16,
                "float16": torch.float16,
                "fp16": torch.float16,
                "half": torch.float16,
                "float32": torch.float32,
                "fp32": torch.float32,
                "float": torch.float32,
            }
            dtype_lower = self.dtype.lower()
            if dtype_lower in dtype_map:
                self.dtype = dtype_map[dtype_lower]
            else:
                logger.warning(f"Unknown dtype string '{self.dtype}', defaulting to bfloat16")
                self.dtype = torch.bfloat16

        # Convert cache_config dict to DiffusionCacheConfig if needed
        if isinstance(self.cache_config, dict):
            self.cache_config = DiffusionCacheConfig.from_dict(self.cache_config)
        elif not isinstance(self.cache_config, DiffusionCacheConfig):
            # If it's neither dict nor DiffusionCacheConfig, convert to empty config
            self.cache_config = DiffusionCacheConfig()

        if self.video_output_transport is None:
            self.video_output_transport = VideoOutputTransportConfig()
        elif isinstance(self.video_output_transport, Mapping):
            self.video_output_transport = VideoOutputTransportConfig(**dict(self.video_output_transport))
        elif not isinstance(self.video_output_transport, VideoOutputTransportConfig):
            raise TypeError("video_output_transport must be a VideoOutputTransportConfig or mapping")

        # Auto-detect quantization from TransformerConfig if not explicitly set.
        # This covers the case where tf_model_config is passed at construction
        # time. For late (post-construction) assignment, callers should use
        # set_tf_model_config() which propagates quant_config automatically.
        self._propagate_quantization_from_tf_config(self.tf_model_config)

        # Resolve quantization_config: str/dict -> QuantizationConfig via build_quant_config.
        if self.quantization_config is not None:
            if isinstance(self.quantization_config, QuantizationConfig):
                pass  # Already built
            elif isinstance(self.quantization_config, str):
                self.quantization_config = build_quant_config(self.quantization_config)
            elif isinstance(self.quantization_config, Mapping):
                self.quantization_config = build_quant_config(dict(self.quantization_config))
            else:
                raise TypeError(
                    f"quantization_config must be str, dict, QuantizationConfig, or None, "
                    f"got {type(self.quantization_config)!r}"
                )

        # Match vLLM's config flow: parse entrypoint shorthands before the
        # config object is built, and keep a single runtime truth source.
        self.diffusion_attention_config = build_attention_config(self.diffusion_attention_config)
        self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps)
        self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers)

        if self.max_cpu_loras is None:
            self.max_cpu_loras = 1
        elif self.max_cpu_loras < 1:
            raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA")

        validate_host_weight_runtime_options(
            mode=self.host_weight_runtime_mode,
            root=self.host_weight_runtime_root,
        )
        self.dlo_host_registration_limit_gib = validate_dlo_host_registration_options(
            limit_gib=self.dlo_host_registration_limit_gib,
            enable_dlo=self.enable_distributed_layerwise_offload,
            use_allgather=self.dlo_use_allgather,
            hwr_mode=self.host_weight_runtime_mode,
        )

        if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
            raise ValueError(
                "diffusers_load_kwargs and diffusers_call_kwargs are only "
                "valid together with diffusion_load_format=diffusers"
            )

        # when use hf offline, replace model to local model path
        # align with ar stage behavior, see vllm/engine/arg_utils.py
        if huggingface_hub.constants.HF_HUB_OFFLINE and self.model:
            model_id = self.model
            self.model = get_model_path(self.model, self.revision)
            if model_id != self.model:
                logger.info(
                    "HF_HUB_OFFLINE is True, replace model_id [%s] to model_path [%s]",
                    model_id,
                    self.model,
                )

    @property
    def is_single_file(self) -> bool:
        return isinstance(self.model, str) and os.path.isfile(self.model)

    def _propagate_quantization_from_tf_config(self, tf_config: "TransformerConfig") -> None:
        if tf_config.quant_config is None:
            return

        is_checkpoint_fp8 = bool(getattr(tf_config.quant_config, "is_checkpoint_fp8_serialized", False))
        is_checkpoint_nvfp4 = bool(getattr(tf_config.quant_config, "is_checkpoint_nvfp4_serialized", False))
        should_use_checkpoint_config = (
            self.quantization_config is None
            or (is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config))
            or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config))
        )
        if should_use_checkpoint_config:
            if self.quantization_config is None:
                self.quantization_config_is_auto_detected = True
            self.quantization_config = tf_config.quant_config
            logger.info(
                "Auto-detected quantization '%s' from model config",
                tf_config.quant_method,
            )

    @staticmethod
    def _is_generic_fp8_quant_config(quant_config: object) -> bool:
        if isinstance(quant_config, str):
            return quant_config.lower() == "fp8"
        if isinstance(quant_config, Mapping):
            method = quant_config.get("method", quant_config.get("quant_method"))
            return isinstance(method, str) and method.lower() == "fp8"
        if hasattr(quant_config, "get_name"):
            return quant_config.get_name() == "fp8"
        return False

    @staticmethod
    def _is_generic_nvfp4_quant_config(quant_config: object) -> bool:
        if isinstance(quant_config, str):
            return quant_config.lower() in {"fp4", "nvfp4", "modelopt_fp4"}
        if isinstance(quant_config, Mapping):
            method = quant_config.get("method", quant_config.get("quant_method"))
            return isinstance(method, str) and method.lower() in {"fp4", "nvfp4", "modelopt_fp4"}
        if hasattr(quant_config, "get_name"):
            return quant_config.get_name() == "modelopt_fp4"
        return False

    def set_tf_model_config(self, tf_config: "TransformerConfig") -> None:
        """Assign `tf_model_config` and propagate quantization if detected.

        In the normal startup flow `OmniDiffusionConfig` is created
        *before* the transformer `config.json` is loaded from disk, so
        `__post_init__` sees an empty `TransformerConfig`.  Callers
        that load the config later should use this method instead of bare
        assignment so that an embedded `quant_config` is propagated to
        `self.quantization_config` automatically.

        Args:
            tf_config: Transformer configuration, typically built via
                `TransformerConfig.from_dict`.
        """
        self.tf_model_config = tf_config
        self._propagate_quantization_from_tf_config(tf_config)
        self._propagate_skip_softmax_calibration(tf_config)

    def _propagate_skip_softmax_calibration(self, tf_config: "TransformerConfig") -> None:
        cfg = getattr(self, "diffusion_attention_config", None)
        if not isinstance(cfg, AttentionConfig):
            return
        specs = [s for s in (cfg.default, *cfg.per_role.values()) if s is not None]

        from vllm_omni.diffusion.attention.backends.trtllm_calibration import (
            propagate_skip_softmax_calibration,
        )

        propagate_skip_softmax_calibration(specs, self.model, tf_config)

    def update_multimodal_support(self) -> None:
        # Resolve serving-visible multimodal behavior from shared metadata
        # instead of importing concrete pipeline modules into the config layer.
        metadata = get_diffusion_model_metadata(self.model_class_name)
        self.supports_multimodal_inputs = metadata.supports_multimodal_inputs
        self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs
        self.supports_mixed_reference_inputs = metadata.supports_mixed_reference_inputs

    def _load_component_transformer_config(self) -> bool:
        """Load a registered pipeline's component-level DiT configuration."""
        from vllm.transformers_utils.config import get_hf_file_to_dict

        from vllm_omni.model_extras import get_transformer_config_subfolder

        transformer_subfolder = get_transformer_config_subfolder(
            self.model_class_name,
            model=self.model,
            revision=self.revision,
        )
        tf_config_dict = get_hf_file_to_dict(
            f"{transformer_subfolder}/config.json",
            self.model,
            revision=self.revision,
        )
        if tf_config_dict is None:
            tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model, revision=self.revision)
        if tf_config_dict is None:
            return False
        self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
        return True

    @staticmethod
    def _looks_like_lance_subfolder(model: str | None) -> bool:
        """Return True when ``--model`` points at a Lance per-component subfolder.

        Lance's HF repo bundles ``Lance_3B/``, ``Lance_3B_Video/`` and
        ``Qwen2.5-VL-ViT/`` under a single top-level ``config.json``; users may
        reasonably hand the AR-style sub-checkpoint path directly.  The
        ``LancePipeline`` constructor knows to walk up to the repo root from
        either subfolder name.
        """
        if not model:
            return False
        base = os.path.basename(str(model).rstrip("/"))
        return base in {"Lance_3B", "Lance_3B_Video"}

    def enrich_config(self) -> None:
        """Load model metadata from HuggingFace and populate config fields.

        Diffusers-style models expose a pipeline index with ``_class_name``.
        Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``,
        so we fall back to reading that and mapping model_type manually.
        """
        from vllm.transformers_utils.config import get_hf_file_to_dict

        from vllm_omni.diffusion.registry import resolve_native_single_file
        from vllm_omni.diffusion.utils.hf_utils import (
            get_diffusion_model_index,
            resolve_native_diffusion_model_class,
        )

        assert self.model is not None

        native_single_file_model = resolve_native_single_file(self.model_class_name)
        if self.is_single_file and native_single_file_model is not None:
            self.diffusion_load_format = "default"
            self.model_class_name = native_single_file_model
            self.diffusers_pipeline_cls = None
            self.set_tf_model_config(TransformerConfig())
            return

        try:
            config_dict = get_diffusion_model_index(
                self.model,
                revision=self.revision,
            )
            if config_dict is not None:
                if self.model_class_name is None:
                    self.model_class_name = config_dict.get("_class_name", None)
                    if self.model_class_name is None and self.diffusion_load_format == "diffusers":
                        self.model_class_name = "DiffusersAdapterPipeline"
                self.update_multimodal_support()

                # Skip transformer config loading for diffusers adapter
                # (non-DiT models don't have a separate transformer folder/config)
                if self.diffusion_load_format == "diffusers":
                    self.set_tf_model_config(TransformerConfig())
                    try:
                        diffusers_pipeline_cls_name = config_dict["_class_name"]
                        self.diffusers_pipeline_cls = getattr(diffusers, diffusers_pipeline_cls_name)
                    except (KeyError, AttributeError) as exc:
                        logger.warning(
                            "Could not find a valid _class_name in the Diffusers pipeline index: %s. "
                            "Without the underlying pipeline class the dummy run may omit required inputs.",
                            exc,
                        )
                else:
                    if not self._load_component_transformer_config():
                        self.set_tf_model_config(TransformerConfig())
            else:
                raise FileNotFoundError("Diffusers pipeline index not found")
        except (AttributeError, OSError, ValueError, FileNotFoundError):
            # Skip transformer config loading for diffusers adapter
            # (non-DiT models don't have a separate transformer folder/config)
            if self.diffusion_load_format == "diffusers":
                if self.model_class_name is None:
                    self.model_class_name = "DiffusersAdapterPipeline"
                self.set_tf_model_config(TransformerConfig())
                logger.warning(
                    "Could not find a valid pipeline index per Diffusers format. "
                    "This model is likely unsupported by the diffusers backend. "
                    "Also, without knowing the underlying pipeline class from its index, "
                    "the dummy run will input only text prompt, which may cause errors for pipelines "
                    "that require additional inputs."
                )
            else:
                cfg = get_hf_file_to_dict("config.json", self.model, revision=self.revision)
                if cfg is None:
                    native_model_class = resolve_native_diffusion_model_class(self.model)
                    if native_model_class is not None:
                        self.model_class_name = native_model_class
                        self.set_tf_model_config(TransformerConfig())
                        self.update_multimodal_support()
                        return
                    # Lance ships its top-level config.json one directory above
                    # the per-checkpoint subfolders (``Lance_3B/`` or
                    # ``Lance_3B_Video/``).  Try to recover that case before
                    # raising.
                    if self._looks_like_lance_subfolder(self.model):
                        self.model_class_name = "LancePipeline"
                        self.set_tf_model_config(TransformerConfig())
                        self.update_multimodal_support()
                        return
                    # An explicit topology or CLI override can select a native,
                    # registered pipeline whose repository contains component
                    # configs but no root HF config or Diffusers index. Trust
                    # that explicit class selection only when its component
                    # transformer config is present.
                    from vllm_omni.diffusion.registry import DiffusionModelRegistry

                    if (
                        self.model_class_name in DiffusionModelRegistry.get_supported_archs()
                        and self._load_component_transformer_config()
                    ):
                        logger.info(
                            "Using explicitly selected diffusion pipeline %r with component-level configuration.",
                            self.model_class_name,
                        )
                        self.update_multimodal_support()
                        return
                    raise ValueError(f"Could not find config.json or a Diffusers pipeline index for {self.model}")

                self.set_tf_model_config(TransformerConfig.from_dict(cfg))
                model_type = cfg.get("model_type")
                architectures = cfg.get("architectures") or []

                from vllm_omni.diffusion.utils.hf_utils import _looks_like_hidream_o1

                assert self.model is not None
                is_hidream_o1 = _looks_like_hidream_o1(self.model, cfg)
                if self.model_class_name == "HiDreamO1ImagePipeline" and not is_hidream_o1:
                    raise ValueError(f"Checkpoint {self.model} does not have the HiDream-O1 signature")

                if is_hidream_o1:
                    self.model_class_name = "HiDreamO1ImagePipeline"
                    self.set_tf_model_config(TransformerConfig())
                    self.update_multimodal_support()
                elif model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
                    self.model_class_name = "BagelPipeline"
                    self.set_tf_model_config(TransformerConfig())
                    self.update_multimodal_support()
                elif (
                    model_type == "lance"
                    or "LancePipeline" in architectures
                    or cfg.get("model_name") == "Lance"
                    or self._looks_like_lance_subfolder(self.model)
                ):
                    # Lance ships a non-HF top-level config.json (model_name only)
                    # plus per-component subfolders; resolve to the Lance pipeline.
                    # Also accept --model pointing directly at the ``Lance_3B`` or
                    # ``Lance_3B_Video`` subfolder by walking up to the repo root.
                    self.model_class_name = "LancePipeline"
                    self.set_tf_model_config(TransformerConfig())
                    self.update_multimodal_support()
                elif model_type == "neo_chat":
                    self.model_class_name = "SenseNovaU1Pipeline"
                    self.tf_model_config = TransformerConfig()
                    self.update_multimodal_support()
                elif "BailingMM2NativeForConditionalGeneration" in architectures or model_type in (
                    "bailingmm_moe_v2_lite",
                    "ming_flash_omni",
                    "ming_flash_omni_thinker",
                ):
                    # Ming-flash-omni-2.0 — imagegen stage uses the custom
                    # ``MingImagePipeline`` (ZImage DiT + Qwen2 connector). See
                    # vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py.
                    self.model_class_name = "MingImagePipeline"
                    self.tf_model_config = TransformerConfig()
                    self.update_multimodal_support()
                elif model_type == "nextstep":
                    if self.model_class_name is None:
                        self.model_class_name = "NextStep11Pipeline"
                    self.set_tf_model_config(TransformerConfig())
                    self.update_multimodal_support()
                elif model_type == "s2v":
                    if self.model_class_name is None:
                        self.model_class_name = "WanS2VPipeline"
                    self.tf_model_config = TransformerConfig()
                    self.update_multimodal_support()
                elif model_type == "Gr00tN1d7" or "Gr00tN1d7" in architectures:
                    self.model_class_name = "Gr00tN1d7Pipeline"
                    self.set_tf_model_config(TransformerConfig())
                    self.update_multimodal_support()
                elif model_type == "vla":
                    from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero

                    assert self.model is not None
                    if _looks_like_dreamzero(self.model):
                        self.model_class_name = "DreamZeroPipeline"
                        self.set_tf_model_config(TransformerConfig())
                        self.update_multimodal_support()
                    else:
                        raise
                elif cfg.get("type") == "pi0":
                    # π0 (Pi-Zero) VLA — the LeRobot config.json uses ``type: "pi0"``.
                    if self.model_class_name is None:
                        self.model_class_name = "Pi0Pipeline"
                    self.set_tf_model_config(TransformerConfig())
                    self.update_multimodal_support()
                elif cfg.get("type") == "pi05":
                    if self.model_class_name is None:
                        self.model_class_name = "Pi05Pipeline"
                    self.set_tf_model_config(TransformerConfig())
                    self.update_multimodal_support()
                elif architectures and len(architectures) == 1:
                    architecture = architectures[0]
                    from vllm_omni.diffusion.registry import DiffusionModelRegistry

                    if (
                        self.model_class_name is None
                        or DiffusionModelRegistry._try_load_model_cls(architecture) is not None
                    ):
                        self.model_class_name = architecture
                    self.update_multimodal_support()
                else:
                    raise

    @classmethod
    def normalize_init_kwargs(cls, raw_kwargs: Mapping[str, Any]) -> dict[str, Any]:
        valid_fields = frozenset(f.name for f in fields(cls))
        config_kwargs = normalize_omni_diffusion_kwargs(raw_kwargs)
        validate_omni_diffusion_kwargs(config_kwargs, valid_fields)
        # Remaining ``None`` values mean "unset" at the CLI/deploy boundary.
        # Drop them so non-optional dataclass defaults are not overwritten.
        # Fields where ``None`` has normalization semantics (for example dtype
        # and nullable container inputs) are handled above before this filter.
        return {key: value for key, value in config_kwargs.items() if key in valid_fields and value is not None}

    @classmethod
    def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
        return cls(**cls.normalize_init_kwargs(kwargs))


DIFFUSION_REQUEST_LIFECYCLE_KEY = "_diffusion_request_lifecycle"
DIFFUSION_REQUEST_STARTED = "started"


def is_diffusion_request_started_output(output: Any) -> bool:
    custom_output = getattr(output, "custom_output", None)
    return (
        isinstance(custom_output, dict)
        and custom_output.get(DIFFUSION_REQUEST_LIFECYCLE_KEY) == DIFFUSION_REQUEST_STARTED
    )


@dataclass
class DiffusionOutput:
    """
    Final output (after pipeline completion)
    """

    output: torch.Tensor | tuple[Any, ...] | dict[str, Any] | None = None

    # Legacy compatibility fields. Decoded video uses ``media``; other new
    # payloads should be carried by output["payload"] instead.
    trajectory_timesteps: torch.Tensor | dict[str, Any] | None = None
    trajectory_latents: torch.Tensor | dict[str, Any] | None = None
    trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None
    trajectory_decoded: list[Image.Image] | None = None
    async_output_id: str | None = None
    error: str | None = None
    error_status_code: int | None = None
    error_type: str | None = None
    aborted: bool = False
    abort_message: str | None = None

    post_process_func: Callable[..., Any] | None = None

    # logged timings info, directly from Req.timings
    # timings: Optional["RequestTimings"] = None

    # Streaming info (the defaults should make sense for non-streaming mode)
    finished: bool = True
    chunk_index: int = 0
    total_chunks: int = 1
    started_event_ids: list[str] = field(default_factory=list)
    active_event_ids: list[str] = field(default_factory=list)
    completed_event_ids: list[str] = field(default_factory=list)

    # logged duration of stages
    stage_durations: dict[str, float] = field(default_factory=dict)

    # memory usage info
    peak_memory_mb: float = 0.0

    # KV-recv wall-clock (ms) from the runner; feeds vllm_omni:diffusion_kv_load_s.
    kv_recv_ms: float = 0.0

    # When True, move tensor fields to CPU at construction time. Useful when
    # the output is shipped across process boundaries (e.g. step-execution
    # mode) and the receiving side must not initialise a stray CUDA context.
    to_cpu: bool = False

    # Internal control-plane event emitted on first scheduler admission.
    request_started: bool = False

    # Typed video-media contract. Appended so the pre-existing positional
    # constructor order (output, trajectory_timesteps, ...) that out-of-tree
    # pipelines rely on is preserved. Mutually exclusive with ``output``.
    media: DiffusionMediaOutput | None = None

    # Compatibility adapter for joint (video, audio) outputs that have not
    # migrated to typed media. Only this tuple entry is eligible for video
    # transport optimizations; unmarked legacy outputs retain their old path.
    video_output_index: int | None = None

    def __post_init__(self) -> None:
        if self.media is not None and not isinstance(self.media, DiffusionMediaOutput):
            raise TypeError(f"media must be DiffusionMediaOutput, got {type(self.media).__name__}")
        if self.media is not None and self.output is not None:
            raise ValueError("DiffusionOutput cannot contain both media and legacy output")
        if self.media is not None and self.post_process_func is not None:
            raise ValueError("Typed diffusion media cannot carry a model-specific post_process_func")
        if not self.to_cpu:
            return

        def _maybe_to_cpu(value: Any) -> Any:
            if isinstance(value, torch.Tensor):
                return value.detach().cpu()
            if isinstance(value, dict):
                return {key: _maybe_to_cpu(item) for key, item in value.items()}
            if isinstance(value, list):
                return [_maybe_to_cpu(item) for item in value]
            if isinstance(value, tuple):
                return tuple(_maybe_to_cpu(item) for item in value)
            return value

        self.output = _maybe_to_cpu(self.output)
        if self.media is not None:
            if not self.media.prepared_for_transport:
                raise ValueError("Diffusion media must be prepared before to_cpu=True")
            self.media = self.media.to_cpu()
        self.trajectory_timesteps = _maybe_to_cpu(self.trajectory_timesteps)
        self.trajectory_latents = _maybe_to_cpu(self.trajectory_latents)
        self.trajectory_log_probs = _maybe_to_cpu(self.trajectory_log_probs)

    @classmethod
    def from_exception(cls, exc: BaseException) -> "DiffusionOutput":
        status_code, error_type = client_error_metadata(exc)
        return cls(
            error=str(exc),
            error_status_code=status_code,
            error_type=error_type,
        )


class AsyncOutputKind(Enum):
    """Message kind for ``AsyncDiffusionOutput`` — routes async results.

    * ``RPC_RESULT`` — ordinary RPC return (sleep, wake, profile, etc.)
    * ``COMPUTE_DONE`` — worker forward finished, GPU can start next request
    * ``OUTPUT_READY`` — background D2H/SHM packing finished, final output
      is available via ``async_output_id``
    """

    RPC_RESULT = "rpc_result"
    COMPUTE_DONE = "compute_done"
    OUTPUT_READY = "output_ready"


@dataclass
class AsyncDiffusionOutput:
    """Async protocol envelope for ``result_mq`` messages.

    In request-mode (``step_execution=False``), all ``result_mq``
    messages use this envelope.  The ``kind`` field routes the message
    to the correct consumer.
    """

    kind: AsyncOutputKind
    rpc_id: str | None = None
    async_output_id: str | None = None
    result: Any | None = None
    output: DiffusionOutput | None = None
    error: str | None = None


class DiffusionRequestAbortedError(RuntimeError):
    """Raised when a diffusion request ends via user-visible abort."""


def _in_range(value: Any, name: str, lo: float, hi: float | None) -> float | None:
    """Validate an optional numeric control at config-parse time."""
    if value is None:
        return None
    x = float(value)
    if not math.isfinite(x) or x < lo or (hi is not None and x > hi):
        rng = f"in [{lo}, {hi}]" if hi is not None else f">= {lo}"
        raise ValueError(f"{name} must be finite and {rng}; got {value!r}.")
    return x


@dataclass
class SkipSoftmaxSpec:
    """User-facing skip-softmax controls for the TRTLLM_ATTN backend.

    ``target_sparsity`` uses the checkpoint's calibrated curve (a·exp(b·s)); ``threshold``
    is the calibration-free absolute skip threshold. They are mutually exclusive.
    """

    target_sparsity: float | None = None
    threshold: float | None = None
    disabled_until_timestep: float = 0.0

    def __post_init__(self) -> None:
        self.target_sparsity = _in_range(self.target_sparsity, "skip_softmax.target_sparsity", 0.0, 1.0)
        self.threshold = _in_range(self.threshold, "skip_softmax.threshold", 0.0, None)
        self.disabled_until_timestep = (
            _in_range(self.disabled_until_timestep, "skip_softmax.disabled_until_timestep", 0.0, 1.0) or 0.0
        )
        if self.target_sparsity is not None and self.threshold is not None:
            raise ValueError("skip_softmax: set either target_sparsity or threshold, not both.")

    @property
    def enabled(self) -> bool:
        return self.threshold is not None or self.target_sparsity is not None


@dataclass
class AttnQuantSpec:
    dtype_qk: str | None = None
    dtype_vo: str | None = None
    q_block_size: int = 1
    k_block_size: int = 16
    flashinfer_backend: str | None = None

    _VALID_DTYPES = frozenset({"float16", "bfloat16", "int8", "fp8_e4m3"})
    _VALID_BLOCK_SIZES = frozenset({1, 4, 16})

    def __post_init__(self) -> None:
        for name, v in (("dtype_qk", self.dtype_qk), ("dtype_vo", self.dtype_vo)):
            if v is not None and v not in self._VALID_DTYPES:
                raise ValueError(f"quant.{name}={v!r} unsupported; use one of {sorted(self._VALID_DTYPES)}.")
        for block_name, block_value in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)):
            if block_value not in self._VALID_BLOCK_SIZES:
                raise ValueError(
                    f"quant.{block_name}={block_value!r} unsupported; "
                    f"kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}."
                )

    @property
    def enabled(self) -> bool:
        # Include flashinfer_backend so a variant pin without dtypes is serialized.
        return self.dtype_qk is not None or self.dtype_vo is not None or self.flashinfer_backend is not None


# Backends that select key blocks instead of attending densely, and so accept
# a ``block_sparse`` spec. Each maps the same knobs onto its own kernel.
BLOCK_SPARSE_BACKENDS = frozenset({"RAINFUSION_ATTN"})


class RainFusionPrecision(str, Enum):
    """Execution precision for block-sparse RainFusion (rf_v3) attention.

    ``bf16``: no quantization, pure BF16 sparse attention (official baseline).
    ``fp8``: BSA FP8 path - Hadamard rotation then full FP8 block quantization
        of Q/K/V before the BSA kernel.
    ``mix``: EagleQBSA mixed precision - Q/K per-block INT8 + V per-channel FP8.
    """

    BF16 = "bf16"
    FP8 = "fp8"
    MIX = "mix"


@dataclass
class BlockSparseSpec:
    """User-facing controls shared by block-sparse attention backends.

    ``sparsity`` is the nominal fraction of key blocks dropped per query block;
    ``start_step`` keeps the first N denoise steps dense and ``skip_layers`` (an
    index selector such as "0-3,38") exempts individual DiT blocks. Those two are
    the accuracy knobs to trade back quality at a fixed ``sparsity``.
    ``precision`` selects the kernel precision mode (see ``RainFusionPrecision``).
    """

    sparsity: float = 0.8
    start_step: int = 0
    end_step: int = 0
    precision: str = RainFusionPrecision.BF16.value
    skip_layers: str | list[int] | None = None
    skip_layer_indices: set[int] | None = field(default=None, repr=False)

    def __post_init__(self) -> None:
        self.sparsity = _in_range(self.sparsity, "block_sparse.sparsity", 0.0, 1.0) or 0.0
        if self.start_step < 0:
            raise ValueError(f"block_sparse.start_step must be >= 0; got {self.start_step!r}.")
        if self.end_step < 0:
            raise ValueError(f"block_sparse.end_step must be >= 0; got {self.end_step!r}.")
        precision = self.precision.value if isinstance(self.precision, RainFusionPrecision) else self.precision
        valid = {member.value for member in RainFusionPrecision}
        if precision not in valid:
            raise ValueError(f"block_sparse.precision must be one of {sorted(valid)}; got {self.precision!r}.")
        self.precision = precision
        self.skip_layer_indices = parse_kv_cache_skip_selector(self.skip_layers)


@dataclass
class AttentionSpec:
    """Specifies a backend and its typed backend-specific config for one attention role."""

    backend: str
    skip_softmax: SkipSoftmaxSpec | None = None
    quant: AttnQuantSpec | None = None
    fastvideo_vsa_topk: int | None = None
    block_sparse: BlockSparseSpec | None = None
    skip_calibration: dict | None = field(default=None, repr=False)

    def __post_init__(self) -> None:
        if not isinstance(self.backend, str):
            raise TypeError(f"Expected str for AttentionSpec.backend, got {type(self.backend)!r}")
        self.skip_softmax = self._coerce(self.skip_softmax, SkipSoftmaxSpec, "skip_softmax")
        self.quant = self._coerce(self.quant, AttnQuantSpec, "quant")
        self.block_sparse = self._coerce(self.block_sparse, BlockSparseSpec, "block_sparse")
        if self.skip_softmax is not None and self.backend.upper() != "TRTLLM_ATTN":
            raise ValueError(
                f"skip_softmax is only supported by the TRTLLM_ATTN backend, but backend={self.backend!r}. "
                "Remove skip_softmax or set backend to TRTLLM_ATTN."
            )
        if self.quant is not None and self.backend.upper() not in ("TRTLLM_ATTN", "FLASHINFER_ATTN"):
            raise ValueError(
                f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but "
                f"backend={self.backend!r}. Remove quant or set a supported backend."
            )
        if self.fastvideo_vsa_topk is not None:
            if self.backend.upper() != "FASTVIDEO_VSA":
                raise ValueError("fastvideo_vsa_topk is only supported by the FASTVIDEO_VSA backend.")
            if self.fastvideo_vsa_topk <= 0:
                raise ValueError("fastvideo_vsa_topk must be positive.")
        if self.backend.upper() in BLOCK_SPARSE_BACKENDS:
            # Selecting the backend is the opt-in; without an explicit block the
            # defaults apply rather than silently running dense.
            self.block_sparse = self.block_sparse or BlockSparseSpec()
        elif self.block_sparse is not None:
            raise ValueError(
                f"block_sparse is only supported by the {sorted(BLOCK_SPARSE_BACKENDS)} backends, but "
                f"backend={self.backend!r}. Remove block_sparse or set a supported backend."
            )

    @staticmethod
    def _coerce(value: Any, cls: type, field_name: str) -> Any:
        if value is None or isinstance(value, cls):
            return value
        if isinstance(value, Mapping):
            return cls(**dict(value))
        raise TypeError(f"Expected dict or {cls.__name__} for {field_name}, got {type(value)!r}")

    def backend_kwargs(self) -> dict[str, Any] | None:
        """Serialize typed backend config into the kwargs dict the backend impl consumes."""
        kw: dict[str, Any] = {}
        if self.skip_softmax is not None:
            ss = self.skip_softmax
            if ss.threshold is not None:
                kw["skip_softmax_threshold"] = ss.threshold
            if ss.target_sparsity is not None:
                kw["target_sparsity"] = ss.target_sparsity
            if ss.disabled_until_timestep:
                kw["disabled_until_timestep"] = ss.disabled_until_timestep
        if self.quant is not None:
            q = self.quant
            quant_kw: dict[str, Any] = {}
            if q.dtype_qk is not None or q.dtype_vo is not None:
                quant_kw["dtype_qk"] = q.dtype_qk
                quant_kw["q_block_size"] = q.q_block_size
                quant_kw["k_block_size"] = q.k_block_size
                if q.dtype_vo is not None:
                    quant_kw["dtype_vo"] = q.dtype_vo
            if q.flashinfer_backend is not None:
                quant_kw["flashinfer_backend"] = q.flashinfer_backend
            if quant_kw:
                kw["quant"] = quant_kw
        if self.fastvideo_vsa_topk is not None:
            kw["topk"] = self.fastvideo_vsa_topk
        if self.block_sparse is not None:
            bs = self.block_sparse
            kw["sparsity"] = bs.sparsity
            kw["start_step"] = bs.start_step
            kw["end_step"] = bs.end_step
            kw["precision"] = bs.precision
            if bs.skip_layer_indices:
                kw["skip_layers"] = sorted(bs.skip_layer_indices)

        return kw or None


@dataclass
class AttentionConfig:
    """Per-role attention backend configuration.

    Lookup precedence for a given (role, role_category):
      1. per_role[role]         — exact match
      2. per_role[role_category] — category fallback (e.g. "ltx2.audio_to_video" → "cross")
      3. default                — global default
      4. platform default       — unchanged platform logic
    """

    default: AttentionSpec | None = None
    per_role: dict[str, AttentionSpec] = field(default_factory=dict)

    def __post_init__(self) -> None:
        if self.default is not None:
            self.default = self._coerce_spec_or_none(self.default, "default")

        normalized_per_role: dict[str, AttentionSpec] = {}
        for role_key, spec_data in self._normalize_per_role_mapping(self.per_role).items():
            spec = self._coerce_spec_or_none(spec_data, f"per_role[{role_key!r}]")
            if spec is not None:
                normalized_per_role[role_key] = spec
        self.per_role = normalized_per_role

    @staticmethod
    def _coerce_spec(spec_data: Any, field_name: str) -> AttentionSpec:
        if isinstance(spec_data, AttentionSpec):
            return spec_data
        if isinstance(spec_data, str):
            return AttentionSpec(backend=spec_data)
        if isinstance(spec_data, Mapping):
            return AttentionSpec(**dict(spec_data))
        raise TypeError(f"Expected str, dict, or AttentionSpec for {field_name}, got {type(spec_data)!r}")

    @classmethod
    def _coerce_spec_or_none(cls, spec_data: Any, field_name: str) -> AttentionSpec | None:
        spec = cls._coerce_spec(spec_data, field_name)
        if spec.backend.lower() == "auto":
            return None
        return spec

    @classmethod
    def _normalize_per_role_mapping(cls, raw_per_role: Any) -> dict[str, Any]:
        if raw_per_role is None:
            return {}
        if not isinstance(raw_per_role, Mapping):
            raise TypeError(f"Expected dict for AttentionConfig.per_role, got {type(raw_per_role)!r}")

        normalized: dict[str, Any] = {}
        for role_key, spec_data in raw_per_role.items():
            cls._flatten_per_role_entry([role_key], spec_data, normalized)
        return normalized

    @classmethod
    def _flatten_per_role_entry(
        cls,
        path: list[str],
        node: Any,
        normalized: dict[str, Any],
    ) -> None:
        role = ".".join(path)
        if not isinstance(node, Mapping):
            normalized[role] = node
            return

        spec_keys = {"backend", "skip_softmax", "quant", "fastvideo_vsa_topk", "block_sparse"}
        node_dict = dict(node)
        node_keys = set(node_dict)
        if node_keys & spec_keys:
            if not node_keys <= spec_keys:
                raise ValueError(
                    f"Invalid per_role entry for role {role!r}: cannot mix backend/skip_softmax with nested role keys."
                )
            normalized[role] = node_dict
            return

        if not node_dict:
            raise ValueError(f"Empty per_role entry for role {role!r}")

        for child_key, child_value in node_dict.items():
            cls._flatten_per_role_entry([*path, child_key], child_value, normalized)

    def resolve_with_source(
        self,
        role: str = "self",
        role_category: str | None = None,
    ) -> tuple[AttentionSpec | None, str | None]:
        """Resolve the AttentionSpec and report which config entry matched."""
        spec = self.per_role.get(role)
        if spec is not None:
            return spec, f"attention_config.per_role[{role!r}]"
        if role_category is not None:
            spec = self.per_role.get(role_category)
            if spec is not None:
                return spec, f"attention_config.per_role[{role_category!r}] (role_category fallback)"
        if self.default is not None:
            return self.default, "attention_config.default"
        return None, None


def parse_attention_config(
    attention_config: AttentionConfig | Mapping[str, Any] | None = None,
    *,
    attention_backend: str | None = None,
    fastvideo_vsa_topk: int | None = None,
) -> AttentionConfig:
    """Pure type-conversion: coerce *attention_config* to an AttentionConfig.

    Optionally merges an ``attention_backend`` shorthand into the config's
    ``default`` field.  This does **not** read environment variables —
    use :func:`build_attention_config` for the full normalisation that
    should happen exactly once in ``OmniDiffusionConfig.__post_init__``.
    """
    if attention_config is None:
        normalized = AttentionConfig()
    elif isinstance(attention_config, AttentionConfig):
        normalized = copy.deepcopy(attention_config)
    elif isinstance(attention_config, Mapping):
        normalized = AttentionConfig(**dict(attention_config))
    else:
        raise TypeError(
            f"attention_config must be an AttentionConfig, mapping, or None; got {type(attention_config)!r}"
        )

    if attention_backend is not None:
        if normalized.default is not None:
            raise ValueError(
                "--diffusion-attention-backend is mutually exclusive with --diffusion-attention-config.default.backend."
            )
        if attention_backend.lower() != "auto":
            normalized.default = AttentionSpec(backend=attention_backend)

    if fastvideo_vsa_topk is not None:
        if normalized.default is None:
            raise ValueError("--fastvideo-vsa-topk requires --diffusion-attention-backend FASTVIDEO_VSA.")
        if normalized.default.backend.upper() != "FASTVIDEO_VSA":
            raise ValueError("--fastvideo-vsa-topk is only valid with the FASTVIDEO_VSA backend.")
        normalized.default.fastvideo_vsa_topk = fastvideo_vsa_topk
        normalized.default.__post_init__()

    return normalized


def build_attention_config(
    attention_config: AttentionConfig | Mapping[str, Any] | None = None,
) -> AttentionConfig:
    """Normalize diffusion attention config — the single authoritative entry point.

    Called exactly once in ``OmniDiffusionConfig.__post_init__``.
    Handles type-conversion **and** env-var fallback
    (``DIFFUSION_ATTENTION_BACKEND`` and ``DIFFUSION_ATTENTION_QUANT``).
    """
    normalized = parse_attention_config(attention_config)

    if normalized.default is not None:
        return normalized

    env_attention_backend = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
    if env_attention_backend is None:
        return normalized

    if env_attention_backend.lower() == "auto":
        return normalized

    # If config file doesn't explicitly setup quantization specs, optionally set
    # from the env-var.
    quant = None
    env_attention_quant = os.environ.get("DIFFUSION_ATTENTION_QUANT")
    if env_attention_quant is not None:
        fields = [field.strip() for field in env_attention_quant.split(":")]
        if len(fields) not in (2, 4) or not all(fields):
            raise ValueError(
                "DIFFUSION_ATTENTION_QUANT must have the format <dtype_qk>:<dtype_vo>[:<q_block_size>:<k_block_size>]."
            )
        quant_kwargs: dict[str, Any] = {
            "dtype_qk": fields[0],
            "dtype_vo": fields[1],
        }
        if len(fields) == 4:
            q_block_size, k_block_size = int(fields[2]), int(fields[3])
            if (q_block_size, k_block_size) != (0, 0):
                quant_kwargs.update(
                    q_block_size=q_block_size,
                    k_block_size=k_block_size,
                )
        quant = AttnQuantSpec(**quant_kwargs)

    normalized.default = AttentionSpec(backend=env_attention_backend, quant=quant)
    logger.info(
        "Parsed attention config from DIFFUSION_ATTENTION_BACKEND '%s': default=%s, per_role=%s",
        env_attention_backend,
        normalized.default,
        {k: v.backend for k, v in normalized.per_role.items()},
    )
    return normalized


@dataclass
class OmniACK:
    """
    Handshake payload from Workers to Orchestrator.
    """

    task_id: str
    status: str
    stage_id: int | None = None
    rank: int | None = None
    freed_bytes: int = 0
    metadata: dict[str, Any] = field(default_factory=dict)
    """
    Additional telemetry such as:
    - max_contiguous_block: for fragmentation analysis.
    - cuda_graph_recalled: boolean if graphs were successfully destroyed/rebuilt.
    - latency_ms: time taken for the D2H/H2D transfer.
    """
    error_msg: str | None = None


@dataclass
class OmniSleepTask:
    """Structured sleep instruction."""

    task_id: str
    level: int = 2
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class OmniWakeTask:
    """Structured wake-up instruction."""

    task_id: str
    tags: list[str] | None = None


class CuMemTag(str, Enum):
    """Tags representing specific CuMem allocations for sleep/wake state tracking."""

    WEIGHTS = "weights"
    KV_CACHE = "kv_cache"


# Special message broadcast via scheduler queues to signal worker shutdown.
SHUTDOWN_MESSAGE = {"type": "shutdown"}