Hongsheng Liu
f1abb13958
Avoid recurrent mmap-to-pinned packing in no-AllGather DLO by registering transform-complete runtime-cache mappings under an explicit per-worker budget. Copy mapped tensors directly into the rotating device buffers while preserving bounded-staging fallback and ordered shutdown cleanup. Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
1856 lines
80 KiB
Python
1856 lines
80 KiB
Python
# adapted from sglang and fastvideo
|
||
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
import copy
|
||
import math
|
||
import os
|
||
import random
|
||
from collections.abc import Callable, Mapping
|
||
from dataclasses import dataclass, field, fields
|
||
from enum import Enum
|
||
from typing import TYPE_CHECKING, Any
|
||
|
||
import diffusers
|
||
import huggingface_hub
|
||
import torch
|
||
from PIL import Image
|
||
from pydantic import Field, model_validator
|
||
from typing_extensions import Self
|
||
from vllm.config.utils import config
|
||
from vllm.logger import init_logger
|
||
from vllm.model_executor.layers.quantization.base_config import (
|
||
QuantizationConfig,
|
||
)
|
||
from vllm.transformers_utils.repo_utils import get_model_path
|
||
|
||
from vllm_omni.diffusion.diffusion_kv.config import (
|
||
DiffusionKVCacheMode,
|
||
parse_diffusion_kv_cache_mode,
|
||
)
|
||
from vllm_omni.diffusion.lora.manager import LoRABackend
|
||
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
|
||
from vllm_omni.diffusion.utils.network_utils import is_port_available
|
||
from vllm_omni.errors import client_error_metadata
|
||
from vllm_omni.quantization import build_quant_config
|
||
|
||
if TYPE_CHECKING:
|
||
from vllm.config import ProfilerConfig
|
||
|
||
# Import after TYPE_CHECKING to avoid circular imports at runtime
|
||
# The actual import is deferred to __post_init__ to avoid import order issues
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]:
|
||
"""Normalize legacy diffusion kwargs before config construction."""
|
||
normalized = dict(kwargs)
|
||
|
||
# Backwards-compatibility: older callers may use a diffusion-specific
|
||
# "static_lora_scale" kwarg. Normalize it to the canonical "lora_scale".
|
||
if "static_lora_scale" in normalized:
|
||
if "lora_scale" not in normalized:
|
||
normalized["lora_scale"] = normalized["static_lora_scale"]
|
||
normalized.pop("static_lora_scale", None)
|
||
|
||
# Backwards-compatibility: map "quantization" to "quantization_config"
|
||
# so callers using the old field name still work.
|
||
if "quantization" in normalized and normalized.get("quantization_config", None) is None:
|
||
normalized["quantization_config"] = normalized.pop("quantization")
|
||
else:
|
||
normalized.pop("quantization", None)
|
||
|
||
# Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype.
|
||
if normalized.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in normalized:
|
||
normalized["diffusion_kv_cache_dtype"] = normalized.pop("kv_cache_dtype")
|
||
else:
|
||
normalized.pop("kv_cache_dtype", None)
|
||
if normalized.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in normalized:
|
||
normalized["diffusion_kv_cache_skip_steps"] = normalized.pop("kv_cache_skip_steps")
|
||
else:
|
||
normalized.pop("kv_cache_skip_steps", None)
|
||
if normalized.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in normalized:
|
||
normalized["diffusion_kv_cache_skip_layers"] = normalized.pop("kv_cache_skip_layers")
|
||
else:
|
||
normalized.pop("kv_cache_skip_layers", None)
|
||
|
||
# Handle "diffusion_attention_backend" shorthand: merge into
|
||
# diffusion_attention_config before field filtering.
|
||
diffusion_attn_backend = normalized.pop("diffusion_attention_backend", None)
|
||
if diffusion_attn_backend is not None:
|
||
existing = normalized.get("diffusion_attention_config")
|
||
normalized["diffusion_attention_config"] = parse_attention_config(
|
||
existing,
|
||
attention_backend=diffusion_attn_backend,
|
||
)
|
||
|
||
# Check environment variable as fallback for cache_backend.
|
||
# Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND.
|
||
if "cache_backend" not in normalized:
|
||
cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER")
|
||
normalized["cache_backend"] = cache_backend.lower() if cache_backend else "none"
|
||
|
||
# Convert optional YAML null values to empty containers.
|
||
for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"):
|
||
if key in normalized and normalized[key] is None:
|
||
normalized[key] = {}
|
||
|
||
return normalized
|
||
|
||
|
||
def parse_kv_cache_skip_selector(
|
||
selector: str | list[int] | tuple[int, ...] | set[int] | None,
|
||
) -> set[int] | None:
|
||
"""Parse a non-negative index selector such as "0-9,20,25-30"."""
|
||
if selector is None:
|
||
return None
|
||
if isinstance(selector, set):
|
||
values = selector
|
||
elif isinstance(selector, (list, tuple)):
|
||
values = set(selector)
|
||
elif isinstance(selector, str):
|
||
text = selector.strip()
|
||
if not text:
|
||
return None
|
||
values: set[int] = set()
|
||
for chunk in text.split(","):
|
||
token = chunk.strip()
|
||
if not token:
|
||
continue
|
||
if "-" in token:
|
||
start_str, end_str = token.split("-", 1)
|
||
try:
|
||
start = int(start_str.strip())
|
||
end = int(end_str.strip())
|
||
except ValueError as exc:
|
||
raise ValueError(f"Invalid range token '{token}' in selector '{selector}'.") from exc
|
||
if start < 0 or end < 0 or start > end:
|
||
raise ValueError(f"Invalid range token '{token}' in selector '{selector}'.")
|
||
values.update(range(start, end + 1))
|
||
else:
|
||
try:
|
||
index = int(token)
|
||
except ValueError as exc:
|
||
raise ValueError(f"Invalid index token '{token}' in selector '{selector}'.") from exc
|
||
if index < 0:
|
||
raise ValueError(f"Negative index '{index}' is not allowed in selector '{selector}'.")
|
||
values.add(index)
|
||
else:
|
||
raise TypeError(f"Unsupported selector type: {type(selector)!r}")
|
||
|
||
for idx in values:
|
||
if not isinstance(idx, int):
|
||
raise TypeError(f"Selector index must be int, got {type(idx)!r}")
|
||
if idx < 0:
|
||
raise ValueError("Selector indices must be non-negative.")
|
||
return values
|
||
|
||
|
||
@config
|
||
@dataclass
|
||
class DiffusionParallelConfig:
|
||
"""Configuration for diffusion model distributed execution."""
|
||
|
||
pipeline_parallel_size: int = 1
|
||
"""Number of pipeline parallel stages."""
|
||
|
||
data_parallel_size: int = 1
|
||
"""Number of data parallel groups."""
|
||
|
||
tensor_parallel_size: int = 1
|
||
"""Number of tensor parallel groups."""
|
||
|
||
enable_expert_parallel: bool = False
|
||
"""Enable expert parallelism for MoE layers (TP is still used for non-MoE layers)."""
|
||
|
||
sequence_parallel_size: int | None = None
|
||
"""Number of sequence parallel groups.
|
||
sequence_parallel_size = ulysses_degree * ring_degree, or allgather_degree
|
||
when AllGather-KV is enabled."""
|
||
|
||
ulysses_degree: int = 1
|
||
"""Number of GPUs used for ulysses sequence parallelism."""
|
||
|
||
ring_degree: int = 1
|
||
"""Number of GPUs used for ring sequence parallelism."""
|
||
|
||
allgather_degree: int = 1
|
||
"""Number of GPUs used for AllGather-KV sequence parallelism (causal=False only)."""
|
||
|
||
ulysses_mode: str = "strict"
|
||
"""Ulysses sequence-parallel mode.
|
||
|
||
- "strict": Require divisibility constraints (fastest, default).
|
||
- "advanced_uaa": Enable UAA ("Ulysses Anything Attention") to support
|
||
uneven sequence lengths and non-divisible head counts.
|
||
|
||
Note:
|
||
- Ring attention does not support `attention_mask`, so models that rely on
|
||
mask-based auto-padding are still incompatible with Ring.
|
||
- When used in hybrid Ulysses+Ring, Ring requires consistent per-rank
|
||
sequence shapes across the ring group.
|
||
"""
|
||
|
||
cfg_parallel_size: int = 1
|
||
"""Number of ranks used to execute guidance passes in parallel."""
|
||
|
||
vae_patch_parallel_size: int = 1
|
||
"""Number of ranks used for VAE patch/tile parallelism (decode/encode)."""
|
||
|
||
text_encoder_tp_size: int = 1
|
||
"""Number of ranks used to tensor-parallel shard the diffusion text encoder.
|
||
|
||
Ranks are the first ``text_encoder_tp_size`` DiT ranks. Defaults to 1,
|
||
which keeps the encoder fully resident on the DiT main rank (historical
|
||
behavior). Values > 1 shard the Qwen3-VL encoder across the encoder TP
|
||
ranks and run the encode with distributed collectives over the encoder
|
||
process group.
|
||
"""
|
||
|
||
vae_parallel_mode: str = "tile"
|
||
"""VAE parallel decode strategy.
|
||
|
||
- "tile": Patch/tile parallel decode (default). Each rank decodes a subset
|
||
of spatial tiles and the results are stitched on rank 0.
|
||
- "spatial_shard_height": Spatially-sharded decode that splits decoder
|
||
feature maps along height and exchanges halo rows around spatial
|
||
convolutions.
|
||
- "spatial_shard_width": Same as "spatial_shard_height" but sharded along width.
|
||
|
||
The "spatial_shard_*" modes are decode-only and currently require
|
||
``vae_patch_parallel_size`` to match the DiT group size; otherwise the VAE
|
||
falls back to tile-parallel decode at runtime.
|
||
"""
|
||
|
||
use_hsdp: bool = False
|
||
"""Enable Hybrid Sharded Data Parallel (HSDP) for model weight sharding."""
|
||
|
||
mask_sp_padding: bool = False
|
||
"""If True, generate a boolean attention mask for zero-padded SP tokens
|
||
when sequence length is not divisible by the SP world size. The mask
|
||
routes attention through the varlen path (unpad→kernel→repad), which is
|
||
correct but carries additional overhead. When False (default), padding
|
||
tokens are left unmasked; since _shard_with_auto_pad always pads with
|
||
zeros, their contribution to attention output is negligible."""
|
||
|
||
hsdp_shard_size: int = -1
|
||
"""Number of GPUs to shard weights across within each replica group. -1 means auto-calculate."""
|
||
|
||
hsdp_replicate_size: int = 1
|
||
"""Number of replica groups for HSDP. Each replica holds a full sharded copy."""
|
||
|
||
@model_validator(mode="after")
|
||
def _validate_parallel_config(self) -> Self:
|
||
"""Validates the config relationships among the parallel strategies."""
|
||
assert self.pipeline_parallel_size > 0, "Pipeline parallel size must be > 0"
|
||
assert self.data_parallel_size > 0, "Data parallel size must be > 0"
|
||
assert self.tensor_parallel_size > 0, "Tensor parallel size must be > 0"
|
||
assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0"
|
||
assert self.ulysses_degree > 0, "Ulysses degree must be > 0"
|
||
assert self.ring_degree > 0, "Ring degree must be > 0"
|
||
assert self.allgather_degree > 0, "AllGather degree must be > 0"
|
||
assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0"
|
||
assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0"
|
||
assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, (
|
||
"vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, "
|
||
f"but got {self.vae_parallel_mode!r}."
|
||
)
|
||
if self.allgather_degree > 1:
|
||
assert self.ulysses_degree == 1 and self.ring_degree == 1, (
|
||
"AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. "
|
||
f"Got ulysses_degree={self.ulysses_degree}, ring_degree={self.ring_degree}, "
|
||
f"allgather_degree={self.allgather_degree}."
|
||
)
|
||
expected_sp_size = (
|
||
self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
|
||
)
|
||
assert self.sequence_parallel_size == expected_sp_size, (
|
||
f"Sequence parallel size must be {expected_sp_size}, but got {self.sequence_parallel_size}"
|
||
)
|
||
assert self.ulysses_mode in {"strict", "advanced_uaa"}, (
|
||
f"ulysses_mode must be one of {{'strict','advanced_uaa'}}, but got {self.ulysses_mode!r}."
|
||
)
|
||
|
||
# Validate HSDP configuration
|
||
if self.use_hsdp:
|
||
assert self.hsdp_replicate_size > 0, "HSDP replicate size must be > 0"
|
||
assert self.hsdp_shard_size > 0, "HSDP shard size must be > 0 (should be set in __post_init__)"
|
||
return self
|
||
|
||
def __post_init__(self) -> None:
|
||
if self.sequence_parallel_size is None:
|
||
self.sequence_parallel_size = (
|
||
self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree
|
||
)
|
||
|
||
# Calculate world_size from other parallelism dimensions
|
||
other_parallel_world_size = (
|
||
self.pipeline_parallel_size
|
||
* self.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:
|
||
if self.tensor_parallel_size > 1 or self.data_parallel_size > 1:
|
||
raise ValueError(
|
||
"HSDP (use_hsdp=True) cannot be used with TP or DP "
|
||
f"(tensor_parallel_size={self.tensor_parallel_size}, "
|
||
f"data_parallel_size={self.data_parallel_size}). "
|
||
"Set tensor_parallel_size=1 and data_parallel_size=1 when using HSDP."
|
||
)
|
||
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
|
||
|
||
@classmethod
|
||
def from_dict(cls, data: dict[str, Any]) -> "DiffusionParallelConfig":
|
||
"""
|
||
Create DiffusionParallelConfig from a dictionary.
|
||
|
||
Args:
|
||
data: Dictionary containing parallel configuration parameters
|
||
|
||
Returns:
|
||
DiffusionParallelConfig instance with parameters set from dict
|
||
"""
|
||
if not isinstance(data, dict):
|
||
raise TypeError(f"Expected parallel config dict, got {type(data)!r}")
|
||
return cls(**data)
|
||
|
||
|
||
@dataclass
|
||
class TransformerConfig:
|
||
"""Container for raw transformer configuration dictionaries."""
|
||
|
||
params: dict[str, Any] = field(default_factory=dict)
|
||
quant_method: str | None = None
|
||
quant_config: "QuantizationConfig | None" = None
|
||
|
||
@classmethod
|
||
def from_dict(cls, data: dict[str, Any]) -> "TransformerConfig":
|
||
if not isinstance(data, dict):
|
||
raise TypeError(f"Expected transformer config dict, got {type(data)!r}")
|
||
params = dict(data) # copy to avoid mutating caller's dict
|
||
|
||
quant_method: str | None = None
|
||
quant_config: QuantizationConfig | None = None
|
||
disk_qc = params.get("quantization_config")
|
||
if isinstance(disk_qc, dict):
|
||
raw_quant_method = disk_qc.get("quant_method", disk_qc.get("method"))
|
||
quant_config = build_quant_config(disk_qc)
|
||
if quant_config is not None:
|
||
quant_method = raw_quant_method if raw_quant_method is not None else quant_config.get_name()
|
||
|
||
return cls(params=params, quant_method=quant_method, quant_config=quant_config)
|
||
|
||
def to_dict(self) -> dict[str, Any]:
|
||
return dict(self.params)
|
||
|
||
def get(self, key: str, default: Any | None = None) -> Any:
|
||
return self.params.get(key, default)
|
||
|
||
def __getattr__(self, item: str) -> Any:
|
||
params = object.__getattribute__(self, "params")
|
||
try:
|
||
return params[item]
|
||
except KeyError as exc:
|
||
raise AttributeError(item) from exc
|
||
|
||
|
||
@dataclass
|
||
class DiffusionCacheConfig:
|
||
"""
|
||
Configuration for cache adapters (TeaCache, cache-dit, MagCache, etc.).
|
||
|
||
This dataclass provides a unified interface for cache configuration parameters.
|
||
It can be initialized from a dictionary and accessed via attributes.
|
||
|
||
Common parameters:
|
||
- TeaCache: rel_l1_thresh, coefficients (optional)
|
||
- cache-dit: Fn_compute_blocks, Bn_compute_blocks, max_warmup_steps,
|
||
residual_diff_threshold, enable_taylorseer, taylorseer_order,
|
||
scm_steps_mask_policy, scm_steps_policy
|
||
- MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio,
|
||
mag_ratios, mag_calibrate
|
||
- step_cache: step_cache_dit_enabled, velocity_sim_thresholds,
|
||
velocity_skip_countdowns, step_cache_dit_min_history
|
||
|
||
Example:
|
||
>>> # From dict (user-facing API) - partial config uses defaults for missing keys
|
||
>>> config = DiffusionCacheConfig.from_dict({"rel_l1_thresh": 0.3})
|
||
>>> # Access via attribute
|
||
>>> print(config.rel_l1_thresh) # 0.3 (from dict)
|
||
>>> print(config.Fn_compute_blocks) # 8 (default)
|
||
>>> # Empty dict defers model-specific defaults to the TeaCache backend
|
||
>>> default_config = DiffusionCacheConfig.from_dict({})
|
||
>>> print(default_config.rel_l1_thresh) # None
|
||
"""
|
||
|
||
# TeaCache parameters [tea_cache only]
|
||
# None defers to the model-specific TeaCache default (0.2 fallback).
|
||
rel_l1_thresh: float | None = None
|
||
coefficients: list[float] | None = None # Uses model-specific defaults if None
|
||
|
||
# MagCache parameters [mag_cache only]
|
||
# Default: 0.24 threshold for accumulated magnitude error
|
||
mag_threshold: float = 0.24
|
||
# Default: 5 maximum consecutive skip steps (K)
|
||
mag_max_skip_steps: int = 5
|
||
# Default: 0.1 fraction of initial steps where skipping is disabled (stability)
|
||
mag_retention_ratio: float = 0.1
|
||
# Default: None magnitude ratios (model-specific, required for inference)
|
||
mag_ratios: list[float] | None = None
|
||
# Default: False calibration mode (computes mag_ratios on first run)
|
||
mag_calibrate: bool = False
|
||
|
||
# cache-dit parameters [cache-dit only]
|
||
# Default: 1 forward compute block (optimized for single-transformer models)
|
||
# Use 1 as default instead of cache-dit's 8, optimized for single-transformer models
|
||
# This provides better performance while maintaining quality for most use cases
|
||
Fn_compute_blocks: int = 1
|
||
# Default: 0 backward compute blocks (no fusion by default)
|
||
Bn_compute_blocks: int = 0
|
||
# Default: 4 warmup steps (optimized for few-step distilled models like Z-Image with 8 steps)
|
||
# Use 4 as default warmup steps instead of 8 in cache-dit, making DBCache work
|
||
# for few-step distilled models (e.g., Z-Image with 8 steps)
|
||
max_warmup_steps: int = 4
|
||
# Default: -1 (unlimited cached steps) - DBCache disables caching when previous cached steps exceed this value
|
||
# to prevent precision degradation. Set to -1 for unlimited caching (cache-dit default).
|
||
max_cached_steps: int = -1
|
||
# Default: 0.24 residual difference threshold (higher for more aggressive caching)
|
||
# Use a relatively higher residual diff threshold (0.24) as default to allow more
|
||
# aggressive caching. This is safe because we have max_continuous_cached_steps limit.
|
||
# Without this limit, a lower threshold like 0.12 would be needed.
|
||
residual_diff_threshold: float = 0.24
|
||
# Default: Limit consecutive cached steps to 3 to prevent precision degradation
|
||
# This allows us to use a higher residual_diff_threshold for more aggressive caching
|
||
max_continuous_cached_steps: int = 3
|
||
# Default: Disable TaylorSeer (not suitable for few-step distilled models)
|
||
# TaylorSeer is not suitable for few-step distilled models, so we disable it by default.
|
||
# References:
|
||
# - From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers
|
||
# - Forecast then Calibrate: Feature Caching as ODE for Efficient Diffusion Transformers
|
||
enable_taylorseer: bool = False
|
||
# Default: 1st order TaylorSeer polynomial
|
||
taylorseer_order: int = 1
|
||
# Default: None SCM mask policy (disabled by default)
|
||
scm_steps_mask_policy: str | None = None
|
||
# Default: "dynamic" steps policy for adaptive caching
|
||
scm_steps_policy: str = "dynamic"
|
||
# Used by cache-dit for scm mask generation. If this value changes during inference,
|
||
# we will re-generate the scm mask and refresh the cache context.
|
||
num_inference_steps: int | None = None
|
||
# Force refresh the cache at a specific step index hint, useful for models like
|
||
# GLM-Image (image preprocessing step in editing mode).
|
||
force_refresh_step_hint: int | None = None
|
||
# Policy for force refresh: "once" refreshes only at the hint step,
|
||
# "repeat" refreshes every force_refresh_step_hint steps.
|
||
force_refresh_step_policy: str = "once"
|
||
|
||
# step_cache parameters [step_cache only] — DreamZero velocity schedule
|
||
step_cache_dit_enabled: bool = True
|
||
velocity_sim_thresholds: list[float] = field(default_factory=lambda: [0.95, 0.93])
|
||
velocity_skip_countdowns: list[int] = field(default_factory=lambda: [4, 2])
|
||
step_cache_dit_min_history: int = 2
|
||
step_cache_dit_max_history: int = 2
|
||
|
||
# Additional parameters that may be passed but not explicitly defined
|
||
_extra_params: dict[str, Any] = field(default_factory=dict, repr=False)
|
||
|
||
@classmethod
|
||
def from_dict(cls, data: dict[str, Any]) -> "DiffusionCacheConfig":
|
||
"""
|
||
Create DiffusionCacheConfig from a dictionary.
|
||
|
||
Args:
|
||
data: Dictionary containing cache configuration parameters
|
||
|
||
Returns:
|
||
DiffusionCacheConfig instance with parameters set from dict
|
||
"""
|
||
if not isinstance(data, dict):
|
||
raise TypeError(f"Expected cache config dict, got {type(data)!r}")
|
||
|
||
# Get all dataclass field names automatically
|
||
field_names = {f.name for f in fields(cls)}
|
||
|
||
# Extract parameters that match dataclass fields (excluding private fields)
|
||
known_params = {k: v for k, v in data.items() if k in field_names and not k.startswith("_")}
|
||
|
||
# Store extra parameters
|
||
extra_params = {k: v for k, v in data.items() if k not in field_names}
|
||
|
||
# Create instance with known params (missing ones will use defaults)
|
||
# Then update _extra_params after creation since it's a private field
|
||
instance = cls(**known_params, _extra_params=extra_params)
|
||
return instance
|
||
|
||
def __getattr__(self, item: str) -> Any:
|
||
"""
|
||
Allow access to extra parameters via attribute access.
|
||
|
||
This enables accessing parameters that weren't explicitly defined
|
||
in the dataclass fields but were passed in the dict.
|
||
"""
|
||
if item == "_extra_params" or item.startswith("_"):
|
||
return object.__getattribute__(self, item)
|
||
|
||
extra = object.__getattribute__(self, "_extra_params")
|
||
if item in extra:
|
||
return extra[item]
|
||
|
||
raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'")
|
||
|
||
|
||
def resolve_model_class_name(
|
||
model: str | None,
|
||
diffusion_load_format: str = "default",
|
||
revision: str | None = None,
|
||
) -> str | None:
|
||
"""Resolve the diffusion pipeline class name from the model config.
|
||
|
||
Read-only counterpart of ``OmniDiffusionConfig.enrich_config``, safe to call
|
||
client-side. Returns ``None`` if the pipeline can't be determined.
|
||
"""
|
||
from vllm.transformers_utils.config import get_hf_file_to_dict
|
||
|
||
from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
|
||
|
||
if not model:
|
||
return None
|
||
is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"}
|
||
|
||
# Diffusers models: read _class_name from the pipeline index.
|
||
model_index = get_diffusion_model_index(model, revision=revision)
|
||
if model_index is not None:
|
||
return model_index.get("_class_name")
|
||
if diffusion_load_format == "diffusers":
|
||
return "DiffusersAdapterPipeline"
|
||
|
||
# Other models: map model_type / architecture from config.json.
|
||
try:
|
||
cfg = get_hf_file_to_dict("config.json", model, revision=revision) or {}
|
||
except Exception:
|
||
cfg = {}
|
||
model_type = cfg.get("model_type")
|
||
architectures = cfg.get("architectures") or []
|
||
|
||
if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
|
||
return "BagelPipeline"
|
||
if (
|
||
model_type == "lance"
|
||
or "LancePipeline" in architectures
|
||
or cfg.get("model_name") == "Lance"
|
||
or is_lance_subfolder
|
||
):
|
||
return "LancePipeline"
|
||
if model_type == "neo_chat":
|
||
return "SenseNovaU1Pipeline"
|
||
if "BailingMM2NativeForConditionalGeneration" in architectures or model_type in (
|
||
"bailingmm_moe_v2_lite",
|
||
"ming_flash_omni",
|
||
"ming_flash_omni_thinker",
|
||
):
|
||
return "MingImagePipeline"
|
||
if model_type == "nextstep":
|
||
return "NextStep11Pipeline"
|
||
if model_type == "s2v":
|
||
return "WanS2VPipeline"
|
||
if model_type == "vla":
|
||
from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
|
||
|
||
return "DreamZeroPipeline" if _looks_like_dreamzero(model) else None
|
||
if len(architectures) == 1:
|
||
return architectures[0]
|
||
return None
|
||
|
||
|
||
@dataclass
|
||
class OmniDiffusionConfig:
|
||
# Model and path configuration (for convenience)
|
||
stage_id: int = 0
|
||
|
||
model: str | None = None
|
||
|
||
model_class_name: str | None = None
|
||
|
||
# Optional model-defined startup task. Pipelines may use this to select
|
||
# task-specific components or weights before serving requests.
|
||
task_type: str | None = None
|
||
|
||
dtype: torch.dtype = torch.bfloat16
|
||
|
||
model_config: dict[str, Any] = field(default_factory=dict)
|
||
tf_model_config: TransformerConfig = field(default_factory=TransformerConfig)
|
||
|
||
# Attention
|
||
diffusion_attention_config: "AttentionConfig" = field(default_factory=lambda: AttentionConfig())
|
||
fa_deterministic: bool = False
|
||
|
||
# Running mode
|
||
# mode: ExecutionMode = ExecutionMode.INFERENCE
|
||
|
||
# Workload type
|
||
# workload_type: WorkloadType = WorkloadType.T2V
|
||
|
||
# Cache strategy (legacy)
|
||
cache_strategy: str = "none"
|
||
parallel_config: DiffusionParallelConfig = field(default_factory=DiffusionParallelConfig)
|
||
|
||
# Cache backend configuration (NEW)
|
||
cache_backend: str = "none" # "tea_cache", "deep_cache", etc.
|
||
cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict)
|
||
enable_cache_dit_summary: bool = False
|
||
|
||
# Prompt-embedding cache. When enabled, ``DiffusionModelRunner`` wraps the
|
||
# pipeline's ``encode_prompt`` so repeated calls with identical prompt
|
||
# arguments (e.g. GRPO rollouts that sample the same prompt many times
|
||
# with different seeds) reuse the text-encoder output instead of re-running
|
||
# it. Safe against inputs that cannot be hashed (tensors, PIL images):
|
||
# those calls transparently bypass the cache.
|
||
enable_prompt_embed_cache: bool = False
|
||
prompt_embed_cache_size: int = 32
|
||
|
||
# Opt-in: route per-session world-model state through the shared
|
||
# SessionStateManager (RFC #4480) instead of the model's bespoke cache.
|
||
# Default off; the bespoke path remains the default. A *set*
|
||
# OMNI_DIFFUSION_SESSION_STATE_MANAGER environment variable overrides this
|
||
# in both directions. See docs/features/session_state_manager.md.
|
||
enable_session_state_manager: bool = False
|
||
|
||
# Distributed executor backend
|
||
distributed_executor_backend: str = "mp"
|
||
nccl_port: int | None = None
|
||
|
||
# Engine backend selection, resolved by ``DiffusionEngine.resolve_engine_class``
|
||
# (mirrors ``DiffusionExecutor.get_class``). Config files use a string:
|
||
# "default" -> DiffusionEngine, or an import path (set e.g. by a deploy
|
||
# config). Programmatic callers may pass a DiffusionEngine subclass
|
||
# directly — hence ``str | type`` (structured-config mirrors should expose
|
||
# the string form only).
|
||
engine_backend: str | type = "default"
|
||
|
||
# Local Diffusion KV ownership and cache-layout mode.
|
||
diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
|
||
|
||
# Optional override for the diffusion model runner class (import path).
|
||
# Precedence in the worker: this override > the runner declared by the
|
||
# selected engine class (``default_diffusion_model_runner_cls``) > the
|
||
# platform default. Never mutated by engines.
|
||
diffusion_model_runner_cls: str | None = None
|
||
|
||
# HuggingFace specific parameters
|
||
trust_remote_code: bool = False
|
||
revision: str | None = None
|
||
|
||
num_gpus: int | None = None
|
||
|
||
dist_timeout: int | None = None # timeout for torch.distributed
|
||
|
||
# pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False)
|
||
|
||
# LoRA parameters
|
||
lora_path: str | list[str] | None = None
|
||
lora_scale: float | list[float] = 1.0
|
||
lora_backend: LoRABackend = LoRABackend.PEFT # available choices: ["peft", "distill"]
|
||
max_cpu_loras: int | None = None
|
||
|
||
output_type: str = "pil"
|
||
|
||
# CPU offload parameters
|
||
# When enabled, DiT and encoders swap GPU access (mutual exclusion):
|
||
# - Text encoders run on GPU while DiT is on CPU
|
||
# - DiT runs on GPU while encoders are on CPU
|
||
enable_cpu_offload: bool = False
|
||
# Layer-wise offloading (block-level offloading) parameters
|
||
enable_layerwise_offload: bool = False
|
||
# Distributed layer-wise offloading with H2D + AllGather overlap (RFC-1)
|
||
enable_distributed_layerwise_offload: bool = False
|
||
# If True: shard weights 1/dp_size + AllGather (saves CPU memory, requires
|
||
# concurrent requests in DP mode). If False: each rank streams the standard
|
||
# loader's rank-local tensors (including TP-local shards) via H2D only.
|
||
# This avoids AllGather synchronization, while host memory follows the
|
||
# loader's existing rank-local layout instead of adding a second DP shard.
|
||
dlo_use_allgather: bool = True
|
||
# Build or join a node-local mmap cache after ordinary loading. This is
|
||
# used when direct checkpoint mmap is unavailable, or when direct CUDA
|
||
# registration requires a transform-complete final layout. The first
|
||
# version is no-AllGather only.
|
||
dlo_enable_runtime_cache: bool = False
|
||
# Shared local-disk root. None selects ~/.cache/vllm-omni/dlo-runtime-weights.
|
||
dlo_runtime_cache_dir: str | None = None
|
||
# Maximum wait for the per-layout POSIX writer lock.
|
||
dlo_runtime_cache_lock_timeout: float = 600.0
|
||
# Maximum file-backed memory a worker may register with CUDA. Zero keeps
|
||
# the bounded mmap-to-pinned staging path.
|
||
dlo_runtime_cache_pin_limit_gib: float = 0.0
|
||
# Leading main-DiT blocks kept resident by distributed layerwise offload.
|
||
dlo_resident_layers: int = 0
|
||
|
||
pin_cpu_memory: bool = True # Use pinned memory for faster transfers when offloading
|
||
|
||
# VAE memory optimization parameters
|
||
vae_use_slicing: bool = False
|
||
vae_use_tiling: bool = False
|
||
|
||
# STA (Sliding Tile Attention) parameters
|
||
mask_strategy_file_path: str | None = None
|
||
# STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
|
||
skip_time_steps: int = 15
|
||
|
||
# MoE kernel backend selection
|
||
moe_backend: str = "auto"
|
||
|
||
# Compilation
|
||
enforce_eager: bool = False
|
||
# Controls the generic compilation path used when a pipeline does not
|
||
# provide its own setup_compile() implementation.
|
||
diffusion_compile_granularity: str = "regional"
|
||
diffusion_compile_dynamic: bool = True
|
||
|
||
# Parallel weight loading (for faster diffusion model startup)
|
||
enable_multithread_weight_load: bool = True
|
||
num_weight_load_threads: int = 4
|
||
|
||
# Enable sleep mode
|
||
enable_sleep_mode: bool = False
|
||
|
||
disable_autocast: bool = False
|
||
|
||
# VSA parameters
|
||
VSA_sparsity: float = 0.0 # inference/validation sparsity
|
||
|
||
# V-MoBA parameters
|
||
moba_config_path: str | None = None
|
||
# moba_config: dict[str, Any] = field(default_factory=dict)
|
||
|
||
# Master port for distributed inference
|
||
# TODO: do not hard code
|
||
master_port: int | None = None
|
||
|
||
# Worker extension class for custom functionality
|
||
worker_extension_cls: str | None = None
|
||
|
||
# Custom pipeline arguments for custom pipelines
|
||
custom_pipeline_args: dict[str, Any] | None = None
|
||
|
||
# Diffusion model loading format
|
||
# "default", "custom_pipeline", "dummy", "diffusers" (HF diffusers adapter)
|
||
diffusion_load_format: str = "default"
|
||
|
||
# Diffusers adapter kwargs
|
||
# kwargs forwarded to DiffusionPipeline.from_pretrained()
|
||
diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict)
|
||
# kwargs forwarded to pipeline.__call__()
|
||
diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict)
|
||
# Actual diffusers pipeline object (to determine inputs of the dummy run)
|
||
diffusers_pipeline_cls: type[diffusers.DiffusionPipeline] | None = None # pyright: ignore[reportPrivateImportUsage]
|
||
|
||
# http server endpoint config, would be ignored in local mode
|
||
host: str | None = None
|
||
port: int | None = None
|
||
|
||
scheduler_port: int = 5555
|
||
|
||
# Stage verification
|
||
enable_stage_verification: bool = True
|
||
|
||
# Prompt text file for batch processing
|
||
prompt_file_path: str | None = None
|
||
|
||
# model paths for correct deallocation
|
||
model_paths: dict[str, str] = field(default_factory=dict)
|
||
model_loaded: dict[str, bool] = field(
|
||
default_factory=lambda: {
|
||
"transformer": True,
|
||
"vae": True,
|
||
}
|
||
)
|
||
override_transformer_cls_name: str | None = None
|
||
|
||
# # DMD parameters
|
||
# dmd_denoising_steps: List[int] | None = field(default=None)
|
||
|
||
# MoE parameters used by Wan2.2
|
||
boundary_ratio: float | None = None
|
||
# Scheduler flow_shift for Wan2.2 (12.0 for 480p, 5.0 for 720p)
|
||
flow_shift: float | None = None
|
||
|
||
# Support multi-image inputs and expose any model-specific request limit
|
||
# through a generic config field so serving code stays model-agnostic.
|
||
supports_multimodal_inputs: bool = False
|
||
max_multimodal_image_inputs: int | None = None
|
||
supports_mixed_reference_inputs: bool = False
|
||
|
||
log_level: str = "info"
|
||
|
||
# Omni configuration (injected from stage config)
|
||
omni_kv_config: dict[str, Any] = field(default_factory=dict)
|
||
additional_config: dict[str, Any] = field(default_factory=dict)
|
||
|
||
profiler_config: "ProfilerConfig | dict[str, Any] | None" = None
|
||
|
||
# Model-specific function for collecting CFG KV caches (set at runtime)
|
||
cfg_kv_collect_func: Any | None = None
|
||
|
||
# Quantization: str method name, dict config, QuantizationConfig, or None.
|
||
# str is resolved to {"method": <str>} internally.
|
||
# Per-component: {"transformer": {"method": "fp8"}, "vae": None}
|
||
quantization_config: str | QuantizationConfig | dict[str, Any] | None = None
|
||
# Explicit runtime override for ModelOpt FP8 diffusion checkpoints. This
|
||
# does not enable FP8 by itself; it only selects CUTLASS once the checkpoint
|
||
# has already resolved to vLLM's ModelOpt FP8 linear method.
|
||
force_cutlass_fp8: bool = False
|
||
|
||
# Diffusion attention KV cache dtype (not vLLM's --kv-cache-dtype for AR models).
|
||
# None = native dtype (no quantization).
|
||
# "fp8" = dynamic FP8 (float8_e4m3fn) quantization per forward pass.
|
||
# On Hopper+FA3: native FP8 attention (memory + compute savings).
|
||
# On other backends: no benefit, backends skip quantization.
|
||
diffusion_kv_cache_dtype: str | None = None
|
||
# Optional skip selectors for KV-cache quantization. Format: "0-9,20,25-30".
|
||
# Listed steps/layers skip quantization; others keep quantized execution.
|
||
diffusion_kv_cache_skip_steps: str | None = None
|
||
diffusion_kv_cache_skip_layers: str | None = None
|
||
diffusion_kv_cache_skip_step_indices: set[int] | None = None
|
||
diffusion_kv_cache_skip_layer_indices: set[int] | None = None
|
||
|
||
# Diffusion pipeline Profiling config
|
||
enable_diffusion_pipeline_profiler: bool = False
|
||
|
||
# Step mode settings
|
||
step_execution: bool = False
|
||
|
||
# Streaming mode settings
|
||
streaming_output: bool = False # Start (video) generation with initial prompt, but streaming output in chunks
|
||
|
||
# Maximum number of sequences to generate in a batch
|
||
max_num_seqs: int = 1
|
||
|
||
# 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 not math.isfinite(self.dlo_runtime_cache_lock_timeout) or self.dlo_runtime_cache_lock_timeout <= 0:
|
||
raise ValueError("dlo_runtime_cache_lock_timeout must be positive")
|
||
if not math.isfinite(self.dlo_runtime_cache_pin_limit_gib) or self.dlo_runtime_cache_pin_limit_gib < 0:
|
||
raise ValueError("dlo_runtime_cache_pin_limit_gib must be non-negative")
|
||
if self.dlo_enable_runtime_cache and not self.enable_distributed_layerwise_offload:
|
||
raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload")
|
||
if self.dlo_enable_runtime_cache and self.dlo_use_allgather:
|
||
raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False")
|
||
if self.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache:
|
||
raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True")
|
||
if self.diffusion_compile_granularity not in {"regional", "full"}:
|
||
raise ValueError(
|
||
"diffusion_compile_granularity must be 'regional' or 'full', "
|
||
f"got {self.diffusion_compile_granularity!r}"
|
||
)
|
||
if not isinstance(self.diffusion_compile_dynamic, bool):
|
||
raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
|
||
self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
|
||
|
||
if self.omni_kv_config is None:
|
||
self.omni_kv_config = {}
|
||
elif isinstance(self.omni_kv_config, Mapping):
|
||
self.omni_kv_config = dict(self.omni_kv_config)
|
||
else:
|
||
raise TypeError("omni_kv_config must be a mapping")
|
||
if self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER and self.omni_kv_config.get(
|
||
"need_recv_cache", False
|
||
):
|
||
raise ValueError(
|
||
"paged_scheduler Diffusion KV does not support imported AR KV in Phase 1; "
|
||
"disable need_recv_cache until connector-aware admission is implemented"
|
||
)
|
||
|
||
self.master_port = self._resolve_master_port()
|
||
self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0)
|
||
if not math.isfinite(self.request_batch_max_wait_ms) or self.request_batch_max_wait_ms < 0:
|
||
raise ValueError(
|
||
f"request_batch_max_wait_ms must be a finite non-negative number, got {self.request_batch_max_wait_ms}."
|
||
)
|
||
|
||
if isinstance(self.profiler_config, dict):
|
||
from vllm.config import ProfilerConfig
|
||
|
||
self.profiler_config = ProfilerConfig(**self.profiler_config)
|
||
|
||
if self.additional_config is None:
|
||
self.additional_config = {}
|
||
elif isinstance(self.additional_config, Mapping):
|
||
self.additional_config = dict(self.additional_config)
|
||
else:
|
||
raise TypeError(f"additional_config must be a mapping or None, got {type(self.additional_config)!r}")
|
||
|
||
# Convert parallel_config dict/DictConfig to DiffusionParallelConfig
|
||
# Use Mapping to handle both plain dicts and OmegaConf DictConfig
|
||
if isinstance(self.parallel_config, Mapping):
|
||
self.parallel_config = DiffusionParallelConfig.from_dict(dict(self.parallel_config))
|
||
elif not isinstance(self.parallel_config, DiffusionParallelConfig):
|
||
self.parallel_config = DiffusionParallelConfig()
|
||
|
||
if self.num_gpus is None:
|
||
if self.parallel_config is not None:
|
||
self.num_gpus = self.parallel_config.world_size
|
||
else:
|
||
self.num_gpus = 1
|
||
|
||
if self.num_gpus < self.parallel_config.world_size:
|
||
raise ValueError(
|
||
f"num_gpus ({self.num_gpus}) < parallel_config.world_size ({self.parallel_config.world_size})"
|
||
)
|
||
|
||
if self.diffusion_compile_granularity == "full":
|
||
incompatible_features = []
|
||
if self.parallel_config.use_hsdp:
|
||
incompatible_features.append("HSDP")
|
||
if self.parallel_config.sequence_parallel_size > 1:
|
||
incompatible_features.append("sequence parallelism")
|
||
if self.enable_cpu_offload:
|
||
incompatible_features.append("CPU offload")
|
||
if self.enable_layerwise_offload:
|
||
incompatible_features.append("layerwise offload")
|
||
if incompatible_features:
|
||
features = ", ".join(incompatible_features)
|
||
raise ValueError(
|
||
"diffusion_compile_granularity='full' is incompatible with "
|
||
f"{features}; use 'regional' compilation instead"
|
||
)
|
||
|
||
# Convert string dtype to torch.dtype if needed
|
||
if isinstance(self.dtype, str):
|
||
dtype_map = {
|
||
"auto": torch.bfloat16,
|
||
"bfloat16": torch.bfloat16,
|
||
"bf16": torch.bfloat16,
|
||
"float16": torch.float16,
|
||
"fp16": torch.float16,
|
||
"half": torch.float16,
|
||
"float32": torch.float32,
|
||
"fp32": torch.float32,
|
||
"float": torch.float32,
|
||
}
|
||
dtype_lower = self.dtype.lower()
|
||
if dtype_lower in dtype_map:
|
||
self.dtype = dtype_map[dtype_lower]
|
||
else:
|
||
logger.warning(f"Unknown dtype string '{self.dtype}', defaulting to bfloat16")
|
||
self.dtype = torch.bfloat16
|
||
|
||
# Convert cache_config dict to DiffusionCacheConfig if needed
|
||
if isinstance(self.cache_config, dict):
|
||
self.cache_config = DiffusionCacheConfig.from_dict(self.cache_config)
|
||
elif not isinstance(self.cache_config, DiffusionCacheConfig):
|
||
# If it's neither dict nor DiffusionCacheConfig, convert to empty config
|
||
self.cache_config = DiffusionCacheConfig()
|
||
|
||
# Auto-detect quantization from TransformerConfig if not explicitly set.
|
||
# This covers the case where tf_model_config is passed at construction
|
||
# time. For late (post-construction) assignment, callers should use
|
||
# set_tf_model_config() which propagates quant_config automatically.
|
||
self._propagate_quantization_from_tf_config(self.tf_model_config)
|
||
|
||
# Resolve quantization_config: str/dict -> QuantizationConfig via build_quant_config.
|
||
if self.quantization_config is not None:
|
||
if isinstance(self.quantization_config, QuantizationConfig):
|
||
pass # Already built
|
||
elif isinstance(self.quantization_config, str):
|
||
self.quantization_config = build_quant_config(self.quantization_config)
|
||
elif isinstance(self.quantization_config, Mapping):
|
||
self.quantization_config = build_quant_config(dict(self.quantization_config))
|
||
else:
|
||
raise TypeError(
|
||
f"quantization_config must be str, dict, QuantizationConfig, or None, "
|
||
f"got {type(self.quantization_config)!r}"
|
||
)
|
||
|
||
# Match vLLM's config flow: parse entrypoint shorthands before the
|
||
# config object is built, and keep a single runtime truth source.
|
||
self.diffusion_attention_config = build_attention_config(self.diffusion_attention_config)
|
||
self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps)
|
||
self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers)
|
||
|
||
if self.max_cpu_loras is None:
|
||
self.max_cpu_loras = 1
|
||
elif self.max_cpu_loras < 1:
|
||
raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA")
|
||
|
||
if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs):
|
||
raise ValueError(
|
||
"diffusers_load_kwargs and diffusers_call_kwargs are only "
|
||
"valid together with diffusion_load_format=diffusers"
|
||
)
|
||
|
||
# when use hf offline, replace model to local model path
|
||
# align with ar stage behavior, see vllm/engine/arg_utils.py
|
||
if huggingface_hub.constants.HF_HUB_OFFLINE and self.model:
|
||
model_id = self.model
|
||
self.model = get_model_path(self.model, self.revision)
|
||
if model_id != self.model:
|
||
logger.info(
|
||
"HF_HUB_OFFLINE is True, replace model_id [%s] to model_path [%s]",
|
||
model_id,
|
||
self.model,
|
||
)
|
||
|
||
def _propagate_quantization_from_tf_config(self, tf_config: "TransformerConfig") -> None:
|
||
if tf_config.quant_config is None:
|
||
return
|
||
|
||
is_checkpoint_fp8 = bool(getattr(tf_config.quant_config, "is_checkpoint_fp8_serialized", False))
|
||
is_checkpoint_nvfp4 = bool(getattr(tf_config.quant_config, "is_checkpoint_nvfp4_serialized", False))
|
||
should_use_checkpoint_config = (
|
||
self.quantization_config is None
|
||
or (is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config))
|
||
or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config))
|
||
)
|
||
if should_use_checkpoint_config:
|
||
self.quantization_config = tf_config.quant_config
|
||
logger.info(
|
||
"Auto-detected quantization '%s' from model config",
|
||
tf_config.quant_method,
|
||
)
|
||
|
||
@staticmethod
|
||
def _is_generic_fp8_quant_config(quant_config: object) -> bool:
|
||
if isinstance(quant_config, str):
|
||
return quant_config.lower() == "fp8"
|
||
if isinstance(quant_config, Mapping):
|
||
method = quant_config.get("method", quant_config.get("quant_method"))
|
||
return isinstance(method, str) and method.lower() == "fp8"
|
||
if hasattr(quant_config, "get_name"):
|
||
return quant_config.get_name() == "fp8"
|
||
return False
|
||
|
||
@staticmethod
|
||
def _is_generic_nvfp4_quant_config(quant_config: object) -> bool:
|
||
if isinstance(quant_config, str):
|
||
return quant_config.lower() in {"fp4", "nvfp4", "modelopt_fp4"}
|
||
if isinstance(quant_config, Mapping):
|
||
method = quant_config.get("method", quant_config.get("quant_method"))
|
||
return isinstance(method, str) and method.lower() in {"fp4", "nvfp4", "modelopt_fp4"}
|
||
if hasattr(quant_config, "get_name"):
|
||
return quant_config.get_name() == "modelopt_fp4"
|
||
return False
|
||
|
||
def set_tf_model_config(self, tf_config: "TransformerConfig") -> None:
|
||
"""Assign `tf_model_config` and propagate quantization if detected.
|
||
|
||
In the normal startup flow `OmniDiffusionConfig` is created
|
||
*before* the transformer `config.json` is loaded from disk, so
|
||
`__post_init__` sees an empty `TransformerConfig`. Callers
|
||
that load the config later should use this method instead of bare
|
||
assignment so that an embedded `quant_config` is propagated to
|
||
`self.quantization_config` automatically.
|
||
|
||
Args:
|
||
tf_config: Transformer configuration, typically built via
|
||
`TransformerConfig.from_dict`.
|
||
"""
|
||
self.tf_model_config = tf_config
|
||
self._propagate_quantization_from_tf_config(tf_config)
|
||
self._propagate_skip_softmax_calibration(tf_config)
|
||
|
||
def _propagate_skip_softmax_calibration(self, tf_config: "TransformerConfig") -> None:
|
||
cfg = getattr(self, "diffusion_attention_config", None)
|
||
if not isinstance(cfg, AttentionConfig):
|
||
return
|
||
specs = [s for s in (cfg.default, *cfg.per_role.values()) if s is not None]
|
||
|
||
from vllm_omni.diffusion.attention.backends.trtllm_calibration import (
|
||
propagate_skip_softmax_calibration,
|
||
)
|
||
|
||
propagate_skip_softmax_calibration(specs, self.model, tf_config)
|
||
|
||
def update_multimodal_support(self) -> None:
|
||
# Resolve serving-visible multimodal behavior from shared metadata
|
||
# instead of importing concrete pipeline modules into the config layer.
|
||
metadata = get_diffusion_model_metadata(self.model_class_name)
|
||
self.supports_multimodal_inputs = metadata.supports_multimodal_inputs
|
||
self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs
|
||
self.supports_mixed_reference_inputs = metadata.supports_mixed_reference_inputs
|
||
|
||
@staticmethod
|
||
def _looks_like_lance_subfolder(model: str | None) -> bool:
|
||
"""Return True when ``--model`` points at a Lance per-component subfolder.
|
||
|
||
Lance's HF repo bundles ``Lance_3B/``, ``Lance_3B_Video/`` and
|
||
``Qwen2.5-VL-ViT/`` under a single top-level ``config.json``; users may
|
||
reasonably hand the AR-style sub-checkpoint path directly. The
|
||
``LancePipeline`` constructor knows to walk up to the repo root from
|
||
either subfolder name.
|
||
"""
|
||
if not model:
|
||
return False
|
||
base = os.path.basename(str(model).rstrip("/"))
|
||
return base in {"Lance_3B", "Lance_3B_Video"}
|
||
|
||
def enrich_config(self) -> None:
|
||
"""Load model metadata from HuggingFace and populate config fields.
|
||
|
||
Diffusers-style models expose a pipeline index with ``_class_name``.
|
||
Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``,
|
||
so we fall back to reading that and mapping model_type manually.
|
||
"""
|
||
from vllm.transformers_utils.config import get_hf_file_to_dict
|
||
|
||
from vllm_omni.diffusion.utils.hf_utils import get_diffusion_model_index
|
||
|
||
# Default model_class_name for diffusers adapter
|
||
if self.model_class_name is None and self.diffusion_load_format == "diffusers":
|
||
self.model_class_name = "DiffusersAdapterPipeline"
|
||
|
||
try:
|
||
config_dict = get_diffusion_model_index(
|
||
self.model,
|
||
revision=self.revision,
|
||
)
|
||
if config_dict is not None:
|
||
if self.model_class_name is None:
|
||
self.model_class_name = config_dict.get("_class_name", None)
|
||
self.update_multimodal_support()
|
||
|
||
# Skip transformer config loading for diffusers adapter
|
||
# (non-DiT models don't have a separate transformer folder/config)
|
||
if self.diffusion_load_format == "diffusers":
|
||
self.set_tf_model_config(TransformerConfig())
|
||
try:
|
||
diffusers_pipeline_cls_name = config_dict["_class_name"]
|
||
self.diffusers_pipeline_cls = getattr(diffusers, diffusers_pipeline_cls_name)
|
||
except (KeyError, AttributeError) as exc:
|
||
logger.warning(
|
||
"Could not find a valid _class_name in the Diffusers pipeline index: %s. "
|
||
"Without the underlying pipeline class the dummy run may omit required inputs.",
|
||
exc,
|
||
)
|
||
else:
|
||
from vllm_omni.model_extras import get_transformer_config_subfolder
|
||
|
||
transformer_subfolder = get_transformer_config_subfolder(
|
||
self.model_class_name,
|
||
model=self.model,
|
||
revision=self.revision,
|
||
)
|
||
tf_config_dict = get_hf_file_to_dict(
|
||
f"{transformer_subfolder}/config.json",
|
||
self.model,
|
||
revision=self.revision,
|
||
)
|
||
if tf_config_dict is None:
|
||
tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model, revision=self.revision)
|
||
if tf_config_dict is not None:
|
||
self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
|
||
else:
|
||
self.set_tf_model_config(TransformerConfig())
|
||
else:
|
||
raise FileNotFoundError("Diffusers pipeline index not found")
|
||
except (AttributeError, OSError, ValueError, FileNotFoundError):
|
||
# Skip transformer config loading for diffusers adapter
|
||
# (non-DiT models don't have a separate transformer folder/config)
|
||
if self.diffusion_load_format == "diffusers":
|
||
self.set_tf_model_config(TransformerConfig())
|
||
logger.warning(
|
||
"Could not find a valid pipeline index per Diffusers format. "
|
||
"This model is likely unsupported by the diffusers backend. "
|
||
"Also, without knowing the underlying pipeline class from its index, "
|
||
"the dummy run will input only text prompt, which may cause errors for pipelines "
|
||
"that require additional inputs."
|
||
)
|
||
else:
|
||
cfg = get_hf_file_to_dict("config.json", self.model, revision=self.revision)
|
||
if cfg is None:
|
||
# Lance ships its top-level config.json one directory above
|
||
# the per-checkpoint subfolders (``Lance_3B/`` or
|
||
# ``Lance_3B_Video/``). Try to recover that case before
|
||
# raising.
|
||
if self._looks_like_lance_subfolder(self.model):
|
||
self.model_class_name = "LancePipeline"
|
||
self.set_tf_model_config(TransformerConfig())
|
||
self.update_multimodal_support()
|
||
return
|
||
raise ValueError(f"Could not find config.json or a Diffusers pipeline index for {self.model}")
|
||
|
||
self.set_tf_model_config(TransformerConfig.from_dict(cfg))
|
||
model_type = cfg.get("model_type")
|
||
architectures = cfg.get("architectures") or []
|
||
|
||
if model_type == "bagel" or "BagelForConditionalGeneration" in architectures:
|
||
self.model_class_name = "BagelPipeline"
|
||
self.set_tf_model_config(TransformerConfig())
|
||
self.update_multimodal_support()
|
||
elif (
|
||
model_type == "lance"
|
||
or "LancePipeline" in architectures
|
||
or cfg.get("model_name") == "Lance"
|
||
or self._looks_like_lance_subfolder(self.model)
|
||
):
|
||
# Lance ships a non-HF top-level config.json (model_name only)
|
||
# plus per-component subfolders; resolve to the Lance pipeline.
|
||
# Also accept --model pointing directly at the ``Lance_3B`` or
|
||
# ``Lance_3B_Video`` subfolder by walking up to the repo root.
|
||
self.model_class_name = "LancePipeline"
|
||
self.set_tf_model_config(TransformerConfig())
|
||
self.update_multimodal_support()
|
||
elif model_type == "neo_chat":
|
||
self.model_class_name = "SenseNovaU1Pipeline"
|
||
self.tf_model_config = TransformerConfig()
|
||
self.update_multimodal_support()
|
||
elif "BailingMM2NativeForConditionalGeneration" in architectures or model_type in (
|
||
"bailingmm_moe_v2_lite",
|
||
"ming_flash_omni",
|
||
"ming_flash_omni_thinker",
|
||
):
|
||
# Ming-flash-omni-2.0 — imagegen stage uses the custom
|
||
# ``MingImagePipeline`` (ZImage DiT + Qwen2 connector). See
|
||
# vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py.
|
||
self.model_class_name = "MingImagePipeline"
|
||
self.tf_model_config = TransformerConfig()
|
||
self.update_multimodal_support()
|
||
elif model_type == "nextstep":
|
||
if self.model_class_name is None:
|
||
self.model_class_name = "NextStep11Pipeline"
|
||
self.set_tf_model_config(TransformerConfig())
|
||
self.update_multimodal_support()
|
||
elif model_type == "s2v":
|
||
if self.model_class_name is None:
|
||
self.model_class_name = "WanS2VPipeline"
|
||
self.tf_model_config = TransformerConfig()
|
||
self.update_multimodal_support()
|
||
elif model_type == "Gr00tN1d7" or "Gr00tN1d7" in architectures:
|
||
self.model_class_name = "Gr00tN1d7Pipeline"
|
||
self.set_tf_model_config(TransformerConfig())
|
||
self.update_multimodal_support()
|
||
elif model_type == "vla":
|
||
from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero
|
||
|
||
if _looks_like_dreamzero(self.model):
|
||
self.model_class_name = "DreamZeroPipeline"
|
||
self.set_tf_model_config(TransformerConfig())
|
||
self.update_multimodal_support()
|
||
else:
|
||
raise
|
||
elif architectures and len(architectures) == 1:
|
||
architecture = architectures[0]
|
||
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
||
|
||
if (
|
||
self.model_class_name is None
|
||
or DiffusionModelRegistry._try_load_model_cls(architecture) is not None
|
||
):
|
||
self.model_class_name = architecture
|
||
else:
|
||
raise
|
||
|
||
@classmethod
|
||
def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
|
||
kwargs = normalize_omni_diffusion_kwargs(kwargs)
|
||
|
||
# Filter kwargs to only include valid fields
|
||
valid_fields = {f.name for f in fields(cls)}
|
||
filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}
|
||
|
||
instance = cls(**filtered_kwargs)
|
||
return instance
|
||
|
||
|
||
@dataclass
|
||
class DiffusionOutput:
|
||
"""
|
||
Final output (after pipeline completion)
|
||
"""
|
||
|
||
output: torch.Tensor | tuple[Any, ...] | dict[str, Any] | None = None
|
||
|
||
# Legacy compatibility fields. New pipeline-specific payloads should be
|
||
# carried by output["payload"] instead.
|
||
trajectory_timesteps: torch.Tensor | dict[str, Any] | None = None
|
||
trajectory_latents: torch.Tensor | dict[str, Any] | None = None
|
||
trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None
|
||
trajectory_decoded: list[Image.Image] | None = None
|
||
async_output_id: str | None = None
|
||
error: str | None = None
|
||
error_status_code: int | None = None
|
||
error_type: str | None = None
|
||
aborted: bool = False
|
||
abort_message: str | None = None
|
||
|
||
post_process_func: Callable[..., Any] | None = None
|
||
|
||
# logged timings info, directly from Req.timings
|
||
# timings: Optional["RequestTimings"] = None
|
||
|
||
# Streaming info (the defaults should make sense for non-streaming mode)
|
||
finished: bool = True
|
||
chunk_index: int = 0
|
||
total_chunks: int = 1
|
||
started_event_ids: list[str] = field(default_factory=list)
|
||
active_event_ids: list[str] = field(default_factory=list)
|
||
completed_event_ids: list[str] = field(default_factory=list)
|
||
|
||
# logged duration of stages
|
||
stage_durations: dict[str, float] = field(default_factory=dict)
|
||
|
||
# memory usage info
|
||
peak_memory_mb: float = 0.0
|
||
|
||
# When True, move tensor fields to CPU at construction time. Useful when
|
||
# the output is shipped across process boundaries (e.g. step-execution
|
||
# mode) and the receiving side must not initialise a stray CUDA context.
|
||
to_cpu: bool = False
|
||
|
||
def __post_init__(self) -> None:
|
||
if not self.to_cpu:
|
||
return
|
||
|
||
def _maybe_to_cpu(value: Any) -> Any:
|
||
if isinstance(value, torch.Tensor):
|
||
return value.detach().cpu()
|
||
if isinstance(value, dict):
|
||
return {key: _maybe_to_cpu(item) for key, item in value.items()}
|
||
if isinstance(value, list):
|
||
return [_maybe_to_cpu(item) for item in value]
|
||
if isinstance(value, tuple):
|
||
return tuple(_maybe_to_cpu(item) for item in value)
|
||
return value
|
||
|
||
self.output = _maybe_to_cpu(self.output)
|
||
self.trajectory_timesteps = _maybe_to_cpu(self.trajectory_timesteps)
|
||
self.trajectory_latents = _maybe_to_cpu(self.trajectory_latents)
|
||
self.trajectory_log_probs = _maybe_to_cpu(self.trajectory_log_probs)
|
||
|
||
@classmethod
|
||
def from_exception(cls, exc: BaseException) -> "DiffusionOutput":
|
||
status_code, error_type = client_error_metadata(exc)
|
||
return cls(
|
||
error=str(exc),
|
||
error_status_code=status_code,
|
||
error_type=error_type,
|
||
)
|
||
|
||
|
||
class AsyncOutputKind(Enum):
|
||
"""Message kind for ``AsyncDiffusionOutput`` — routes async results.
|
||
|
||
* ``RPC_RESULT`` — ordinary RPC return (sleep, wake, profile, etc.)
|
||
* ``COMPUTE_DONE`` — worker forward finished, GPU can start next request
|
||
* ``OUTPUT_READY`` — background D2H/SHM packing finished, final output
|
||
is available via ``async_output_id``
|
||
"""
|
||
|
||
RPC_RESULT = "rpc_result"
|
||
COMPUTE_DONE = "compute_done"
|
||
OUTPUT_READY = "output_ready"
|
||
|
||
|
||
@dataclass
|
||
class AsyncDiffusionOutput:
|
||
"""Async protocol envelope for ``result_mq`` messages.
|
||
|
||
In request-mode (``step_execution=False``), all ``result_mq``
|
||
messages use this envelope. The ``kind`` field routes the message
|
||
to the correct consumer.
|
||
"""
|
||
|
||
kind: AsyncOutputKind
|
||
rpc_id: str | None = None
|
||
async_output_id: str | None = None
|
||
result: Any | None = None
|
||
output: DiffusionOutput | None = None
|
||
error: str | None = None
|
||
|
||
|
||
class DiffusionRequestAbortedError(RuntimeError):
|
||
"""Raised when a diffusion request ends via user-visible abort."""
|
||
|
||
|
||
def _in_range(value: Any, name: str, lo: float, hi: float | None) -> float | None:
|
||
"""Validate an optional numeric control at config-parse time."""
|
||
if value is None:
|
||
return None
|
||
x = float(value)
|
||
if not math.isfinite(x) or x < lo or (hi is not None and x > hi):
|
||
rng = f"in [{lo}, {hi}]" if hi is not None else f">= {lo}"
|
||
raise ValueError(f"{name} must be finite and {rng}; got {value!r}.")
|
||
return x
|
||
|
||
|
||
@dataclass
|
||
class SkipSoftmaxSpec:
|
||
"""User-facing skip-softmax controls for the TRTLLM_ATTN backend.
|
||
|
||
``target_sparsity`` uses the checkpoint's calibrated curve (a·exp(b·s)); ``threshold``
|
||
is the calibration-free absolute skip threshold. They are mutually exclusive.
|
||
"""
|
||
|
||
target_sparsity: float | None = None
|
||
threshold: float | None = None
|
||
disabled_until_timestep: float = 0.0
|
||
|
||
def __post_init__(self) -> None:
|
||
self.target_sparsity = _in_range(self.target_sparsity, "skip_softmax.target_sparsity", 0.0, 1.0)
|
||
self.threshold = _in_range(self.threshold, "skip_softmax.threshold", 0.0, None)
|
||
self.disabled_until_timestep = (
|
||
_in_range(self.disabled_until_timestep, "skip_softmax.disabled_until_timestep", 0.0, 1.0) or 0.0
|
||
)
|
||
if self.target_sparsity is not None and self.threshold is not None:
|
||
raise ValueError("skip_softmax: set either target_sparsity or threshold, not both.")
|
||
|
||
@property
|
||
def enabled(self) -> bool:
|
||
return self.threshold is not None or self.target_sparsity is not None
|
||
|
||
|
||
@dataclass
|
||
class AttnQuantSpec:
|
||
dtype_qk: str | None = None
|
||
dtype_vo: str | None = None
|
||
q_block_size: int = 1
|
||
k_block_size: int = 16
|
||
flashinfer_backend: str | None = None
|
||
|
||
_VALID_DTYPES = frozenset({"float16", "bfloat16", "int8", "fp8_e4m3"})
|
||
_VALID_BLOCK_SIZES = frozenset({1, 4, 16})
|
||
|
||
def __post_init__(self) -> None:
|
||
for name, v in (("dtype_qk", self.dtype_qk), ("dtype_vo", self.dtype_vo)):
|
||
if v is not None and v not in self._VALID_DTYPES:
|
||
raise ValueError(f"quant.{name}={v!r} unsupported; use one of {sorted(self._VALID_DTYPES)}.")
|
||
for name, v in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)):
|
||
if v not in self._VALID_BLOCK_SIZES:
|
||
raise ValueError(
|
||
f"quant.{name}={v!r} unsupported; kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}."
|
||
)
|
||
|
||
@property
|
||
def enabled(self) -> bool:
|
||
return self.dtype_qk is not None or self.dtype_vo is not None
|
||
|
||
|
||
# Backends that select key blocks instead of attending densely, and so accept
|
||
# a ``block_sparse`` spec. Each maps the same knobs onto its own kernel.
|
||
BLOCK_SPARSE_BACKENDS = frozenset({"RAINFUSION_ATTN"})
|
||
|
||
|
||
@dataclass
|
||
class BlockSparseSpec:
|
||
"""User-facing controls shared by block-sparse attention backends.
|
||
|
||
``sparsity`` is the nominal fraction of key blocks dropped per query block;
|
||
``start_step`` keeps the first N denoise steps dense and ``skip_layers`` (an
|
||
index selector such as "0-3,38") exempts individual DiT blocks. Those two are
|
||
the accuracy knobs to trade back quality at a fixed ``sparsity``.
|
||
"""
|
||
|
||
sparsity: float = 0.8
|
||
start_step: int = 0
|
||
skip_layers: str | list[int] | None = None
|
||
skip_layer_indices: set[int] | None = field(default=None, repr=False)
|
||
|
||
def __post_init__(self) -> None:
|
||
self.sparsity = _in_range(self.sparsity, "block_sparse.sparsity", 0.0, 1.0) or 0.0
|
||
if self.start_step < 0:
|
||
raise ValueError(f"block_sparse.start_step must be >= 0; got {self.start_step!r}.")
|
||
self.skip_layer_indices = parse_kv_cache_skip_selector(self.skip_layers)
|
||
|
||
|
||
@dataclass
|
||
class AttentionSpec:
|
||
"""Specifies a backend and its typed backend-specific config for one attention role."""
|
||
|
||
backend: str
|
||
skip_softmax: SkipSoftmaxSpec | None = None
|
||
quant: AttnQuantSpec | None = None
|
||
block_sparse: BlockSparseSpec | None = None
|
||
skip_calibration: dict | None = field(default=None, repr=False)
|
||
|
||
def __post_init__(self) -> None:
|
||
if not isinstance(self.backend, str):
|
||
raise TypeError(f"Expected str for AttentionSpec.backend, got {type(self.backend)!r}")
|
||
self.skip_softmax = self._coerce(self.skip_softmax, SkipSoftmaxSpec, "skip_softmax")
|
||
self.quant = self._coerce(self.quant, AttnQuantSpec, "quant")
|
||
self.block_sparse = self._coerce(self.block_sparse, BlockSparseSpec, "block_sparse")
|
||
if self.skip_softmax is not None and self.backend.upper() != "TRTLLM_ATTN":
|
||
raise ValueError(
|
||
f"skip_softmax is only supported by the TRTLLM_ATTN backend, but backend={self.backend!r}. "
|
||
"Remove skip_softmax or set backend to TRTLLM_ATTN."
|
||
)
|
||
if self.quant is not None and self.backend.upper() not in ("TRTLLM_ATTN", "FLASHINFER_ATTN"):
|
||
raise ValueError(
|
||
f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but "
|
||
f"backend={self.backend!r}. Remove quant or set a supported backend."
|
||
)
|
||
if self.backend.upper() in BLOCK_SPARSE_BACKENDS:
|
||
# Selecting the backend is the opt-in; without an explicit block the
|
||
# defaults apply rather than silently running dense.
|
||
self.block_sparse = self.block_sparse or BlockSparseSpec()
|
||
elif self.block_sparse is not None:
|
||
raise ValueError(
|
||
f"block_sparse is only supported by the {sorted(BLOCK_SPARSE_BACKENDS)} backends, but "
|
||
f"backend={self.backend!r}. Remove block_sparse or set a supported backend."
|
||
)
|
||
|
||
@staticmethod
|
||
def _coerce(value: Any, cls: type, field_name: str) -> Any:
|
||
if value is None or isinstance(value, cls):
|
||
return value
|
||
if isinstance(value, Mapping):
|
||
return cls(**dict(value))
|
||
raise TypeError(f"Expected dict or {cls.__name__} for {field_name}, got {type(value)!r}")
|
||
|
||
def backend_kwargs(self) -> dict[str, Any] | None:
|
||
"""Serialize typed backend config into the kwargs dict the backend impl consumes."""
|
||
kw: dict[str, Any] = {}
|
||
if self.skip_softmax is not None:
|
||
ss = self.skip_softmax
|
||
if ss.threshold is not None:
|
||
kw["skip_softmax_threshold"] = ss.threshold
|
||
if ss.target_sparsity is not None:
|
||
kw["target_sparsity"] = ss.target_sparsity
|
||
if ss.disabled_until_timestep:
|
||
kw["disabled_until_timestep"] = ss.disabled_until_timestep
|
||
if self.quant is not None and self.quant.enabled:
|
||
q = self.quant
|
||
quant_kw: dict[str, Any] = {
|
||
"dtype_qk": q.dtype_qk,
|
||
"q_block_size": q.q_block_size,
|
||
"k_block_size": q.k_block_size,
|
||
}
|
||
if q.dtype_vo is not None:
|
||
quant_kw["dtype_vo"] = q.dtype_vo
|
||
if q.flashinfer_backend is not None:
|
||
quant_kw["flashinfer_backend"] = q.flashinfer_backend
|
||
kw["quant"] = quant_kw
|
||
if self.block_sparse is not None:
|
||
bs = self.block_sparse
|
||
kw["sparsity"] = bs.sparsity
|
||
kw["start_step"] = bs.start_step
|
||
if bs.skip_layer_indices:
|
||
kw["skip_layers"] = sorted(bs.skip_layer_indices)
|
||
return kw or None
|
||
|
||
|
||
@dataclass
|
||
class AttentionConfig:
|
||
"""Per-role attention backend configuration.
|
||
|
||
Lookup precedence for a given (role, role_category):
|
||
1. per_role[role] — exact match
|
||
2. per_role[role_category] — category fallback (e.g. "ltx2.audio_to_video" → "cross")
|
||
3. default — global default
|
||
4. platform default — unchanged platform logic
|
||
"""
|
||
|
||
default: AttentionSpec | None = None
|
||
per_role: dict[str, AttentionSpec] = field(default_factory=dict)
|
||
|
||
def __post_init__(self) -> None:
|
||
if self.default is not None:
|
||
self.default = self._coerce_spec_or_none(self.default, "default")
|
||
|
||
normalized_per_role: dict[str, AttentionSpec] = {}
|
||
for role_key, spec_data in self._normalize_per_role_mapping(self.per_role).items():
|
||
spec = self._coerce_spec_or_none(spec_data, f"per_role[{role_key!r}]")
|
||
if spec is not None:
|
||
normalized_per_role[role_key] = spec
|
||
self.per_role = normalized_per_role
|
||
|
||
@staticmethod
|
||
def _coerce_spec(spec_data: Any, field_name: str) -> AttentionSpec:
|
||
if isinstance(spec_data, AttentionSpec):
|
||
return spec_data
|
||
if isinstance(spec_data, str):
|
||
return AttentionSpec(backend=spec_data)
|
||
if isinstance(spec_data, Mapping):
|
||
return AttentionSpec(**dict(spec_data))
|
||
raise TypeError(f"Expected str, dict, or AttentionSpec for {field_name}, got {type(spec_data)!r}")
|
||
|
||
@classmethod
|
||
def _coerce_spec_or_none(cls, spec_data: Any, field_name: str) -> AttentionSpec | None:
|
||
spec = cls._coerce_spec(spec_data, field_name)
|
||
if spec.backend.lower() == "auto":
|
||
return None
|
||
return spec
|
||
|
||
@classmethod
|
||
def _normalize_per_role_mapping(cls, raw_per_role: Any) -> dict[str, Any]:
|
||
if raw_per_role is None:
|
||
return {}
|
||
if not isinstance(raw_per_role, Mapping):
|
||
raise TypeError(f"Expected dict for AttentionConfig.per_role, got {type(raw_per_role)!r}")
|
||
|
||
normalized: dict[str, Any] = {}
|
||
for role_key, spec_data in raw_per_role.items():
|
||
cls._flatten_per_role_entry([role_key], spec_data, normalized)
|
||
return normalized
|
||
|
||
@classmethod
|
||
def _flatten_per_role_entry(
|
||
cls,
|
||
path: list[str],
|
||
node: Any,
|
||
normalized: dict[str, Any],
|
||
) -> None:
|
||
role = ".".join(path)
|
||
if not isinstance(node, Mapping):
|
||
normalized[role] = node
|
||
return
|
||
|
||
spec_keys = {"backend", "skip_softmax", "quant", "block_sparse"}
|
||
node_dict = dict(node)
|
||
node_keys = set(node_dict)
|
||
if node_keys & spec_keys:
|
||
if not node_keys <= spec_keys:
|
||
raise ValueError(
|
||
f"Invalid per_role entry for role {role!r}: cannot mix backend/skip_softmax with nested role keys."
|
||
)
|
||
normalized[role] = node_dict
|
||
return
|
||
|
||
if not node_dict:
|
||
raise ValueError(f"Empty per_role entry for role {role!r}")
|
||
|
||
for child_key, child_value in node_dict.items():
|
||
cls._flatten_per_role_entry([*path, child_key], child_value, normalized)
|
||
|
||
def resolve_with_source(
|
||
self,
|
||
role: str = "self",
|
||
role_category: str | None = None,
|
||
) -> tuple[AttentionSpec | None, str | None]:
|
||
"""Resolve the AttentionSpec and report which config entry matched."""
|
||
spec = self.per_role.get(role)
|
||
if spec is not None:
|
||
return spec, f"attention_config.per_role[{role!r}]"
|
||
if role_category is not None:
|
||
spec = self.per_role.get(role_category)
|
||
if spec is not None:
|
||
return spec, f"attention_config.per_role[{role_category!r}] (role_category fallback)"
|
||
if self.default is not None:
|
||
return self.default, "attention_config.default"
|
||
return None, None
|
||
|
||
|
||
def parse_attention_config(
|
||
attention_config: AttentionConfig | Mapping[str, Any] | None = None,
|
||
*,
|
||
attention_backend: str | None = None,
|
||
) -> AttentionConfig:
|
||
"""Pure type-conversion: coerce *attention_config* to an AttentionConfig.
|
||
|
||
Optionally merges an ``attention_backend`` shorthand into the config's
|
||
``default`` field. This does **not** read environment variables —
|
||
use :func:`build_attention_config` for the full normalisation that
|
||
should happen exactly once in ``OmniDiffusionConfig.__post_init__``.
|
||
"""
|
||
if attention_config is None:
|
||
normalized = AttentionConfig()
|
||
elif isinstance(attention_config, AttentionConfig):
|
||
normalized = copy.deepcopy(attention_config)
|
||
elif isinstance(attention_config, Mapping):
|
||
normalized = AttentionConfig(**dict(attention_config))
|
||
else:
|
||
raise TypeError(
|
||
f"attention_config must be an AttentionConfig, mapping, or None; got {type(attention_config)!r}"
|
||
)
|
||
|
||
if attention_backend is not None:
|
||
if normalized.default is not None:
|
||
raise ValueError(
|
||
"--diffusion-attention-backend is mutually exclusive with --diffusion-attention-config.default.backend."
|
||
)
|
||
if attention_backend.lower() != "auto":
|
||
normalized.default = AttentionSpec(backend=attention_backend)
|
||
|
||
return normalized
|
||
|
||
|
||
def build_attention_config(
|
||
attention_config: AttentionConfig | Mapping[str, Any] | None = None,
|
||
) -> AttentionConfig:
|
||
"""Normalize diffusion attention config — the single authoritative entry point.
|
||
|
||
Called exactly once in ``OmniDiffusionConfig.__post_init__``.
|
||
Handles type-conversion **and** env-var fallback
|
||
(``DIFFUSION_ATTENTION_BACKEND``).
|
||
"""
|
||
normalized = parse_attention_config(attention_config)
|
||
|
||
if normalized.default is not None:
|
||
return normalized
|
||
|
||
env_attention_backend = os.environ.get("DIFFUSION_ATTENTION_BACKEND")
|
||
if env_attention_backend is None:
|
||
return normalized
|
||
|
||
if env_attention_backend.lower() == "auto":
|
||
return normalized
|
||
|
||
normalized.default = AttentionSpec(backend=env_attention_backend)
|
||
logger.info(
|
||
"Parsed attention config from DIFFUSION_ATTENTION_BACKEND '%s': default=%s, per_role=%s",
|
||
env_attention_backend,
|
||
normalized.default,
|
||
{k: v.backend for k, v in normalized.per_role.items()},
|
||
)
|
||
return normalized
|
||
|
||
|
||
@dataclass
|
||
class OmniACK:
|
||
"""
|
||
Handshake payload from Workers to Orchestrator.
|
||
"""
|
||
|
||
task_id: str
|
||
status: str
|
||
stage_id: int | None = None
|
||
rank: int | None = None
|
||
freed_bytes: int = 0
|
||
metadata: dict[str, Any] = field(default_factory=dict)
|
||
"""
|
||
Additional telemetry such as:
|
||
- max_contiguous_block: for fragmentation analysis.
|
||
- cuda_graph_recalled: boolean if graphs were successfully destroyed/rebuilt.
|
||
- latency_ms: time taken for the D2H/H2D transfer.
|
||
"""
|
||
error_msg: str | None = None
|
||
|
||
|
||
@dataclass
|
||
class OmniSleepTask:
|
||
"""Structured sleep instruction."""
|
||
|
||
task_id: str
|
||
level: int = 2
|
||
metadata: dict[str, Any] = field(default_factory=dict)
|
||
|
||
|
||
@dataclass
|
||
class OmniWakeTask:
|
||
"""Structured wake-up instruction."""
|
||
|
||
task_id: str
|
||
tags: list[str] | None = None
|
||
|
||
|
||
class CuMemTag(str, Enum):
|
||
"""Tags representing specific CuMem allocations for sleep/wake state tracking."""
|
||
|
||
WEIGHTS = "weights"
|
||
KV_CACHE = "kv_cache"
|
||
|
||
|
||
# Special message broadcast via scheduler queues to signal worker shutdown.
|
||
SHUTDOWN_MESSAGE = {"type": "shutdown"}
|