2387 lines
106 KiB
Python
2387 lines
106 KiB
Python
# 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"}
|