Configuration module for vLLM-Omni.
Modules:
| Name | Description |
|---|---|
composable_parallel | Composable parallel strategies for vLLM-Omni. |
config_factory | Config factories for vllm-omni, e.g., StageConfigFactory. |
endpoint_policy | Endpoint restriction policy for omni pipelines. |
lora | |
model | |
omni_config | Structured vLLM-Omni configuration classes. |
pipeline_registry | Pipeline registry and factory for vllm-omni. |
server_settings | |
stage_config | Stage configuration system for vLLM-Omni. |
yaml_util | Centralized OmegaConf wrapper for vLLM-Omni. |
module-attribute ¶
module-attribute ¶
StageConfigType: TypeAlias = (
VllmOmniARStageConfig
| VllmOmniGenerationStageConfig
| VllmOmniDiffusionStageConfig
)
Common structured config contract shared by all Omni stage realizations.
class-attribute instance-attribute ¶
cache_config: OmniStageCacheConfig = field(
default_factory=OmniStageCacheConfig
)
class-attribute instance-attribute ¶
connector_config: OmniStageConnectorConfig = field(
default_factory=OmniStageConnectorConfig
)
class-attribute instance-attribute ¶
load_config: OmniStageLoadConfig = field(
default_factory=OmniStageLoadConfig
)
class-attribute instance-attribute ¶
model_config: OmniStageModelConfig = field(
default_factory=OmniStageModelConfig
)
class-attribute instance-attribute ¶
parallel_config: OmniStageParallelConfig = field(
default_factory=OmniStageParallelConfig
)
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
runtime_config: OmniStageRuntimeConfig = field(
default_factory=OmniStageRuntimeConfig
)
class-attribute instance-attribute ¶
scheduler_config: OmniStageSchedulerConfig = field(
default_factory=OmniStageSchedulerConfig
)
dataclass ¶
Loaded from deploy/
Top-level fields (trust_remote_code, distributed_executor_backend, dtype, quantization, enable_prefix_caching, enable_chunked_prefill, data_parallel_size, pipeline_parallel_size) are pipeline-wide: they apply uniformly to every stage. Fields that legitimately vary per stage live in the individual StageDeployConfig entries under stages:.
class-attribute instance-attribute ¶
distributed_executor_backend: str | None = None
class-attribute instance-attribute ¶
enable_chunked_prefill: bool | None = None
class-attribute instance-attribute ¶
enable_prefix_caching: bool | None = None
class-attribute instance-attribute ¶
pipeline_parallel_size: int | None = None
class-attribute instance-attribute ¶
stages: list[StageDeployConfig] = field(
default_factory=list
)
Bases: ModelConfig
Configuration for Omni models, extending the base ModelConfig.
This configuration class extends the base vLLM ModelConfig with omni-specific fields for multi-stage pipeline processing.
Attributes: hf_config: The model's HF Transformers config (default: None) hf_text_config: The sub text_config of the model's hf_config (default: None) stage_id: Identifier for the stage in a multi-stage pipeline (default: 0) async_chunk: If set to True, perform async chunk model_stage: Stage type identifier, e.g., "thinker" or "talker" (default: "thinker") model_arch: Model architecture name (default: "Qwen2_5OmniForConditionalGeneration") worker_type: Model Type, e.g., "ar" or "generation" engine_output_type: Optional output type specification for the engine. Used to route outputs to appropriate processors (e.g., "image", "audio", "latents"). If None, output type is inferred. stage_connector_config: Stage connector configuration dictionary. Contains "name" (connector name), "extra" (extra connector config). task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base). If not specified, will be inferred from model path.
The correct way to initialize this class is via vLLM config, as most of the logic for handling values is in the ModelConfig's post_init.
Example: >>> config = OmniModelConfig.from_vllm_model_config( ... vllm_config, ... stage_id=0, ... model_stage="thinker", ... model_arch="Qwen2_5OmniForConditionalGeneration" ... )
class-attribute instance-attribute ¶
custom_process_next_stage_input_func: str | None = None
class-attribute instance-attribute ¶
stage_connector_config: dict[str, Any] = field(
default_factory=lambda: {
"name": "SharedMemoryConnector",
"extra": {},
}
)
class-attribute instance-attribute ¶
classmethod ¶
Because we init the OmniModelConfig with new to sidestep expensive validation, we need to be careful to ensure fields with default factories are initialized, otherwise we will get an AttributeError when we use it.
To work around this issue, we explicitly add defaults to the omni_kwargs dict provided to ensure all fields are defined correctly.
NOTE: omni_kwargs are mutated in place.
classmethod ¶
Create OmniModelConfig from an existing vLLM ModelConfig and additional Omni specific kwargs.
NOTE: The validation and post_init for ModelConfig is expensive; to avoid calling it a second time, we explicitly retrieve defaults from dataclass attributes for values not passed to omni_kwargs, and use that to initialize a new instance. This is significantly faster than creating the OmniModelConfig directly from the ModelConfig, and saves us from having to pass all kwargs to the OmniModelConfig.
Per-stage engine cache and memory behavior.
This is separate from _DiffusionConfigProjection.cache_config, which configures vLLM-Omni diffusion-specific cache backends such as TeaCache and Cache-DiT.
Bases: OmniStageParallelConfig
Diffusion-stage distributed parallelism behavior.
class-attribute instance-attribute ¶
cfg_parallel_size: int = Field(default=1, ge=1, le=3)
class-attribute instance-attribute ¶
hsdp_replicate_size: int = Field(default=1, ge=1)
Per-stage loading behavior.
Per-stage model behavior.
class-attribute instance-attribute ¶
active_stream_window: int = Field(default=0, ge=0)
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
enable_flashinfer_autotune: bool | None = None
class-attribute instance-attribute ¶
enable_multithread_weight_load: bool = True
class-attribute instance-attribute ¶
num_weight_load_threads: int = Field(default=4, ge=1)
class-attribute instance-attribute ¶
Common per-stage distributed parallelism behavior.
class-attribute instance-attribute ¶
data_parallel_size: int = Field(default=1, ge=1)
class-attribute instance-attribute ¶
pipeline_parallel_size: int = Field(default=1, ge=1)
Per-stage process placement and runtime behavior.
dataclass ¶
Complete pipeline topology for a model (frozen).
class-attribute instance-attribute ¶
endpoint_restrictions: tuple[EndpointRestriction, ...] = ()
class-attribute instance-attribute ¶
dataclass ¶
Per-stage config (legacy path). Used by both new and legacy loaders.
TODO(@lishunyang12): replace with ResolvedStageConfig once all models are migrated.
class-attribute instance-attribute ¶
custom_process_input_func: str | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
Factory that loads pipeline YAML and merges CLI overrides.
Handles both single-stage and multi-stage models.
Pipelines are declared in vllm_omni/config/pipeline_registry.py and where keys in OMNI_PIPELINES map to either a PipelineConfig, or a callable which accepts a Transformers config as an arg & resolves to a PipelineConfig.
NOTE: Models with generic HF model_type collisions (e.g. MiMo Audio reports qwen2) should declare hf_architectures=(...) on their PipelineConfig so the factory can disambiguate via hf_config.architectures.
classmethod ¶
Single-stage diffusion - no YAML needed.
Creates a default diffusion stage configuration for single-stage diffusion models. Returns a legacy OmegaConf-compatible dict for backward compatibility with OmniStage.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
kwargs | dict[str, Any] | Engine arguments from CLI/API. | required |
Returns:
| Type | Description |
|---|---|
list[dict[str, Any]] | List containing a single config dict for the diffusion stage. |
classmethod ¶
create_from_model(
model: str,
*,
trust_remote_code: bool = False,
cli_overrides: dict[str, Any] | None = None,
deploy_config_path: str | None = None,
strategy_specs: Mapping[Any, Any] | None = None,
**deprecated_kwargs: Any,
) -> tuple[list[StageConfig] | None, str | None]
Load pipeline + deploy config, merge with CLI overrides.
Checks OMNI_PIPELINES first, since supported models should be explicitly registered. If a model is not registered in OMNI_PIPELINES, tries to fall back to using the Transformers config & finding pipelines that have overlapping supported architectures.
When strategy_specs is provided (a mapping of role -> list of StrategySpec), the derived parallel sizing is overlaid onto the merged stages (see vllm_omni.config.composable_parallel). This is opt-in: omitting it leaves the existing merge path untouched.
Returns (stages, omni_lb_policy): the merged stages (None when the model is not in the pipeline registry and the caller should fall back to the legacy YAML path) and the strategy-derived, pipeline-wide omni_lb_policy (None when no stage_replica axis set one). The policy is returned rather than threaded through a mutable out-param so the engine can apply it without every intermediate call carrying the dict.
cached classmethod ¶
Fetch the HF config (if it exists) from the model directory.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model | str | Model name or path. | required |
trust_remote_code | bool | Whether to trust remote code for HF config loading. | required |
Returns:
| Type | Description |
|---|---|
PretrainedConfig | None | the model's config or None. |
cached classmethod ¶
get_pipeline_config(
model: str,
trust_remote_code: bool,
deploy_config_path: str | None = None,
) -> PipelineConfig | None
Resolve the PipelineConfig for a model path/name.
classmethod ¶
get_pipeline_endpoint_restrictions(
model: str,
trust_remote_code: bool,
deploy_config_path: str | None,
) -> tuple[EndpointRestriction, ...]
Given a model string, determine the corresponding endpoint restrictions.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model | str | Model name or path. | required |
trust_remote_code | bool | Whether to trust remote code for HF config loading. | required |
deploy_config_path | str | None | Optional path to the deploy config for the pipeline. | required |
Returns:
| Type | Description |
|---|---|
tuple[EndpointRestriction, ...] | A tuple of model specific endpoint restrictions. |
staticmethod ¶
resolve_pipeline_config(
model_type: str,
hf_config: PretrainedConfig | None = None,
) -> PipelineConfig | None
Given a model type, resolve to the pipeline to be used. If the pipeline maps to a callable we resolve based on the HF config.
cached classmethod ¶
Auto-detect model_type from model directory and apply any model specific patches to get the correct model_type str. If we are unable to infer it from the model directory, we fall back to the PipelineConfig.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model | str | Model name or path. | required |
trust_remote_code | bool | Whether to trust remote code for HF config loading. | required |
Returns:
| Type | Description |
|---|---|
str | None | model_type as a string; may be None on failure. |
dataclass ¶
Per-stage deployment knobs.
Only fields whose value legitimately varies across stages of the same pipeline live here (e.g. max_num_seqs on thinker vs talker, devices for GPU placement). Pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level of DeployConfig and propagated to every stage.
class-attribute instance-attribute ¶
auxiliary_text_encoder: str | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
diffusion_attention_backend: str | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
diffusion_kv_cache_dtype: str | None = None
class-attribute instance-attribute ¶
diffusion_kv_cache_skip_layers: str | None = None
class-attribute instance-attribute ¶
diffusion_kv_cache_skip_steps: str | None = None
class-attribute instance-attribute ¶
diffusion_quantization_config: str | None = None
class-attribute instance-attribute ¶
disable_hybrid_kv_cache_manager: bool | None = None
class-attribute instance-attribute ¶
enable_cache_dit_summary: bool | None = None
class-attribute instance-attribute ¶
enable_diffusion_pipeline_profiler: bool | None = None
class-attribute instance-attribute ¶
enable_expert_parallel: bool | None = None
class-attribute instance-attribute ¶
enable_flashinfer_autotune: bool | None = None
class-attribute instance-attribute ¶
enable_layerwise_offload: bool | None = None
class-attribute instance-attribute ¶
enable_multithread_weight_load: bool | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
gpu_memory_utilization: float | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
max_generated_image_size: int | None = None
class-attribute instance-attribute ¶
max_num_batched_tokens: int | None = None
class-attribute instance-attribute ¶
mm_processor_cache_gb: float | None = None
class-attribute instance-attribute ¶
num_weight_load_threads: int | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
sequence_parallel_size: int | None = None
class-attribute instance-attribute ¶
class-attribute instance-attribute ¶
tts_max_instructions_length: int | None = None
dataclass ¶
Fixed topology for one stage (frozen, not user-configurable).
class-attribute instance-attribute ¶
async_chunk_process_next_stage_input_func: str | None = None
class-attribute instance-attribute ¶
custom_process_input_func: str | None = None
class-attribute instance-attribute ¶
custom_process_next_stage_input_func: str | None = None
class-attribute instance-attribute ¶
execution_type: StageExecutionType = (
StageExecutionType.LLM_AR
)
class-attribute instance-attribute ¶
requires_multimodal_data: bool = False
class-attribute instance-attribute ¶
Top-level structured Omni config built once from registry inputs.
class-attribute instance-attribute ¶
orchestrator_config: VllmOmniOrchestratorConfig = field(
default_factory=VllmOmniOrchestratorConfig
)
Bases: BaseVllmOmniStageConfig
Structured config for diffusion stages.
class-attribute instance-attribute ¶
diffusion_config: _DiffusionConfigProjection = field(
default_factory=_DiffusionConfigProjection
)
class-attribute instance-attribute ¶
parallel_config: OmniStageDiffusionParallelConfig = field(
default_factory=OmniStageDiffusionParallelConfig
)
Configuration consumed by the orchestrator process only.
class-attribute instance-attribute ¶
omni_dp_size_local: int = Field(default=1, ge=1)
class-attribute instance-attribute ¶
omni_heartbeat_timeout: float = Field(default=30.0, gt=0.0)
class-attribute instance-attribute ¶
shm_threshold_bytes: int = Field(default=65536, ge=0)
create_config(data: Any) -> DictConfig
Wrap a dict (or list) into a DictConfig.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
data | Any | Dict, list, or other structure to wrap. | required |
Returns:
| Type | Description |
|---|---|
DictConfig | OmegaConf DictConfig / ListConfig. |
load_deploy_config(path: str | Path) -> DeployConfig
Load a deploy YAML (with optional base_config inheritance).
merge_pipeline_deploy(
pipeline: PipelineConfig,
deploy: DeployConfig,
cli_overrides: dict[str, Any] | None = None,
) -> list[StageConfig]
Merge pipeline + deploy + platform overrides → list[StageConfig].
register_pipeline(
pipeline: PipelineConfig | PipelineResolverFunc,
model_type: str | None = None,
)
Register an out of tree pipeline or PipelineResolverFunc to a model_type key. If a PipelineConfig is provided, model_type is optional, and pipeline.model_type will be used by default. If a callable is provided, model_type must be provided, since resolvers can return multiple different PipelineConfigs depending on the consumed config.