module-attribute ¶
module-attribute ¶
dataclass ¶
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.
Bases: 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 requestOUTPUT_READY— background D2H/SHM packing finished, final output is available viaasync_output_id
dataclass ¶
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
dataclass ¶
Specifies a backend and its typed backend-specific config for one attention role.
class-attribute instance-attribute ¶
dataclass ¶
dataclass ¶
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).
class-attribute instance-attribute ¶
dataclass ¶
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
class-attribute instance-attribute ¶
force_refresh_step_hint: int | None = None
class-attribute instance-attribute ¶
force_refresh_step_policy: str = 'once'
class-attribute instance-attribute ¶
max_continuous_cached_steps: int = 3
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
classmethod ¶
from_dict(data: dict[str, Any]) -> DiffusionCacheConfig
Create DiffusionCacheConfig from a dictionary.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
data | dict[str, Any] | Dictionary containing cache configuration parameters | required |
Returns:
| Type | Description |
|---|---|
DiffusionCacheConfig | DiffusionCacheConfig instance with parameters set from dict |
dataclass ¶
Final output (after pipeline completion)
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
trajectory_decoded: list[Image] | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
dataclass ¶
Configuration for diffusion model distributed execution.
class-attribute instance-attribute ¶
allgather_degree: int = 1
Number of GPUs used for AllGather-KV sequence parallelism (causal=False only).
class-attribute instance-attribute ¶
cfg_parallel_size: int = 1
Number of ranks used to execute guidance passes in parallel.
class-attribute instance-attribute ¶
data_parallel_size: int | None = None
Number of data parallel groups. Inferred from WORLD when omitted.
class-attribute instance-attribute ¶
enable_expert_parallel: bool = False
Enable expert parallelism for MoE layers (TP is still used for non-MoE layers).
class-attribute instance-attribute ¶
hsdp_replicate_size: int = 1
Number of replica groups for HSDP. Each replica holds a full sharded copy.
class-attribute instance-attribute ¶
hsdp_shard_size: int = -1
Number of GPUs to shard weights across within each replica group. -1 means auto-calculate.
class-attribute instance-attribute ¶
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.
class-attribute instance-attribute ¶
pipeline_parallel_size: int = 1
Number of pipeline parallel stages.
class-attribute instance-attribute ¶
ring_degree: int = 1
Number of GPUs used for ring sequence parallelism.
class-attribute instance-attribute ¶
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.
class-attribute instance-attribute ¶
tensor_parallel_size: int = 1
Number of tensor parallel groups.
class-attribute instance-attribute ¶
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.
class-attribute instance-attribute ¶
ulysses_a2a_permute: bool = False
Use fused permute-free all-to-all for eligible strict Ulysses exchanges.
class-attribute instance-attribute ¶
ulysses_degree: int = 1
Number of GPUs used for ulysses sequence parallelism.
class-attribute instance-attribute ¶
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.
class-attribute instance-attribute ¶
use_hsdp: bool = False
Enable Hybrid Sharded Data Parallel (HSDP) for model weight sharding.
class-attribute instance-attribute ¶
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.
class-attribute instance-attribute ¶
vae_patch_parallel_size: int = 1
Number of ranks used for VAE patch/tile parallelism (decode/encode).
classmethod ¶
from_dict(data: dict[str, Any]) -> DiffusionParallelConfig
Create DiffusionParallelConfig from a dictionary.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
data | dict[str, Any] | Dictionary containing parallel configuration parameters | required |
Returns:
| Type | Description |
|---|---|
DiffusionParallelConfig | DiffusionParallelConfig instance with parameters set from dict |
classmethod ¶
from_stage_overrides(
kwargs: Mapping[str, Any],
) -> DiffusionParallelConfig
Resolve nested or flat stage inputs through the config's own defaults.
Bases: RuntimeError
Raised when a diffusion request ends via user-visible abort.
dataclass ¶
Handshake payload from Workers to Orchestrator.
class-attribute instance-attribute ¶
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.
dataclass ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
diffusers_pipeline_cls: type[DiffusionPipeline] | None = (
None
)
class-attribute instance-attribute ¶
diffusion_attention_config: AttentionConfig = field(
default_factory=lambda: AttentionConfig()
)
class-attribute instance-attribute ¶
diffusion_compile_dynamic: bool = True
class-attribute instance-attribute ¶
diffusion_compile_granularity: str = 'regional'
class-attribute instance-attribute ¶
diffusion_kv_cache_dtype: str | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
diffusion_kv_cache_skip_layers: str | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
diffusion_kv_cache_skip_steps: str | None = None
class-attribute instance-attribute ¶
diffusion_kv_max_rows_per_request: int | None = None
class-attribute instance-attribute ¶
diffusion_kv_mode: DiffusionKVCacheMode = (
DiffusionKVCacheMode.DENSE_LEGACY
)
class-attribute instance-attribute ¶
diffusion_model_runner_cls: str | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
distributed_executor_backend: str | None = None
class-attribute instance-attribute ¶
dlo_host_registration_limit_gib: float = 0.0
class-attribute instance-attribute ¶
enable_broadcast_weight_load: bool = False
class-attribute instance-attribute ¶
enable_cache_dit_summary: bool = False
class-attribute instance-attribute ¶
enable_diffusion_pipeline_profiler: bool = False
class-attribute instance-attribute ¶
enable_distributed_layerwise_offload: bool = False
class-attribute instance-attribute ¶
enable_layerwise_offload: bool = False
class-attribute instance-attribute ¶
enable_multithread_weight_load: bool = True
class-attribute instance-attribute ¶
enable_prompt_embed_cache: bool = False
class-attribute instance-attribute ¶
enable_session_state_manager: bool = False
class-attribute instance-attribute ¶
enable_stage_verification: bool = True
class-attribute instance-attribute ¶
host_weight_runtime_mode: str = 'disabled'
class-attribute instance-attribute ¶
host_weight_runtime_root: str | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
mask_strategy_file_path: str | None = None
class-attribute instance-attribute ¶
max_multimodal_image_inputs: int | None = None
class-attribute instance-attribute ¶
max_num_batched_tokens: int | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
model_loaded: dict[str, bool] = field(
default_factory=lambda: {
"transformer": True,
"vae": True,
"text_encoder": True,
"vae_encoder": True,
}
)
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
override_transformer_cls_name: str | None = None
class-attribute instance-attribute ¶
parallel_config: DiffusionParallelConfig = field(
default_factory=DiffusionParallelConfig
)
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
quantization_config_is_auto_detected: bool = False
class-attribute instance-attribute ¶
request_batch_max_wait_ms: float = 0.0
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
supports_mixed_reference_inputs: bool = False
class-attribute instance-attribute ¶
supports_multimodal_inputs: bool = False
class-attribute instance-attribute ¶
tf_model_config: TransformerConfig = field(
default_factory=TransformerConfig
)
class-attribute instance-attribute ¶
video_output_transport: VideoOutputTransportConfig = field(
default_factory=VideoOutputTransportConfig
)
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.
classmethod ¶
set_tf_model_config(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.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
tf_config | TransformerConfig | Transformer configuration, typically built via | required |
Find an available port with retry logic.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
port | int | Initial port to check | required |
port_inc | int | Port increment for each attempt | 42 |
max_attempts | int | Maximum number of attempts to find an available port | 100 |
Returns:
| Type | Description |
|---|---|
int | An available port number |
Raises:
| Type | Description |
|---|---|
RuntimeError | If no available port is found after max_attempts |
dataclass ¶
dataclass ¶
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.
dataclass ¶
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.
dataclass ¶
Container for raw transformer configuration dictionaries.
dataclass ¶
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).
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.
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__.
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".
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.
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.
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.
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.