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 project
import copy
import math
import os
import random
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.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 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__)


def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]:
    """Normalize legacy diffusion kwargs before config construction."""
    normalized = dict(kwargs)

    # Backwards-compatibility: older callers may use a diffusion-specific
    # "static_lora_scale" kwarg. Normalize it to the canonical "lora_scale".
    if "static_lora_scale" in normalized:
        if "lora_scale" not in normalized:
            normalized["lora_scale"] = normalized["static_lora_scale"]
        normalized.pop("static_lora_scale", None)

    # Backwards-compatibility: map "quantization" to "quantization_config"
    # so callers using the old field name still work.
    if "quantization" in normalized and normalized.get("quantization_config", None) is None:
        normalized["quantization_config"] = normalized.pop("quantization")
    else:
        normalized.pop("quantization", None)

    # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype.
    if normalized.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in normalized:
        normalized["diffusion_kv_cache_dtype"] = normalized.pop("kv_cache_dtype")
    else:
        normalized.pop("kv_cache_dtype", None)
    if normalized.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in normalized:
        normalized["diffusion_kv_cache_skip_steps"] = normalized.pop("kv_cache_skip_steps")
    else:
        normalized.pop("kv_cache_skip_steps", None)
    if normalized.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in normalized:
        normalized["diffusion_kv_cache_skip_layers"] = normalized.pop("kv_cache_skip_layers")
    else:
        normalized.pop("kv_cache_skip_layers", None)

    # Handle "diffusion_attention_backend" shorthand: merge into
    # diffusion_attention_config before field filtering.
    diffusion_attn_backend = normalized.pop("diffusion_attention_backend", None)
    if diffusion_attn_backend is not None:
        existing = normalized.get("diffusion_attention_config")
        normalized["diffusion_attention_config"] = parse_attention_config(
            existing,
            attention_backend=diffusion_attn_backend,
        )

    # Check environment variable as fallback for cache_backend.
    # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND.
    if "cache_backend" not in normalized:
        cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
        normalized["cache_backend"] = cache_backend.lower() if cache_backend else "none"

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

    return normalized


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
    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[int] = 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.
    """

    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.
    - "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 > 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", "spatial_shard_height", "spatial_shard_width"}, (
            "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
            f"but got {self.vae_parallel_mode!r}."
        )
        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

        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.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)


@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
        - 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

    # 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

    if not model:
        return None
    is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"}

    # Diffusers models: read _class_name from the pipeline index.
    model_index = get_diffusion_model_index(model, revision=revision)
    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


@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"  # "tea_cache", "deep_cache", etc.
    cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict)
    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
    distributed_executor_backend: str = "mp"
    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

    # 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
    # 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

    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

    # 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
    num_weight_load_threads: int = 4

    # 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

    # 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,
        }
    )
    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)
    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

    # 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
    # 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

    # Diffusion attention KV cache dtype (not vLLM's --kv-cache-dtype for AR models).
    # None = native dtype (no quantization).
    # "fp8" = dynamic FP8 (float8_e4m3fn) quantization per forward pass.
    # On Hopper+FA3: native FP8 attention (memory + compute savings).
    # On other backends: no benefit, backends skip quantization.
    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.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 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 in Phase 1; "
                "disable need_recv_cache until connector-aware admission is implemented"
            )

        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)

        if self.diffusion_compile_granularity == "full":
            incompatible_features = []
            if self.parallel_config.use_hsdp:
                incompatible_features.append("HSDP")
            if self.parallel_config.sequence_parallel_size > 1:
                incompatible_features.append("sequence parallelism")
            if self.enable_cpu_offload:
                incompatible_features.append("CPU offload")
            if self.enable_layerwise_offload:
                incompatible_features.append("layerwise offload")
            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()

        # 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")

        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,
                )

    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:
            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

    @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.utils.hf_utils import get_diffusion_model_index

        # Default model_class_name for diffusers adapter
        if self.model_class_name is None and self.diffusion_load_format == "diffusers":
            self.model_class_name = "DiffusersAdapterPipeline"

        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)
                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:
                    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 not None:
                        self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
                    else:
                        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":
                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:
                    # 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
                    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

                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

                    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 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
                else:
                    raise

    @classmethod
    def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
        kwargs = normalize_omni_diffusion_kwargs(kwargs)

        # Filter kwargs to only include valid fields
        valid_fields = {f.name for f in fields(cls)}
        filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}

        instance = cls(**filtered_kwargs)
        return instance


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

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

    # Legacy compatibility fields. New pipeline-specific 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

    # 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

    def __post_init__(self) -> None:
        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)
        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 name, v in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)):
            if v not in self._VALID_BLOCK_SIZES:
                raise ValueError(
                    f"quant.{name}={v!r} unsupported; kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}."
                )

    @property
    def enabled(self) -> bool:
        return self.dtype_qk is not None or self.dtype_vo 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"})


@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``.
    """

    sparsity: float = 0.8
    start_step: int = 0
    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}.")
        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
    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.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 and self.quant.enabled:
            q = self.quant
            quant_kw: dict[str, Any] = {
                "dtype_qk": q.dtype_qk,
                "q_block_size": q.q_block_size,
                "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
            kw["quant"] = quant_kw
        if self.block_sparse is not None:
            bs = self.block_sparse
            kw["sparsity"] = bs.sparsity
            kw["start_step"] = bs.start_step
            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", "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,
) -> 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)

    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``).
    """
    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

    normalized.default = AttentionSpec(backend=env_attention_backend)
    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"}