[Refactor] Migrate dynin_omni to pipeline registry, drop legacy stage… (#4078)

Signed-off-by: dph97 <nenbaying@163.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

examples/offline_inference/dynin_omni/README.md

@@ -42,7 +42,7 @@ pip install \
 ## 3. Hardware and VRAM Requirements
 
 This example uses a 3-stage pipeline on one GPU by default
-([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
+([`dynin_omni.yaml`](../../../vllm_omni/deploy/dynin_omni.yaml)):
 
 - Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
 - Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`

examples/offline_inference/dynin_omni/end2end.py

@@ -878,8 +878,8 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
     parser.add_argument(
         "--stage-config-path",
         type=str,
-        default=str(repo_root / "vllm_omni/model_executor/stage_configs/dynin_omni.yaml"),
-        help="Path to stage config yaml.",
+        default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
+        help="Path to deploy config yaml.",
     )
     parser.add_argument(
         "--dynin-config-path",

examples/online_serving/dynin_omni/README.md

@@ -22,13 +22,14 @@ export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
 export HF_MODULES_CACHE=<transformers_modules_path>
 ```
 
-Run from repository root:
+Run from repository root. The 3-stage pipeline is resolved automatically from
+the model type (`vllm_omni/deploy/dynin_omni.yaml`), so no config path is
+required. Pass `--deploy-config` only to override the defaults:
 
 ```bash
 vllm-omni serve snu-aidas/Dynin-Omni \
   --omni \
-  --port 8091 \
-  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
+  --port 8091
 ```
 
 If `vllm-omni` is not in PATH, run:
@@ -36,8 +37,7 @@ If `vllm-omni` is not in PATH, run:
 ```bash
 PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
   --omni \
-  --port 8091 \
-  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
+  --port 8091
 ```
 
 Wait until the server logs show both `All stages initialized successfully` and

vllm_omni/config/pipeline_registry.py

@@ -149,4 +149,8 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.higgs_audio_v2.pipeline",
         "HIGGS_AUDIO_V2_PIPELINE",
     ),
+    "dynin_omni": (
+        "vllm_omni.model_executor.models.dynin_omni.pipeline",
+        "DYNIN_OMNI_PIPELINE",
+    ),
 }

vllm_omni/deploy/dynin_omni.yaml

@@ -0,0 +1,38 @@
+# Dynin-Omni 3-stage deploy: token2text → token2image → token2audio.
+#
+# Topology (stages, processor functions, model_arch) is frozen in
+# vllm_omni/model_executor/models/dynin_omni/pipeline.py. This file carries
+# only the deploy-time knobs (devices, GPU memory, batched tokens).
+#
+# async_chunk=false: dynin hands the full token payload across stages through
+# the worker-connector data plane (``*_full_payload`` producers), not chunked
+# streaming. Pipeline-wide ``trust_remote_code`` is required by the custom
+# Dynin HF modeling code.
+async_chunk: false
+trust_remote_code: true
+enable_prefix_caching: false
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.5
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false
+
+  - stage_id: 1
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.1
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false
+
+  - stage_id: 2
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.1
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false

vllm_omni/deploy/dynin_omni_ci.yaml

@@ -1,81 +1,35 @@
-# stage config for running dynin_omni with a 3-stage architecture.
-# this config is intended for e2e smoke tests.
+# Dynin-Omni deploy for e2e smoke tests (3-stage, one GPU).
+#
+# Topology is frozen in models/dynin_omni/pipeline.py; ``pipeline: dynin_omni``
+# pins it explicitly so this CI YAML resolves regardless of model_type
+# auto-detection. Differs from the default deploy only in reduced GPU memory
+# and a smaller max_num_batched_tokens so the pipeline fits L4-class CI GPUs.
+pipeline: dynin_omni
+async_chunk: false
+trust_remote_code: true
+enable_prefix_caching: false
 
-stage_args:
+stages:
   - stage_id: 0
-    stage_type: llm
-    runtime:
-      process: true
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2text
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image_full_payload
-      trust_remote_code: true
-      gpu_memory_utilization: 0.5
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 4096
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.5
+    max_num_batched_tokens: 4096
+    enforce_eager: true
+    async_scheduling: false
 
   - stage_id: 1
-    stage_type: llm
-    runtime:
-      process: true
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2image
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio_full_payload
-      trust_remote_code: true
-      gpu_memory_utilization: 0.2
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 4096
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
-    final_output: true
-    final_output_type: image
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.2
+    max_num_batched_tokens: 4096
+    enforce_eager: true
+    async_scheduling: false
 
   - stage_id: 2
-    stage_type: llm
-    runtime:
-      process: true
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2audio
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: true
-      gpu_memory_utilization: 0.2
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 4096
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
-    final_output: true
-    final_output_type: audio
-
-runtime:
-  enabled: true
-  edges:
-    - from: 0
-      to: 1
-    - from: 1
-      to: 2
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.2
+    max_num_batched_tokens: 4096
+    enforce_eager: true
+    async_scheduling: false

vllm_omni/deploy/dynin_omni_multiconnector.yaml

@@ -0,0 +1,65 @@
+# Dynin-Omni 3-stage deploy with cross-node connectors (Mooncake demo).
+#
+# Topology is frozen in models/dynin_omni/pipeline.py (``pipeline: dynin_omni``).
+# This variant wires a MooncakeConnector for each stage hand-off in place of
+# the default SharedMemoryConnector. Update the ``extra`` host/master/metadata
+# fields to match your Mooncake deployment before use.
+pipeline: dynin_omni
+async_chunk: false
+trust_remote_code: false
+enable_prefix_caching: false
+
+connectors:
+  # Mooncake connector for cross-node/intra-node communication.
+  mooncake_connector:
+    name: MooncakeConnector
+    extra:
+      host: "127.0.0.1"
+      metadata_server: "http://10.90.67.86:8080/metadata"
+      master: "10.90.67.86:50051"
+      segment: 512000000    # 512MB
+      localbuf: 64000000    # 64MB
+      proto: "tcp"
+
+  # Yuanrong connector for cross-node/intra-node communication.
+  yuanrong_connector:
+    name: YuanrongConnector
+    extra:
+      host: "127.0.0.1"
+      port: "35000"
+
+  # Alternative SHM connector with an explicit threshold.
+  shared_memory_connector:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536    # 64KB threshold
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 1
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false
+    output_connectors:
+      to_stage_1: mooncake_connector
+
+  - stage_id: 1
+    devices: "0"
+    max_num_seqs: 1
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false
+    input_connectors:
+      from_stage_0: mooncake_connector
+    output_connectors:
+      to_stage_2: mooncake_connector
+
+  - stage_id: 2
+    devices: "0"
+    max_num_seqs: 1
+    max_num_batched_tokens: 32768
+    enforce_eager: true
+    async_scheduling: false
+    input_connectors:
+      from_stage_1: mooncake_connector

vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py

@@ -144,7 +144,6 @@ _DYNIN_CONFIG_CANDIDATE_RELPATHS = (
     "configs/dynin_omni.yaml",
     "models/configs/dynin_omni.yaml",
     "vllm_omni/model_executor/models/dynin_omni/configs/dynin_omni.yaml",
-    "vllm_omni/model_executor/stage_configs/dynin_omni.yaml",
     "dynin_omni.yaml",
 )
 
@@ -482,7 +481,6 @@ def _resolve_config_path(vllm_config: VllmConfig, runtime_info: dict[str, Any])
     for bundled in (
         module_root / "configs" / "dynin_omni.yaml",
         module_root / "models" / "configs" / "dynin_omni.yaml",
-        module_root.parent / "stage_configs" / "dynin_omni.yaml",
     ):
         if bundled.exists():
             return str(bundled)

vllm_omni/model_executor/models/dynin_omni/pipeline.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Dynin-Omni pipeline topology (frozen).
+
+Stage 0: token2text  — multimodal understanding / text generation (comprehension)
+Stage 1: token2image — stage-0 tokens → image latents
+Stage 2: token2audio — stage-1 tokens → audio latents
+
+All three stages run on the generation worker (``LLM_GENERATION``); the
+inter-stage hand-off uses the worker-connector full-payload data plane
+(``*_full_payload`` producers + ``token2text_to_token2image`` /
+``token2image_to_token2audio`` consumers). Deploy knobs (devices, GPU memory,
+batched tokens, connectors) live in ``vllm_omni/deploy/dynin_omni*.yaml``.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.dynin_omni"
+
+DYNIN_OMNI_PIPELINE = PipelineConfig(
+    model_type="dynin_omni",
+    model_arch="DyninOmniForConditionalGeneration",
+    # Arch-fallback safety net: route by hf_config.architectures when the
+    # auto-detected model_type does not match the registry key exactly.
+    hf_architectures=("DyninOmniForConditionalGeneration",),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="token2text",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            engine_output_type="latent",
+            custom_process_next_stage_input_func=f"{_PROC}.token2text_to_token2image_full_payload",
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="token2image",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="image",
+            engine_output_type="latent",
+            custom_process_input_func=f"{_PROC}.token2text_to_token2image",
+            custom_process_next_stage_input_func=f"{_PROC}.token2image_to_token2audio_full_payload",
+        ),
+        StagePipelineConfig(
+            stage_id=2,
+            model_stage="token2audio",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(1,),
+            final_output=True,
+            final_output_type="audio",
+            engine_output_type="latent",
+            custom_process_input_func=f"{_PROC}.token2image_to_token2audio",
+        ),
+    ),
+)

vllm_omni/model_executor/stage_configs/dynin_omni.yaml

@@ -1,76 +0,0 @@
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2text
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image_full_payload
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: true
-      gpu_memory_utilization: 0.5
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-
-  - stage_id: 1
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2image
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: true
-      gpu_memory_utilization: 0.1
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
-    final_output: true
-    final_output_type: image
-
-  - stage_id: 2
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2audio
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: true
-      gpu_memory_utilization: 0.1
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
-    final_output: true
-    final_output_type: audio
-
-# Top-level runtime config (concise): default windows and stage edges
-runtime:
-  enabled: true
-
-  edges:
-    - from: 0
-      to: 1
-    - from: 1
-      to: 2

vllm_omni/model_executor/stage_configs/dynin_omni_multiconnector.yaml

@@ -1,110 +0,0 @@
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2text
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image_full_payload
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: false
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    output_connectors:
-      to_stage_1: mooncake_connector
-    final_output: true
-    final_output_type: text
-
-  - stage_id: 1
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2image
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: false
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    engine_input_source: [0]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
-    final_output: true
-    final_output_type: image
-    input_connectors:
-      from_stage_0: mooncake_connector
-    output_connectors:
-      to_stage_2: mooncake_connector
-
-  - stage_id: 2
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: token2audio
-      model_arch: DyninOmniForConditionalGeneration
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      engine_output_type: latent
-      trust_remote_code: false
-      enforce_eager: true
-      enable_prefix_caching: false
-      async_scheduling: false
-      max_num_batched_tokens: 32768
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
-    final_output: true
-    final_output_type: audio
-    input_connectors:
-      from_stage_1: mooncake_connector
-
-# Top-level runtime config (concise): default windows and stage edges
-runtime:
-  enabled: true
-  ####
-  # same as Qwen2.5_omni version
-  # Distributed connectors configuration (optional)
-  # More connectors will be supported in the future.
-  connectors:
-    # Mooncake connector for cross-node/intra-node communication
-    mooncake_connector:
-      name: MooncakeConnector
-      extra:
-        host: "127.0.0.1"
-        metadata_server: "http://10.90.67.86:8080/metadata"
-        master: "10.90.67.86:50051"
-        segment: 512000000    # 512MB
-        localbuf: 64000000     # 64MB
-        proto: "tcp"
-
-    # Yuanrong connector for cross-node/intra-node communication
-    yuanrong_connector:
-      name: YuanrongConnector
-      extra:
-        host: "127.0.0.1"
-        port: "35000"
-
-    # SharedMemory connector for intra-node communication
-    # Alternative SHM connector with different threshold
-    shared_memory_connector:
-      name: SharedMemoryConnector
-      extra:
-        shm_threshold_bytes: 65536 # 64KB threshold
-    ####
-
-  edges:
-    - from: 0
-      to: 1
-    - from: 1
-      to: 2