[FEAT] support multi-stage deployment (#2396)

Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: Zheng Wengang <zwg0606@gmail.com>
Signed-off-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Signed-off-by: yinpe <11810305@mail.sustech.edu.cn>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: yinpe <11810305@mail.sustech.edu.cn>
Co-authored-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Chenguang Zheng <645327136@qq.com>

.buildkite/test-nightly.yml

@@ -215,6 +215,44 @@ steps:
                       type: DirectoryOrCreate
 
 
+      - label: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100"
+        timeout_in_minutes: 45
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "core_model" --run-level "core_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 4
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
   - group: ":card_index_dividers: TTS Model Test"
     key: nightly-tts-test-group
     depends_on: upload-nightly-pipeline
@@ -297,7 +335,8 @@ steps:
       - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100"
         timeout_in_minutes: 120
         commands:
-          - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy
+          # Keep Bagel expansion and multi-replica tests in their dedicated H100 jobs below.
+          - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_bagel_multi_replicas.py
         agents:
           queue: "mithril-h100-pool"
         plugins:
@@ -370,6 +409,44 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Diffusion X2I(&A&T) · BAGEL Multi-Replica with H100"
+        timeout_in_minutes: 120
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_bagel_multi_replicas.py -m "full_model and diffusion and H100" --run-level "full_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 4
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
       - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
         timeout_in_minutes: 60
         commands:

benchmarks/diffusion/README.md

@@ -32,6 +32,8 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
 
 - The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
 - If you run the server on another host or port, pass `--base-url` accordingly.
+- To compare diffusion multi-replica serving, restart the server with different
+  `runtime.num_replicas` values and keep the benchmark parameters unchanged.
 
 ## 2. Supported Datasets
 
@@ -93,6 +95,35 @@ You can point to your own trace using `--dataset-path`.
 
 ## 3. Benchmark Parameters
 
+### Multi-replica sweep
+
+For replica scaling, use the same model, prompt count, image/video size, and
+step count for every run. Change only the serving stage config:
+
+```yaml
+runtime:
+  devices: "0,1,2,3"
+  num_replicas: 4
+```
+
+Then run a fixed benchmark and write one JSON file per replica count:
+
+```bash
+python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+  --base-url http://localhost:8099 \
+  --model Qwen/Qwen-Image \
+  --task t2i \
+  --dataset random \
+  --num-prompts 32 \
+  --max-concurrency 8 \
+  --request-rate inf \
+  --num-inference-steps 10 \
+  --width 512 --height 512 \
+  --output-file diffusion_replica_4.json
+```
+
+Repeat for `num_replicas=1`, `2`, `3`, and `4`, using matching `runtime.devices`.
+
 ### Basic flags
 
 - `--base-url`: Server address (the script calls `.../v1/chat/completions`).

docs/user_guide/diffusion/multi_replica_serving.md

@@ -0,0 +1,80 @@
+# Diffusion Multi-Replica Serving
+
+Diffusion stages can be replicated in online serving by setting `runtime.num_replicas` on a `stage_type: diffusion` stage. Each replica owns an independent diffusion engine process and receives requests through round-robin routing with request affinity.
+
+## When to Use
+
+Use multi-replica serving when one diffusion process does not saturate the available GPUs, or when you want higher request throughput for many independent image/video generation requests. It improves concurrency throughput; it does not make a single request faster.
+
+## Configuration
+
+A replica consumes the devices required by that diffusion stage. For single-device diffusion pipelines, set `runtime.devices` to one device per replica:
+
+```yaml
+stages:
+  - stage_id: 0
+    stage_type: diffusion
+    runtime:
+      devices: "0,1,2,3"
+      num_replicas: 4
+    engine_args:
+      model: Qwen/Qwen-Image
+      parallel_config:
+        world_size: 1
+```
+
+For multi-device diffusion pipelines, `parallel_config.world_size` is treated as the number of devices per replica. For example, four 2-GPU replicas need eight devices:
+
+```yaml
+stages:
+  - stage_id: 0
+    stage_type: diffusion
+    runtime:
+      devices: "0,1,2,3,4,5,6,7"
+      num_replicas: 4
+    engine_args:
+      model: Wan-AI/Wan2.2-T2V-A14B-Diffusers
+      parallel_config:
+        world_size: 2
+```
+
+The runtime splits the device list into contiguous chunks per replica. With the second example, replica 0 uses `0,1`, replica 1 uses `2,3`, replica 2 uses `4,5`, and replica 3 uses `6,7`.
+
+## Serving
+
+Start the server with the stage config:
+
+```bash
+vllm serve Qwen/Qwen-Image \
+  --omni \
+  --port 8099 \
+  --stage-configs-path /path/to/diffusion_multi_replica.yaml
+```
+
+Then send normal OpenAI-compatible diffusion chat requests. The orchestrator routes new requests across replicas and keeps updates/abort calls on the same replica for each request.
+
+## Benchmarking
+
+Benchmark each replica count by changing `runtime.num_replicas` and `runtime.devices`, restarting the server, then running:
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+  --base-url http://localhost:8099 \
+  --model Qwen/Qwen-Image \
+  --task t2i \
+  --dataset random \
+  --num-prompts 8 \
+  --request-rate inf \
+  --num-inference-steps 2 \
+  --width 256 --height 256 \
+  --output-file diffusion_replica_4.json
+```
+
+For a quick scaling sweep, test `num_replicas` values `1`, `2`, `3`, and `4` with the same prompt count and generation parameters. Compare throughput and end-to-end latency from the JSON outputs.
+
+## Notes and Limitations
+
+- Multi-replica serving is currently supported for diffusion stages. Single-stage mode still rejects non-diffusion replicas.
+- `runtime.devices` must provide exactly `num_replicas * devices_per_replica` devices.
+- Request-level latency may stay similar or slightly increase because each request still runs on one replica.
+- Throughput improves only when enough concurrent requests are issued and each replica has enough GPU memory.

docs/user_guide/diffusion_features.md

@@ -15,7 +15,7 @@ vLLM-Omni supports various advanced features for diffusion models:
 - Acceleration: **cache methods**, **parallelism methods**, **startup optimizations**
 - Memory optimization: **cpu offloading**, **quantization**
 - Extensions: **LoRA inference**, **frame interpolation**
-- Execution modes: **step execution**
+- Execution modes: **step execution**, **multi-replica serving**
 
 ## Supported Features
 
@@ -43,6 +43,7 @@ Parallelism methods distribute computation across GPUs without quality loss (mat
 | **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
 | **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
 | **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
+| **[Multi-Replica Serving](diffusion/multi_replica_serving.md)** | Runs multiple independent diffusion stage replicas and routes requests across them | Online serving throughput with many concurrent requests |
 
 #### Startup Optimization
 

examples/online_serving/qwen3_omni/README.md

@@ -28,6 +28,14 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     --deploy-config /path/to/your_deploy_config.yaml
 ```
 
+For the bundled 3x-GPU multi-replica layout (talker/code2wav scale-out),
+use:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml
+```
+
 ### Launch individual stages (stage-based CLI)
 
 Use the stage-based CLI when you want to run one stage per process.

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -94,3 +95,24 @@ def test_inline_shutdown(client, mock_engine):
 
     assert client._shutting_down
     mock_engine.close.assert_called_once()
+
+
+def test_inline_client_requires_replica_id(mock_engine):
+    metadata = SimpleNamespace(
+        stage_id=0,
+        final_output=True,
+        final_output_type="image",
+        default_sampling_params={},
+        requires_multimodal_data=False,
+        custom_process_input_func=None,
+        engine_input_source=[],
+    )
+    with patch.object(InlineStageDiffusionClient, "_enrich_config"):
+        od_config = MagicMock(spec=OmniDiffusionConfig)
+        with pytest.raises(AttributeError, match="replica_id"):
+            InlineStageDiffusionClient(
+                model="test_model",
+                od_config=od_config,
+                metadata=metadata,
+                batch_size=1,
+            )

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -621,6 +621,29 @@ class TestStageDiffusionClientErrorPropagation:
         assert client._shutting_down is True
         assert client._engine_dead is True
 
+    def test_initialize_client_requires_replica_id(self):
+        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+        client = object.__new__(StageDiffusionClient)
+        metadata = SimpleNamespace(
+            stage_id=0,
+            final_output=True,
+            final_output_type="image",
+            default_sampling_params=None,
+            requires_multimodal_data=False,
+            custom_process_input_func=None,
+            engine_input_source=[],
+        )
+
+        with pytest.raises(AttributeError, match="replica_id"):
+            client._initialize_client(
+                metadata,
+                "tcp://req",
+                "tcp://resp",
+                proc=None,
+                batch_size=1,
+            )
+
 
 # ───────── monitor thread & death sentinel integration tests ─────────
 

tests/e2e/online_serving/test_bagel_multi_replicas.py

@@ -0,0 +1,55 @@
+"""Nightly-only e2e coverage for BAGEL diffusion multi-replica serving.
+
+This test needs 4 H100 GPUs; keep it out of generic test_bagel_* jobs.
+"""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
+MULTI_REPLICA_DEPLOY = get_deploy_config_path("ci/bagel_multi_replicas_4gpu.yaml")
+ROUTE_STRESS_REQUESTS = 6
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+test_params = [
+    OmniServerParams(
+        model=MODEL,
+        stage_config_path=MULTI_REPLICA_DEPLOY,
+        server_args=["--disable-log-stats"],
+        use_stage_cli=True,
+    )
+]
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_multi_stage_diffusion_uses_multi_replica_dit_stage(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": dummy_messages_from_mix_data(
+            content_text="<|im_start|>A small red cube on a white table.<|im_end|>"
+        ),
+        "modalities": ["image"],
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "guidance_scale": 0.0,
+            "seed": 42,
+        },
+    }
+
+    responses = openai_client.send_diffusion_request(request_config, request_num=ROUTE_STRESS_REQUESTS)
+
+    assert len(responses) == ROUTE_STRESS_REQUESTS
+    assert all(response.success for response in responses)
+    assert all(response.images is not None and response.images[0].size == (512, 512) for response in responses)

tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py

@@ -0,0 +1,137 @@
+"""
+Core-model CI guard for Qwen3-Omni multi-replica stage-pool routing on 4 GPUs.
+"""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
+from tests.helpers.runtime import OmniResponse, OmniServerParams, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+MULTI_REPLICA_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe_multi_replicas_4gpu.yaml")
+ROUTE_STRESS_REQUESTS = 6
+MIXED_MODAL_REQUESTS = 4
+
+test_params = [
+    OmniServerParams(
+        model=MODEL,
+        stage_config_path=MULTI_REPLICA_DEPLOY,
+        server_args=["--disable-log-stats"],
+    )
+]
+
+
+def _system_prompt() -> dict[str, object]:
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def _text_messages() -> list[dict[str, object]]:
+    return dummy_messages_from_mix_data(
+        system_prompt=_system_prompt(),
+        content_text="What is the capital of China? Answer in one short sentence.",
+    )
+
+
+def _mixed_messages() -> list[dict[str, object]]:
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
+    return dummy_messages_from_mix_data(
+        system_prompt=_system_prompt(),
+        video_data_url=video_data_url,
+        image_data_url=image_data_url,
+        audio_data_url=audio_data_url,
+        content_text="What is recited in the audio? What is in this image? Describe the video briefly.",
+    )
+
+
+def _assert_batch_size(responses: list[OmniResponse], expected: int) -> None:
+    assert len(responses) == expected, f"Expected {expected} responses, got {len(responses)}"
+    assert all(resp.success for resp in responses), "At least one request failed"
+    assert all(resp.e2e_latency is not None and resp.e2e_latency > 0 for resp in responses), "Missing request latency"
+
+
+def _assert_text_outputs(responses: list[OmniResponse]) -> None:
+    assert all(resp.text_content is not None for resp in responses), "Missing text output"
+    assert all(resp.audio_bytes is None for resp in responses), "Text-only request unexpectedly produced audio"
+
+
+def _assert_audio_outputs(responses: list[OmniResponse], *, expect_text: bool) -> None:
+    assert all(resp.audio_bytes is not None and len(resp.audio_bytes) > 128 for resp in responses), (
+        "Missing audio output"
+    )
+    if expect_text:
+        assert all(resp.text_content is not None for resp in responses), "Missing text output"
+    else:
+        assert all(not (resp.text_content or "").strip() for resp in responses), (
+            "Audio-only request unexpectedly produced text"
+        )
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_only_batch_uses_multi_replica_talker(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _text_messages(),
+        "stream": False,
+        "modalities": ["text"],
+    }
+
+    responses = openai_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS)
+    _assert_batch_size(responses, ROUTE_STRESS_REQUESTS)
+    _assert_text_outputs(responses)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_to_audio_stream_batch_uses_multi_replica_vocoder(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _text_messages(),
+        "stream": True,
+        "modalities": ["audio"],
+    }
+
+    responses = openai_client.send_omni_request(request_config, request_num=ROUTE_STRESS_REQUESTS)
+    _assert_batch_size(responses, ROUTE_STRESS_REQUESTS)
+    _assert_audio_outputs(responses, expect_text=False)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_mixed_modal_stream_batch_generates_text_and_audio(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _mixed_messages(),
+        "stream": True,
+    }
+
+    responses = openai_client.send_omni_request(request_config, request_num=MIXED_MODAL_REQUESTS)
+    _assert_batch_size(responses, MIXED_MODAL_REQUESTS)
+    _assert_audio_outputs(responses, expect_text=True)

tests/engine/test_async_omni_engine_input.py

@@ -59,7 +59,7 @@ def test_build_add_request_message_preserves_additional_information(mocker: Mock
     assert request.additional_information is not None
     assert request.additional_information.entries["text"].list_data == ["hello world"]
     assert request.additional_information.entries["speaker"].list_data == ["vivian"]
-    output_processor.add_request.assert_called_once()
+    output_processor.add_request.assert_not_called()
 
 
 def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixture):

tests/engine/test_async_omni_engine_stage_init.py

@@ -1,15 +1,138 @@
 import importlib
 import os
 import threading
+import time
 import types
 
 import pytest
 
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.engine.stage_init_utils import (
+    LogicalStageInitPlan,
+    ReplicaInitPlan,
+    build_stage0_input_processor,
+    compute_replica_layout,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _make_llm_metadata(
+    stage_id: int,
+    *,
+    replica_id: int = 0,
+    final_output: bool = False,
+    final_output_type: str | None = None,
+    is_comprehension: bool = False,
+):
+    return types.SimpleNamespace(
+        stage_id=stage_id,
+        stage_type="llm",
+        runtime_cfg={},
+        prompt_expand_func=None,
+        final_output=final_output,
+        final_output_type=final_output_type,
+        default_sampling_params=types.SimpleNamespace(name=f"sp-{stage_id}-{replica_id}"),
+        custom_process_input_func=None,
+        engine_input_source=[] if stage_id == 0 else [stage_id - 1],
+        engine_output_type="token_ids",
+        replica_id=replica_id,
+        is_comprehension=is_comprehension,
+    )
+
+
+def _make_diffusion_metadata(stage_id: int, *, replica_id: int = 0, final_output_type: str = "image"):
+    return types.SimpleNamespace(
+        stage_id=stage_id,
+        stage_type="diffusion",
+        runtime_cfg={"devices": str(replica_id)},
+        prompt_expand_func=None,
+        final_output=True,
+        final_output_type=final_output_type,
+        default_sampling_params=types.SimpleNamespace(name=f"dsp-{stage_id}-{replica_id}"),
+        custom_process_input_func=None,
+        engine_input_source=[],
+        cfg_kv_collect_func=None,
+        replica_id=replica_id,
+    )
+
+
+def _make_llm_plan(
+    stage_idx: int,
+    *,
+    configured_stage_id: int,
+    vllm_config: object,
+    num_replicas: int = 1,
+    final_output: bool = False,
+    final_output_type: str | None = None,
+    is_comprehension: bool = False,
+):
+    replicas: list[ReplicaInitPlan] = []
+    for replica_id in range(num_replicas):
+        stage_cfg = types.SimpleNamespace(
+            stage_id=configured_stage_id,
+            stage_type="llm",
+            runtime=types.SimpleNamespace(devices=str(replica_id)),
+            engine_args={},
+        )
+        replicas.append(
+            ReplicaInitPlan(
+                replica_id=replica_id,
+                num_replicas=num_replicas,
+                launch_mode="local",
+                stage_cfg=stage_cfg,
+                metadata=_make_llm_metadata(
+                    configured_stage_id,
+                    replica_id=replica_id,
+                    final_output=final_output,
+                    final_output_type=final_output_type,
+                    is_comprehension=is_comprehension and replica_id == 0,
+                ),
+                stage_connector_spec={},
+                omni_kv_connector=(None, None, None),
+                stage_vllm_config=vllm_config,
+                executor_class=object,
+            )
+        )
+    return LogicalStageInitPlan(
+        stage_idx=stage_idx,
+        configured_stage_id=configured_stage_id,
+        replicas=replicas,
+    )
+
+
+def _make_diffusion_plan(
+    stage_idx: int,
+    *,
+    configured_stage_id: int,
+    num_replicas: int = 1,
+):
+    replicas: list[ReplicaInitPlan] = []
+    for replica_id in range(num_replicas):
+        stage_cfg = types.SimpleNamespace(
+            stage_id=configured_stage_id,
+            stage_type="diffusion",
+            runtime=types.SimpleNamespace(devices=str(replica_id)),
+            engine_args={},
+        )
+        replicas.append(
+            ReplicaInitPlan(
+                replica_id=replica_id,
+                num_replicas=num_replicas,
+                launch_mode="local",
+                stage_cfg=stage_cfg,
+                metadata=_make_diffusion_metadata(configured_stage_id, replica_id=replica_id),
+                stage_connector_spec={},
+                omni_kv_connector=(None, None, None),
+            )
+        )
+    return LogicalStageInitPlan(
+        stage_idx=stage_idx,
+        configured_stage_id=configured_stage_id,
+        replicas=replicas,
+    )
+
+
 def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     """Regression test for forward references in make_async_mp_client."""
     import vllm_omni.engine.stage_engine_core_client as client_mod
@@ -21,15 +144,279 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     )
 
 
-def test_initialize_stages_restores_device_visibility_after_diffusion_init(monkeypatch):
-    """Regression test for stage device env leakage across stage init.
+def test_compute_replica_layout_splits_diffusion_devices_by_world_size():
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="diffusion",
+        engine_args={"parallel_config": {"tensor_parallel_size": 2}},
+        runtime={"devices": "0,1,2,3", "num_replicas": 2},
+    )
 
-    Diffusion init mutates process-level CUDA visibility. Ensure AsyncOmniEngine
-    restores the previous value after diffusion stage setup.
-    """
+    replicas_per_stage, replica_devices_map = compute_replica_layout([stage_cfg])
+
+    assert replicas_per_stage == [2]
+    assert replica_devices_map == {0: ["0,1", "2,3"]}
+
+
+def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
+    shared = types.SimpleNamespace(name="shared")
+    extra = types.SimpleNamespace(name="extra")
+
+    cleanup_clients = AsyncOmniEngine._collect_initialized_clients_for_cleanup(
+        stage_pools=[types.SimpleNamespace(clients=[shared, None])],
+        initialized_clients_by_stage={0: [shared], 1: [extra]},
+    )
+
+    assert cleanup_clients == [shared, extra]
+
+
+def test_initialize_stages_rejects_non_diffusion_replicas_in_single_stage_mode():
+    engine = object.__new__(AsyncOmniEngine)
+    engine.single_stage_mode = True
+    engine.stage_configs = [types.SimpleNamespace(stage_id=0, runtime={"num_replicas": 2})]
+
+    with pytest.raises(
+        ValueError,
+        match="single_stage_mode only supports num_replicas > 1 for diffusion stages",
+    ):
+        engine._validate_single_stage_mode_replica_constraints()
+
+
+def test_initialize_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
     from vllm_omni.platforms import current_omni_platform
 
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.num_stages = 1
+    engine.diffusion_batch_size = 1
+    engine.single_stage_mode = False
+    engine._omni_master_server = None
+    engine.stage_configs = []
+
+    plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0]
+
+    env_var = current_omni_platform.device_control_env_var
+    old_env = os.environ.get(env_var)
+    os.environ[env_var] = "0,1"
+
+    def _fake_setup_stage_devices(_stage_id, _runtime_cfg):
+        current_omni_platform.set_device_control_env_var("1")
+
+    monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices)
+    monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: types.SimpleNamespace())
+
+    try:
+        engine._initialize_diffusion_replica(plan, stage_init_timeout=1, stage_launch_lock=threading.Lock())
+        assert os.environ.get(env_var) == "0,1"
+    finally:
+        if old_env is None:
+            os.environ.pop(env_var, None)
+        else:
+            os.environ[env_var] = old_env
+
+
+def test_initialize_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.num_stages = 1
+    engine.diffusion_batch_size = 4
+    engine.single_stage_mode = False
+    engine._omni_master_server = None
+    engine.stage_configs = []
+
+    plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0]
+
+    captured: dict[str, object] = {}
+
+    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
+
+    def _capture_initialize_diffusion_stage(
+        stage_id, _model, _stage_cfg, _metadata, *, stage_init_timeout, batch_size, use_inline
+    ):
+        captured["stage_id"] = stage_id
+        captured["stage_init_timeout"] = stage_init_timeout
+        captured["batch_size"] = batch_size
+        captured["use_inline"] = use_inline
+        return types.SimpleNamespace()
+
+    monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", _capture_initialize_diffusion_stage)
+
+    engine._initialize_diffusion_replica(plan, stage_init_timeout=302, stage_launch_lock=threading.Lock())
+
+    assert captured == {
+        "stage_id": 0,
+        "stage_init_timeout": 302,
+        "batch_size": 4,
+        "use_inline": True,
+    }
+
+
+def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_input_processor(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.config_path = "dummy-config"
+    engine.num_stages = 2
+    engine.async_chunk = False
+    engine.diffusion_batch_size = 1
+    engine.single_stage_mode = False
+    engine._single_stage_id_filter = None
+    engine._omni_master_server = None
+    engine.stage_configs = [types.SimpleNamespace(), types.SimpleNamespace()]
+
+    cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plans = [
+        _make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2, is_comprehension=True),
+        _make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, final_output=True),
+    ]
+
+    stage0_client_r0 = types.SimpleNamespace(
+        stage_type="llm",
+        is_comprehension=True,
+        final_output=False,
+        final_output_type=None,
+        default_sampling_params=types.SimpleNamespace(name="sp0"),
+    )
+    stage0_client_r1 = types.SimpleNamespace(
+        stage_type="llm",
+        is_comprehension=False,
+        final_output=False,
+        final_output_type=None,
+        default_sampling_params=types.SimpleNamespace(name="sp0r1"),
+    )
+    stage1_client_r0 = types.SimpleNamespace(
+        stage_type="llm",
+        is_comprehension=False,
+        final_output=True,
+        final_output_type=None,
+        default_sampling_params=types.SimpleNamespace(name="sp1"),
+    )
+    initialized_clients = {
+        0: [stage0_client_r0, stage0_client_r1],
+        1: [stage1_client_r0],
+    }
+
+    stage0_output_processor = object()
+    stage1_output_processor = object()
+    top_level_input_processor = object()
+
+    monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
+    monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
+    monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2, 1], {}))
+    monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None))
+    monkeypatch.setattr(engine, "_initialize_stage_replicas", lambda *_: initialized_clients)
+    monkeypatch.setattr(
+        engine_mod,
+        "build_llm_stage_output_processor",
+        lambda plan, _cfg: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor,
+    )
+    monkeypatch.setattr(engine_mod, "build_stage0_input_processor", lambda _cfg: top_level_input_processor)
+
+    engine._initialize_stages(stage_init_timeout=1)
+
+    assert len(engine.stage_pools) == 2
+    assert engine.input_processor is top_level_input_processor
+    assert engine.stage_clients == [stage0_client_r0, stage1_client_r0]
+    assert engine.stage_vllm_configs == [cfg0, cfg1]
+    assert engine.output_processors == [stage0_output_processor, stage1_output_processor]
+    assert engine.default_sampling_params_list == [
+        stage0_client_r0.default_sampling_params,
+        stage1_client_r0.default_sampling_params,
+    ]
+    assert engine.stage_metadata == [
+        {"final_output": False, "final_output_type": None, "stage_type": "llm"},
+        {"final_output": True, "final_output_type": None, "stage_type": "llm"},
+    ]
+
+
+def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.model = "dummy-model"
+    engine.async_chunk = False
+    engine.single_stage_mode = False
+    engine._single_stage_id_filter = None
+    engine.stage_configs = [
+        types.SimpleNamespace(stage_id=0, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="0")),
+        types.SimpleNamespace(
+            stage_id=1, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="1,2,3")
+        ),
+    ]
+
+    metadata_by_stage = {
+        0: _make_llm_metadata(0),
+        1: _make_llm_metadata(1),
+    }
+
+    monkeypatch.setattr(
+        engine_mod,
+        "extract_stage_metadata",
+        lambda cfg: types.SimpleNamespace(**metadata_by_stage[cfg.stage_id].__dict__),
+    )
+    monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {})
+    monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
+    monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
+    monkeypatch.setattr(
+        engine_mod,
+        "build_vllm_config",
+        lambda stage_cfg, *_args, **_kwargs: (types.SimpleNamespace(tag=f"cfg-{stage_cfg.stage_id}"), object),
+    )
+
+    stage_plans, prompt_expand_func = engine._build_logical_stage_init_plans(
+        omni_transfer_config=None,
+        replicas_per_stage=[1, 3],
+        replica_devices_map={1: ["1", "2", "3"]},
+    )
+
+    assert prompt_expand_func is None
+    assert [plan.configured_stage_id for plan in stage_plans] == [0, 1]
+    assert [replica.stage_cfg.runtime.devices for replica in stage_plans[1].replicas] == ["1", "2", "3"]
+    assert [replica.replica_id for replica in stage_plans[1].replicas] == [0, 1, 2]
+    assert all(replica.num_replicas == 3 for replica in stage_plans[1].replicas)
+
+
+def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monkeypatch):
+    engine = object.__new__(AsyncOmniEngine)
+
+    cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plans = [
+        _make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2),
+        _make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, num_replicas=2),
+    ]
+
+    clients = {
+        (0, 0): types.SimpleNamespace(name="stage0-replica0"),
+        (0, 1): types.SimpleNamespace(name="stage0-replica1"),
+        (1, 0): types.SimpleNamespace(name="stage1-replica0"),
+        (1, 1): types.SimpleNamespace(name="stage1-replica1"),
+    }
+
+    def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock):
+        time.sleep(0.02 * (3 - plan.metadata.stage_id - plan.replica_id))
+        return clients[(plan.metadata.stage_id, plan.replica_id)]
+
+    monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica)
+
+    initialized_clients = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=123)
+
+    assert initialized_clients == {
+        0: [clients[(0, 0)], clients[(0, 1)]],
+        1: [clients[(1, 0)], clients[(1, 1)]],
+    }
+
+
+def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
     engine = object.__new__(AsyncOmniEngine)
     engine.model = "dummy-model"
     engine.config_path = "dummy-config"
@@ -39,158 +426,111 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke
     engine.single_stage_mode = False
     engine._single_stage_id_filter = None
     engine._omni_master_server = None
-    engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion")]
+    engine.stage_configs = [types.SimpleNamespace()]
 
-    env_var = current_omni_platform.device_control_env_var
-    old_env = os.environ.get(env_var)
-    os.environ[env_var] = "0,1"
-
-    diffusion_client = types.SimpleNamespace(is_comprehension=False)
-
-    metadata = types.SimpleNamespace(
-        stage_id=0,
-        stage_type="diffusion",
-        runtime_cfg={"devices": "1"},
-        prompt_expand_func=None,
-    )
+    cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)]
+    initialized_client = types.SimpleNamespace(shutdown=lambda: None)
 
     monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
     monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
-    monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata)
-    monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {})
-    monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
+    monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {}))
+    monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None))
 
-    def _fake_setup_stage_devices(_stage_id, _runtime_cfg):
-        # Simulate diffusion setup mutating process-global visibility.
-        current_omni_platform.set_device_control_env_var("1")
+    def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock):
+        if plan.replica_id == 0:
+            return initialized_client
+        time.sleep(0.05)
+        raise RuntimeError("replica launch failed")
 
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices)
-    monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
-    monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: diffusion_client)
-    monkeypatch.setattr(
-        engine_mod,
-        "finalize_initialized_stages",
-        lambda stage_clients, _input_processor: (
-            stage_clients,
-            [types.SimpleNamespace()],
-            [{"final_output_type": "image"}],
-        ),
-    )
+    monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica)
 
-    try:
+    captured_cleanup: list[list[object]] = []
+
+    def _capture_shutdown(clients):
+        captured_cleanup.append(list(clients))
+
+    monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown)
+
+    with pytest.raises(RuntimeError, match="replica launch failed"):
         engine._initialize_stages(stage_init_timeout=1)
-        assert os.environ.get(env_var) == "0,1"
-    finally:
-        if old_env is None:
-            os.environ.pop(env_var, None)
-        else:
-            os.environ[env_var] = old_env
+
+    assert captured_cleanup == [[initialized_client]]
 
 
-def test_initialize_stages_passes_stage_init_timeout_to_diffusion_handshake(monkeypatch):
-    """Regression test for stage_init_timeout passing to complete_diffusion_handshake
-    in the diffusion stage path.
-    """
-    import vllm_omni.diffusion.data as diffusion_data_mod
-    import vllm_omni.diffusion.stage_diffusion_client as client_mod
+def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_replica_failure(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
-    from vllm_omni.platforms import current_omni_platform
 
     engine = object.__new__(AsyncOmniEngine)
-    engine.log_stats = False
     engine.model = "dummy-model"
     engine.config_path = "dummy-config"
-    engine.num_stages = 2
+    engine.num_stages = 1
     engine.async_chunk = False
     engine.diffusion_batch_size = 1
     engine.single_stage_mode = False
+    engine._single_stage_id_filter = None
     engine._omni_master_server = None
-    engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion", engine_args={})]
+    engine.stage_configs = [types.SimpleNamespace()]
 
-    metadata = types.SimpleNamespace(
-        stage_id=0,
-        stage_type="diffusion",
-        runtime_cfg={"devices": "0"},
-        prompt_expand_func=None,
-        final_output=True,
-        final_output_type="image",
-        default_sampling_params=None,
-        custom_process_input_func=None,
-        engine_input_source=None,
-        cfg_kv_collect_func=None,
-    )
-
-    captured_timeout = None
-    device_env_var = current_omni_platform.device_control_env_var
-    prev_device_env = os.environ.get(device_env_var)
-    os.environ[device_env_var] = "0"
+    cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)]
+    initialized_client = types.SimpleNamespace(shutdown=lambda: None)
+    replica1_started = threading.Event()
 
     monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
     monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
-    monkeypatch.setattr(engine_mod, "extract_stage_metadata", lambda _cfg: metadata)
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
-    monkeypatch.setattr(
-        engine_mod,
-        "finalize_initialized_stages",
-        lambda stage_clients, _input_processor: (
-            stage_clients,
-            [types.SimpleNamespace()],
-            [{"final_output_type": "image"}],
-        ),
-    )
-    monkeypatch.setattr(
-        diffusion_data_mod.OmniDiffusionConfig,
-        "from_kwargs",
-        classmethod(lambda cls, **kwargs: types.SimpleNamespace(parallel_config=types.SimpleNamespace(world_size=1))),
-    )
-    monkeypatch.setattr(
-        client_mod,
-        "spawn_diffusion_proc",
-        lambda model, od_cfg: (object(), "ipc://handshake", "ipc://request", "ipc://response"),
-    )
+    monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {}))
+    monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None))
 
-    def _capture_handshake_timeout(proc, handshake_address, handshake_timeout):
-        nonlocal captured_timeout
-        captured_timeout = handshake_timeout
+    def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock):
+        if plan.replica_id == 0:
+            assert replica1_started.wait(timeout=1)
+            raise RuntimeError("replica launch failed")
+        replica1_started.set()
+        time.sleep(0.05)
+        return initialized_client
 
-    monkeypatch.setattr(client_mod, "complete_diffusion_handshake", _capture_handshake_timeout)
-    monkeypatch.setattr(
-        client_mod.zmq,
-        "Context",
-        lambda: types.SimpleNamespace(socket=lambda _: types.SimpleNamespace(connect=lambda _: None)),
-    )
+    monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica)
 
-    try:
-        engine._initialize_stages(stage_init_timeout=302)
-    finally:
-        if prev_device_env is None:
-            os.environ.pop(device_env_var, None)
-        else:
-            os.environ[device_env_var] = prev_device_env
+    captured_cleanup: list[list[object]] = []
 
-    assert captured_timeout == 302
+    def _capture_shutdown(clients):
+        captured_cleanup.append(list(clients))
+
+    monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown)
+
+    with pytest.raises(RuntimeError, match="replica launch failed"):
+        engine._initialize_stages(stage_init_timeout=1)
+
+    assert captured_cleanup == [[initialized_client]]
 
 
-def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch):
-    """Regression test for stage_init_timeout reaching complete_stage_handshake
-    in the LLM stage path.
-    """
+def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch):
     import vllm_omni.engine.async_omni_engine as engine_mod
     from vllm_omni.platforms import current_omni_platform
 
     engine = object.__new__(AsyncOmniEngine)
-    engine.log_stats = False
     engine.model = "dummy-model"
     engine.single_stage_mode = False
     engine._omni_master_server = None
     engine.stage_configs = []
 
-    metadata = types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"})
     fake_vllm_config = types.SimpleNamespace()
-    fake_addresses = types.SimpleNamespace()
+    fake_addresses = types.SimpleNamespace(inputs=["in"], outputs=["out"], frontend_stats_publish_address=None)
     fake_proc = types.SimpleNamespace()
+    captured_timeout: int | None = None
 
-    captured_timeout = None
+    plan = ReplicaInitPlan(
+        replica_id=0,
+        num_replicas=1,
+        launch_mode="local",
+        stage_cfg=types.SimpleNamespace(engine_args={}, runtime=types.SimpleNamespace(devices="0")),
+        metadata=types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"}),
+        stage_connector_spec={},
+        omni_kv_connector=(None, None, None),
+        stage_vllm_config=fake_vllm_config,
+        executor_class=object,
+    )
 
     device_env_var = current_omni_platform.device_control_env_var
     prev_device_env = os.environ.get(device_env_var)
@@ -198,28 +538,22 @@ def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake(
 
     monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
     monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
-    monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (fake_vllm_config, object))
     monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: [])
-    monkeypatch.setattr(
-        engine_mod,
-        "spawn_stage_core",
-        lambda **_: (fake_addresses, fake_proc, "ipc://handshake"),
-    )
+    monkeypatch.setattr(engine_mod, "spawn_stage_core", lambda **_: (fake_addresses, fake_proc, "ipc://handshake"))
 
     def _capture_stage_timeout(_proc, _handshake_addr, _addresses, _vllm_cfg, handshake_timeout):
         nonlocal captured_timeout
         captured_timeout = handshake_timeout
 
     monkeypatch.setattr(engine_mod, "complete_stage_handshake", _capture_stage_timeout)
+    monkeypatch.setattr(
+        engine_mod.StageEngineCoreClientBase,
+        "make_async_mp_client",
+        staticmethod(lambda **_: types.SimpleNamespace(shutdown=lambda: None)),
+    )
 
     try:
-        engine._launch_llm_stage(
-            stage_cfg=types.SimpleNamespace(engine_args={}),
-            metadata=metadata,
-            stage_connector_spec={},
-            stage_init_timeout=302,
-            llm_stage_launch_lock=threading.Lock(),
-        )
+        engine._initialize_llm_replica(plan, 302, threading.Lock())
     finally:
         if prev_device_env is None:
             os.environ.pop(device_env_var, None)
@@ -229,155 +563,27 @@ def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake(
     assert captured_timeout == 302
 
 
-def test_launch_llm_stage_releases_launch_lock_before_complete_stage_handshake(monkeypatch):
-    """Regression test for parallel LLM stage startup during handshake wait."""
-    import vllm_omni.engine.async_omni_engine as engine_mod
-    from vllm_omni.platforms import current_omni_platform
-
-    engine = object.__new__(AsyncOmniEngine)
-    engine.log_stats = False
-    engine.model = "dummy-model"
-    engine.single_stage_mode = False
-    engine._omni_master_server = None
-    engine.stage_configs = []
-
-    fake_vllm_config = types.SimpleNamespace()
-    fake_addresses = types.SimpleNamespace()
-    shared_launch_lock = threading.Lock()
-    counter_lock = threading.Lock()
-    first_handshake_started = threading.Event()
-    second_stage_spawned = threading.Event()
-    allow_first_handshake_to_finish = threading.Event()
-    launch_errors: list[BaseException] = []
-    spawn_count = 0
-
-    device_env_var = current_omni_platform.device_control_env_var
-    prev_device_env = os.environ.get(device_env_var)
-    os.environ[device_env_var] = "0"
-
-    monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
-    monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
-    monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (fake_vllm_config, object))
-    monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: [])
-
-    def _spawn_stage_core(**_):
-        nonlocal spawn_count
-        with counter_lock:
-            spawn_count += 1
-            call_idx = spawn_count
-        if call_idx == 2:
-            second_stage_spawned.set()
-        return fake_addresses, types.SimpleNamespace(), f"ipc://handshake-{call_idx}"
-
-    def _complete_stage_handshake(_proc, handshake_address, _addresses, _vllm_cfg, _timeout):
-        if handshake_address == "ipc://handshake-1":
-            first_handshake_started.set()
-            assert second_stage_spawned.wait(timeout=1), (
-                "second stage did not reach spawn_stage_core while first stage waited in handshake"
-            )
-            assert allow_first_handshake_to_finish.wait(timeout=1), (
-                "second stage did not enter handshake while first stage was still waiting"
-            )
-        else:
-            allow_first_handshake_to_finish.set()
-
-    monkeypatch.setattr(engine_mod, "spawn_stage_core", _spawn_stage_core)
-    monkeypatch.setattr(engine_mod, "complete_stage_handshake", _complete_stage_handshake)
-
-    def _launch_stage(stage_id: int) -> None:
-        metadata = types.SimpleNamespace(stage_id=stage_id, runtime_cfg={"devices": str(stage_id)})
-        try:
-            engine._launch_llm_stage(
-                stage_cfg=types.SimpleNamespace(engine_args={}),
-                metadata=metadata,
-                stage_connector_spec={},
-                stage_init_timeout=302,
-                llm_stage_launch_lock=shared_launch_lock,
-            )
-        except BaseException as exc:  # pragma: no cover - surfaced through assertion below
-            launch_errors.append(exc)
-
-    try:
-        first_thread = threading.Thread(target=_launch_stage, args=(0,))
-        first_thread.start()
-        assert first_handshake_started.wait(timeout=1), "first stage never entered handshake"
-
-        second_thread = threading.Thread(target=_launch_stage, args=(1,))
-        second_thread.start()
-
-        first_thread.join(timeout=3)
-        second_thread.join(timeout=3)
-    finally:
-        if prev_device_env is None:
-            os.environ.pop(device_env_var, None)
-        else:
-            os.environ[device_env_var] = prev_device_env
-
-    assert not first_thread.is_alive()
-    assert not second_thread.is_alive()
-    assert second_stage_spawned.is_set()
-    assert not launch_errors
-
-
-def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
-    """Regression test for GLM-Image t2i preprocessing path.
-
-    Stage-0 InputProcessor must use OmniInputPreprocessor so text prompts with
-    mm_processor_kwargs go through multimodal preprocessing.
-    """
-    import vllm_omni.engine.async_omni_engine as engine_mod
-
-    class DummyStageEngineCoreClient:
-        def __init__(self, **kwargs):
-            self.kwargs = kwargs
-
-        def shutdown(self):
-            return None
-
-    class DummyOutputProcessor:
-        def __init__(self, **kwargs):
-            self.kwargs = kwargs
+def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch):
+    import vllm_omni.engine.stage_init_utils as init_mod
 
     class DummyInputProcessor:
         def __init__(self, vllm_config):
             self.vllm_config = vllm_config
             self.renderer = object()
-            self.input_preprocessor = object()
+            self.input_preprocessor = None
 
     class DummyOmniInputPreprocessor:
         def __init__(self, vllm_config, renderer=None):
             self.vllm_config = vllm_config
             self.renderer = renderer
 
-    monkeypatch.setattr(
-        engine_mod.StageEngineCoreClientBase,
-        "make_async_mp_client",
-        staticmethod(lambda **kwargs: DummyStageEngineCoreClient(**kwargs)),
-    )
-    monkeypatch.setattr(engine_mod, "MultimodalOutputProcessor", DummyOutputProcessor)
-    monkeypatch.setattr(engine_mod, "InputProcessor", DummyInputProcessor)
-    monkeypatch.setattr(engine_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor)
+    monkeypatch.setattr(init_mod, "InputProcessor", DummyInputProcessor)
+    monkeypatch.setattr(init_mod, "OmniInputPreprocessor", DummyOmniInputPreprocessor)
 
-    started = types.SimpleNamespace(
-        stage_id=0,
-        metadata=types.SimpleNamespace(stage_id=0, engine_output_type="token_ids"),
-        vllm_config=types.SimpleNamespace(model_config=types.SimpleNamespace(skip_tokenizer_init=True)),
-        executor_class=object,
-        engine_manager=object(),
-        coordinator=object(),
-        proc=None,
-        addresses=types.SimpleNamespace(
-            inputs=["inproc://input"],
-            outputs=["inproc://output"],
-            frontend_stats_publish_address=None,
-        ),
+    input_processor = build_stage0_input_processor(
+        types.SimpleNamespace(model_config=types.SimpleNamespace(try_get_generation_config=lambda: {}))
     )
 
-    engine = object.__new__(AsyncOmniEngine)
-
-    _stage_client, _out_proc, _vllm_cfg, input_processor = engine._attach_llm_stage(started)
-
-    assert input_processor is not None
     assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor)
     assert input_processor.input_preprocessor.renderer is input_processor.renderer
 

tests/engine/test_cfg_companion_tracker.py

@@ -1,7 +1,6 @@
 import pytest
 
 from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -25,18 +24,6 @@ def test_register_companion_and_cleanup():
     assert tracker.get_companion_request_ids("req1") == {}
 
 
-def test_attach_cfg_request_ids_clones_diffusion_params():
-    tracker = CfgCompanionTracker()
-    tracker.register_companion("req1", "cfg_text", "req1__cfg_text")
-
-    params = OmniDiffusionSamplingParams()
-    updated = tracker.attach_cfg_request_ids("req1", params)
-
-    assert updated is not params
-    assert params.cfg_kv_request_ids is None
-    assert updated.cfg_kv_request_ids == {"cfg_text": "req1__cfg_text"}
-
-
 def test_abort_parent_expands_to_companions_and_cleans_up_deferred_parent():
     tracker = CfgCompanionTracker()
     tracker.register_companion("req1", "cfg_text", "req1__cfg_text")

tests/engine/test_orchestrator.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import concurrent.futures
+import logging
 import queue
 import threading
 import time
@@ -16,8 +17,9 @@ from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.core_client import AsyncMPClient
 
-from vllm_omni.engine.orchestrator import Orchestrator
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
@@ -73,7 +75,7 @@ class FakeStageClient:
     def set_engine_outputs(self, outputs) -> None:
         return None
 
-    def process_engine_inputs(self, stage_list, prompt=None, streaming_context=None):
+    def process_engine_inputs(self, source_outputs, prompt=None, streaming_context=None):
         return list(self.next_inputs)
 
     async def abort_requests_async(self, request_ids: list[str]) -> None:
@@ -90,12 +92,33 @@ class FakeStageClient:
         self._diffusion_outputs.put_nowait(output)
 
 
+class FakeCollectiveRpcStageClient(FakeStageClient):
+    def __init__(self, *args, rpc_result: Any = None, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.rpc_result = rpc_result
+        self.collective_rpc_calls: list[tuple[str, float | None, tuple[Any, ...], dict[str, Any]]] = []
+
+    async def collective_rpc_async(
+        self,
+        *,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        normalized_kwargs = dict(kwargs or {})
+        self.collective_rpc_calls.append((method, timeout, args, normalized_kwargs))
+        return self.rpc_result
+
+
 class FakeOutputProcessor:
     def __init__(self, *, request_outputs: list[object] | None = None) -> None:
         self.request_outputs = list(request_outputs or [])
+        self.add_request_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = []
         self.abort_calls: list[list[str]] = []
 
-    def add_request(self, *_args, **_kwargs) -> None:
+    def add_request(self, *args, **kwargs) -> None:
+        self.add_request_calls.append((args, kwargs))
         return None
 
     def process_outputs(self, *_args, **_kwargs):
@@ -185,17 +208,60 @@ def _build_request_output(
     )
 
 
+def _build_stage_pools(
+    stage_clients: list[list[FakeStageClient]],
+    *,
+    output_processors: list[FakeOutputProcessor] | None = None,
+    stage_vllm_configs: list[object] | None = None,
+) -> list[StagePool]:
+    """Build StagePool list from per-stage replica lists.
+
+    ``stage_clients[i]`` is the list of FakeStageClient replicas for stage i.
+    """
+    num_stages = len(stage_clients)
+    if output_processors is None:
+        output_processors = [FakeOutputProcessor() for _ in stage_clients]
+    if stage_vllm_configs is None:
+        stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) for _ in stage_clients]
+
+    pools: list[StagePool] = []
+    for stage_id in range(num_stages):
+        clients = stage_clients[stage_id]
+        if clients[0].stage_type == "diffusion":
+            pools.append(StagePool(stage_id, clients[0]))
+        else:
+            pools.append(
+                StagePool(
+                    stage_id,
+                    clients,
+                    output_processor=output_processors[stage_id],
+                    stage_vllm_config=stage_vllm_configs[stage_id],
+                )
+            )
+    return pools
+
+
 def _build_harness(
     stage_clients: list[object],
     *,
     output_processors: list[object] | None = None,
     stage_vllm_configs: list[object] | None = None,
     async_chunk: bool = False,
+    stage_pools: list[StagePool] | None = None,
 ) -> OrchestratorFixture:
-    if output_processors is None:
-        output_processors = [FakeOutputProcessor() for _ in stage_clients]
-    if stage_vllm_configs is None:
-        stage_vllm_configs = [SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) for _ in stage_clients]
+    """Build an Orchestrator test harness.
+
+    Accepts either pre-built ``stage_pools`` or flat lists of single-replica
+    clients/processors.
+    """
+    if stage_pools is None:
+        # Wrap flat lists into per-stage single-replica lists.
+        nested_clients = [[c] for c in stage_clients]
+        stage_pools = _build_stage_pools(
+            nested_clients,
+            output_processors=output_processors,
+            stage_vllm_configs=stage_vllm_configs,
+        )
 
     ready_future: concurrent.futures.Future[tuple[Orchestrator, janus.Queue, janus.Queue, janus.Queue]] = (
         concurrent.futures.Future()
@@ -214,9 +280,7 @@ def _build_harness(
                 request_async_queue=request_queue.async_q,
                 output_async_queue=output_queue.async_q,
                 rpc_async_queue=rpc_queue.async_q,
-                stage_clients=stage_clients,
-                output_processors=output_processors,
-                stage_vllm_configs=stage_vllm_configs,
+                stage_pools=stage_pools,
                 async_chunk=async_chunk,
             )
             ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue))
@@ -283,6 +347,18 @@ async def _get_output_message(orchestrator_fixture: OrchestratorFixture, *, time
             return msg
 
 
+async def _get_rpc_message(orchestrator_fixture: OrchestratorFixture, *, timeout: float = 2.0) -> dict:
+    deadline = time.monotonic() + timeout
+    rpc_sync_q = orchestrator_fixture.queues[2].sync_q
+    while True:
+        if time.monotonic() >= deadline:
+            raise AssertionError("Timed out waiting for orchestrator rpc output")
+        try:
+            return rpc_sync_q.get_nowait()
+        except queue.Empty:
+            await asyncio.sleep(0.01)
+
+
 async def _enqueue_add_request(
     orchestrator_fixture: OrchestratorFixture,
     *,
@@ -387,6 +463,11 @@ def orchestrator_factory():
             q.close()
 
 
+# ---------------------------------------------------------------------------
+# Existing single-replica tests (adapted to StagePool interface)
+# ---------------------------------------------------------------------------
+
+
 @pytest.mark.asyncio
 async def test_run_two_stage_llm(orchestrator_factory) -> None:
     stage0 = FakeStageClient(stage_type="llm", final_output=False)
@@ -600,12 +681,425 @@ async def test_run_abort(orchestrator_factory) -> None:
         await _wait_for(lambda: len(stages[0].add_request_calls) == 1)
 
         await _enqueue_abort_request(orchestrator_fixture, ["req-abort"])
-        await _wait_for(lambda: all(stage.abort_calls for stage in stages))
+        await _wait_for(lambda: bool(stages[0].abort_calls))
 
-        for stage in stages:
-            assert stage.abort_calls == [["req-abort"]]
-        for proc in processors:
-            assert proc.abort_calls == [["req-abort"]]
+        assert stages[0].abort_calls == [["req-abort"]]
+        assert stages[1].abort_calls == []
         assert "req-abort" not in orchestrator_fixture.orchestrator.request_states
     finally:
         await _shutdown_orchestrator(orchestrator_fixture)
+
+
+# ---------------------------------------------------------------------------
+# Multi-replica tests
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_multi_replica_round_robin_distribution(orchestrator_factory) -> None:
+    """Two replicas at stage-0, single replica at stage-1.
+
+    Send two requests — they should land on different stage-0 replicas
+    (round-robin), then both forward to the single stage-1 replica.
+    """
+    stage0_r0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage0_r1 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(
+        stage_type="llm",
+        final_output=True,
+        next_inputs=[{"prompt_token_ids": [7, 8]}],
+    )
+
+    proc0 = FakeOutputProcessor(request_outputs=[_build_request_output("req-0", token_ids=[3], finished=True)])
+    proc1 = FakeOutputProcessor(request_outputs=[_build_request_output("req-0", token_ids=[10], finished=True)])
+
+    default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))
+    stage_pools = _build_stage_pools(
+        [[stage0_r0, stage0_r1], [stage1]],
+        output_processors=[proc0, proc1],
+        stage_vllm_configs=[default_vllm_cfg, default_vllm_cfg],
+    )
+
+    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
+
+    try:
+        # Request 0 → should land on replica 0 (RR starts at 0)
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-0",
+            prompt=SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]),
+            original_prompt={"prompt": "hello 0"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+        await _wait_for(lambda: len(stage0_r0.add_request_calls) == 1)
+        assert len(stage0_r1.add_request_calls) == 0
+
+        # Request 1 → should land on replica 1 (RR advances)
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-1",
+            prompt=SimpleNamespace(request_id="req-1", prompt_token_ids=[5, 6]),
+            original_prompt={"prompt": "hello 1"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+        await _wait_for(lambda: len(stage0_r1.add_request_calls) == 1)
+        assert len(stage0_r0.add_request_calls) == 1  # unchanged
+
+        # Complete req-0 at stage-0 replica-0 → should forward to stage-1
+        stage0_r0.push_engine_core_outputs(_engine_core_outputs("s0r0-raw", 1.0))
+        await _wait_for(lambda: len(stage1.add_request_calls) == 1)
+        assert stage1.add_request_calls[0][0].request_id == "req-0"
+
+        # Complete req-0 at stage-1 → final output
+        proc1.request_outputs = [_build_request_output("req-0", token_ids=[10], finished=True)]
+        stage1.push_engine_core_outputs(_engine_core_outputs("s1-raw", 2.0))
+        output_msg = await _get_output_message(orchestrator_fixture)
+
+        assert output_msg["request_id"] == "req-0"
+        assert output_msg["stage_id"] == 1
+        assert output_msg["finished"] is True
+        assert "req-0" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_multi_replica_abort_broadcasts_to_all_replicas(orchestrator_factory) -> None:
+    """Abort must be sent to every replica across all stages."""
+    stage0_r0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage0_r1 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(stage_type="llm", final_output=True)
+
+    proc0 = FakeOutputProcessor()
+    proc1 = FakeOutputProcessor()
+
+    default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))
+    stage_pools = _build_stage_pools(
+        [[stage0_r0, stage0_r1], [stage1]],
+        output_processors=[proc0, proc1],
+        stage_vllm_configs=[default_vllm_cfg, default_vllm_cfg],
+    )
+    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-abort-mr",
+            prompt=SimpleNamespace(request_id="req-abort-mr", prompt_token_ids=[1]),
+            original_prompt={"prompt": "cancel"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+        await _wait_for(lambda: len(stage0_r0.add_request_calls) == 1)
+
+        await _enqueue_abort_request(orchestrator_fixture, ["req-abort-mr"])
+        await _wait_for(lambda: bool(stage0_r0.abort_calls))
+
+        assert stage0_r0.abort_calls == [["req-abort-mr"]]
+        assert stage0_r1.abort_calls == []
+        assert stage1.abort_calls == []
+        assert "req-abort-mr" not in orchestrator_fixture.orchestrator.request_states
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_multi_replica_shutdown_all_replicas(orchestrator_factory) -> None:
+    """Shutdown must shut down every replica across all stages."""
+    stage0_r0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage0_r1 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(stage_type="llm", final_output=True)
+
+    default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))
+    stage_pools = _build_stage_pools(
+        [[stage0_r0, stage0_r1], [stage1]],
+        stage_vllm_configs=[default_vllm_cfg, default_vllm_cfg],
+    )
+    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
+
+    await _shutdown_orchestrator(orchestrator_fixture)
+
+    assert not orchestrator_fixture.thread.is_alive()
+    for client in [stage0_r0, stage0_r1, stage1]:
+        assert client.shutdown_calls == 1
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_submit_update_reuses_existing_binding() -> None:
+    """A request admitted to one replica must keep using that replica on updates."""
+    stage0_r0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage0_r1 = FakeStageClient(stage_type="llm", final_output=False)
+    pool = StagePool(
+        0,
+        [stage0_r0, stage0_r1],
+        output_processor=FakeOutputProcessor(),
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+
+    req0_state = OrchestratorRequestState(
+        request_id="req-0",
+        sampling_params_list=[_sampling_params()],
+        final_stage_id=0,
+    )
+    req1_state = OrchestratorRequestState(
+        request_id="req-1",
+        sampling_params_list=[_sampling_params()],
+        final_stage_id=0,
+    )
+
+    await pool.submit_initial("req-0", req0_state, SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]))
+    await pool.submit_update("req-0", req0_state, SimpleNamespace(request_id="req-0", prompt_token_ids=[3]))
+    await pool.submit_initial("req-1", req1_state, SimpleNamespace(request_id="req-1", prompt_token_ids=[4, 5]))
+    await pool.submit_update("req-1", req1_state, SimpleNamespace(request_id="req-1", prompt_token_ids=[6]))
+
+    assert pool.get_bound_replica_id("req-0") == 0
+    assert pool.get_bound_replica_id("req-1") == 1
+    assert len(stage0_r0.add_request_calls) == 2
+    assert len(stage0_r1.add_request_calls) == 2
+    assert stage0_r0.add_request_calls[0][0].request_id == "req-0"
+    assert stage0_r0.add_request_calls[1][0].request_id == "req-0"
+    assert stage0_r1.add_request_calls[0][0].request_id == "req-1"
+    assert stage0_r1.add_request_calls[1][0].request_id == "req-1"
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_submit_update_refreshes_output_processor_state() -> None:
+    output_processor = FakeOutputProcessor()
+
+    class AssertingStageClient(FakeStageClient):
+        async def add_request_async(self, *args, **kwargs) -> None:
+            if len(self.add_request_calls) == 1:
+                prompts = [call_kwargs["prompt"] for _, call_kwargs in output_processor.add_request_calls]
+                assert prompts == ["seg-1", "seg-2"]
+            await super().add_request_async(*args, **kwargs)
+
+    stage0 = AssertingStageClient(stage_type="llm", final_output=False)
+    pool = StagePool(
+        0,
+        [stage0],
+        output_processor=output_processor,
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+    req_state = OrchestratorRequestState(
+        request_id="req-0",
+        sampling_params_list=[_sampling_params()],
+        final_stage_id=0,
+    )
+
+    await pool.submit_initial(
+        "req-0",
+        req_state,
+        SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]),
+        prompt_text="seg-1",
+    )
+    await pool.submit_update(
+        "req-0",
+        req_state,
+        SimpleNamespace(request_id="req-0", prompt_token_ids=[3], resumable=True),
+        prompt_text="seg-2",
+    )
+
+    assert len(output_processor.add_request_calls) == 2
+    assert output_processor.add_request_calls[1][1]["prompt"] == "seg-2"
+
+
+@pytest.mark.asyncio
+async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> None:
+    class RecordingPool:
+        def __init__(self) -> None:
+            self.calls: list[tuple[str, Any]] = []
+
+        async def submit_update(self, request_id, req_state, request, *, prompt_text=None) -> int:
+            self.calls.append((request_id, prompt_text))
+            return 0
+
+    pool = RecordingPool()
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.request_states = {
+        "req-stream": OrchestratorRequestState(
+            request_id="req-stream",
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+        )
+    }
+    orchestrator.stage_pools = [pool]
+
+    await orchestrator._handle_streaming_update(
+        {
+            "request_id": "req-stream",
+            "prompt": SimpleNamespace(request_id="req-stream", prompt_token_ids=[1], resumable=True),
+            "sampling_params_list": [_sampling_params()],
+            "output_prompt_text": "segment-2",
+        }
+    )
+
+    assert pool.calls == [("req-stream", "segment-2")]
+    assert orchestrator.request_states["req-stream"].streaming.enabled is True
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_submit_initial_rolls_back_output_processor_when_client_submit_fails() -> None:
+    class FailingStageClient(FakeStageClient):
+        async def add_request_async(self, *args, **_kwargs) -> None:
+            raise RuntimeError("submit failed")
+
+    class TrackingOutputProcessor(FakeOutputProcessor):
+        def __init__(self) -> None:
+            super().__init__()
+            self.added_request_ids: list[str] = []
+            self.removed_request_ids: list[str] = []
+
+        def add_request(self, request, *_args, **_kwargs) -> None:
+            self.added_request_ids.append(request.request_id)
+
+        def remove_request(self, request_id: str) -> None:
+            self.removed_request_ids.append(request_id)
+
+    client = FailingStageClient(stage_type="llm", final_output=False)
+    output_processor = TrackingOutputProcessor()
+    pool = StagePool(
+        0,
+        [client],
+        output_processor=output_processor,
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+    req_state = OrchestratorRequestState(
+        request_id="req-0",
+        sampling_params_list=[_sampling_params()],
+        final_stage_id=0,
+    )
+
+    with pytest.raises(RuntimeError, match="submit failed"):
+        await pool.submit_initial("req-0", req_state, SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]))
+
+    assert output_processor.added_request_ids == ["req-0"]
+    assert output_processor.removed_request_ids == ["req-0"]
+    assert pool.get_bound_replica_id("req-0") is None
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_abort_requests_logs_when_binding_is_missing(caplog) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=False)
+    pool = StagePool(
+        0,
+        [stage0],
+        output_processor=FakeOutputProcessor(),
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+
+    target_logger = logging.getLogger("vllm_omni.engine.stage_pool")
+    target_logger.addHandler(caplog.handler)
+    prev_level = target_logger.level
+    target_logger.setLevel(logging.DEBUG)
+    try:
+        await pool.abort_requests(["missing-req"])
+    finally:
+        target_logger.removeHandler(caplog.handler)
+        target_logger.setLevel(prev_level)
+
+    assert not stage0.abort_calls
+    assert "abort: no binding for req=missing-req in stage-0" in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_collective_rpc_ignores_invalid_stage_ids(orchestrator_factory, caplog) -> None:
+    stage0 = FakeCollectiveRpcStageClient(stage_type="llm", final_output=True, rpc_result={"stage": 0})
+    stage1 = FakeCollectiveRpcStageClient(stage_type="llm", final_output=True, rpc_result={"stage": 1})
+    stage_pools = _build_stage_pools(
+        [[stage0], [stage1]],
+        output_processors=[FakeOutputProcessor(), FakeOutputProcessor()],
+        stage_vllm_configs=[
+            SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+            SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+        ],
+    )
+    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
+
+    try:
+        target_logger = logging.getLogger("vllm_omni.engine.orchestrator")
+        target_logger.addHandler(caplog.handler)
+        prev_level = target_logger.level
+        target_logger.setLevel(logging.WARNING)
+        try:
+            orchestrator_fixture.request_sync_q.put_nowait(
+                {
+                    "type": "collective_rpc",
+                    "rpc_id": "rpc-1",
+                    "method": "list_loras",
+                    "stage_ids": [99, 1],
+                }
+            )
+
+            msg = await _get_rpc_message(orchestrator_fixture)
+        finally:
+            target_logger.removeHandler(caplog.handler)
+            target_logger.setLevel(prev_level)
+
+        assert msg["type"] == "collective_rpc_result"
+        assert msg["rpc_id"] == "rpc-1"
+        assert msg["stage_ids"] == [1]
+        assert msg["results"] == [{"stage": 1}]
+        assert not stage0.collective_rpc_calls
+        assert len(stage1.collective_rpc_calls) == 1
+        assert "collective_rpc: ignoring invalid stage_id 99" in caplog.text
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator_factory) -> None:
+    """CFG companions should be routed to the same stage-0 replica as their parent."""
+    stage0_r0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage0_r1 = FakeStageClient(stage_type="llm", final_output=False)
+    default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))
+    stage_pools = _build_stage_pools(
+        [[stage0_r0, stage0_r1]],
+        output_processors=[FakeOutputProcessor()],
+        stage_vllm_configs=[default_vllm_cfg],
+    )
+    orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools)
+
+    try:
+        # Consume replica-0 first so the parent request binds to replica-1.
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="warmup",
+            prompt=SimpleNamespace(request_id="warmup", prompt_token_ids=[0]),
+            original_prompt={"prompt": "warmup"},
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+        )
+        await _wait_for(lambda: len(stage0_r0.add_request_calls) == 1)
+
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="parent",
+            prompt=SimpleNamespace(request_id="parent", prompt_token_ids=[1, 2]),
+            original_prompt={"prompt": "parent"},
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+        )
+        await _wait_for(lambda: len(stage0_r1.add_request_calls) == 1)
+
+        orchestrator_fixture.request_sync_q.put_nowait(
+            {
+                "type": "add_companion_request",
+                "companion_id": "parent-neg",
+                "parent_id": "parent",
+                "role": "negative",
+                "prompt": SimpleNamespace(request_id="parent-neg", prompt_token_ids=[9]),
+                "companion_prompt_text": {"prompt": "negative"},
+                "sampling_params_list": [_sampling_params()],
+            }
+        )
+        await _wait_for(lambda: len(stage0_r1.add_request_calls) == 2)
+
+        assert stage_pools[0].get_bound_replica_id("parent") == 1
+        assert stage_pools[0].get_bound_replica_id("parent-neg") == 1
+        assert len(stage0_r0.add_request_calls) == 1
+        assert stage0_r1.add_request_calls[0][0].request_id == "parent"
+        assert stage0_r1.add_request_calls[1][0].request_id == "parent-neg"
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)

tests/engine/test_orchestrator_error_handling.py

@@ -124,8 +124,8 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa
 @pytest.mark.asyncio
 async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) -> None:
     """When a diffusion stage returns an OmniRequestOutput with a non-None
-    error, the orchestrator must route it as a finished output message and
-    clean up the request state.
+    error, the orchestrator must route it as an error message and clean up
+    the request state.
     """
     stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
     orchestrator_fixture = orchestrator_factory([stage0])
@@ -148,10 +148,10 @@ async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) -
 
         msg = await _get_any_output_message(orchestrator_fixture)
 
-        assert msg["type"] == "output"
+        assert msg["type"] == "error"
         assert msg["request_id"] == "req-err"
-        assert msg["finished"] is True
-        assert msg["engine_outputs"].error == "gpu fault"
+        assert msg["stage_id"] == 0
+        assert msg["error"] == "gpu fault"
 
         # Request state should be cleaned up.
         await _wait_for(lambda: "req-err" not in orchestrator_fixture.orchestrator.request_states)

tests/engine/test_orchestrator_kv_sender_info.py

@@ -7,18 +7,18 @@ from vllm import SamplingParams
 from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
 from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
 from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
+from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _DummySenderStage:
+    stage_type = "llm"
+    final_output = False
+
     def __init__(self, sender_info):
         self._sender_info = sender_info
-        self.engine_outputs = None
-
-    def set_engine_outputs(self, outputs):
-        self.engine_outputs = outputs
 
     def get_kv_sender_info(self):
         return self._sender_info
@@ -26,6 +26,7 @@ class _DummySenderStage:
 
 class _DummyDiffusionStage:
     stage_type = "diffusion"
+    final_output = True
     custom_process_input_func = None
 
     def __init__(self, engine_input_source=None):
@@ -43,6 +44,15 @@ class _DummyDiffusionStage:
         )
 
 
+def _build_sender_pool(stage_id: int, sender_info: dict[str, object]) -> StagePool:
+    return StagePool(
+        stage_id,
+        _DummySenderStage(sender_info),
+        output_processor=object(),
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+
+
 def test_stage_engine_core_client_builds_kv_sender_info_from_tcp_address():
     client = object.__new__(StageEngineCoreClient)
     client.stage_id = 0
@@ -126,14 +136,13 @@ def test_stage_engine_core_client_preserves_explicit_loopback_sender_host():
 
 def test_forward_to_diffusion_attaches_kv_sender_info():
     orchestrator = object.__new__(Orchestrator)
-    sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
     diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+    sender_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151})
+    diffusion_pool = StagePool(1, diffusion_stage)
 
     orchestrator.num_stages = 2
-    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator.stage_pools = [sender_pool, diffusion_pool]
     orchestrator._cfg_tracker = CfgCompanionTracker()
-    orchestrator.stage_vllm_configs = [None, None]
-    orchestrator.output_processors = [None, None]
 
     params = OmniDiffusionSamplingParams()
     req_state = OrchestratorRequestState(
@@ -144,9 +153,8 @@ def test_forward_to_diffusion_attaches_kv_sender_info():
     )
 
     output = SimpleNamespace(request_id="req-1", finished=True)
-    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", 0, output, req_state))
+    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", sender_pool.stage_id, output, req_state))
 
-    assert sender_stage.engine_outputs == [output]
     assert diffusion_stage.calls[0]["request_id"] == "req-1"
     assert diffusion_stage.calls[0]["kv_sender_info"] == {
         0: {"host": "10.0.0.2", "zmq_port": 50151},
@@ -156,15 +164,14 @@ def test_forward_to_diffusion_attaches_kv_sender_info():
 
 def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
     orchestrator = object.__new__(Orchestrator)
-    source_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
-    previous_stage = _DummySenderStage({"host": "10.0.0.9", "zmq_port": 59999})
     diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+    source_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151})
+    previous_pool = _build_sender_pool(1, {"host": "10.0.0.9", "zmq_port": 59999})
+    diffusion_pool = StagePool(2, diffusion_stage)
 
     orchestrator.num_stages = 3
-    orchestrator.stage_clients = [source_stage, previous_stage, diffusion_stage]
+    orchestrator.stage_pools = [source_pool, previous_pool, diffusion_pool]
     orchestrator._cfg_tracker = CfgCompanionTracker()
-    orchestrator.stage_vllm_configs = [None, None, None]
-    orchestrator.output_processors = [None, None, None]
 
     params = OmniDiffusionSamplingParams()
     req_state = OrchestratorRequestState(
@@ -175,9 +182,8 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
     )
 
     output = SimpleNamespace(request_id="req-3", finished=True)
-    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", 1, output, req_state))
+    asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", previous_pool.stage_id, output, req_state))
 
-    assert previous_stage.engine_outputs == [output]
     assert diffusion_stage.calls[0]["kv_sender_info"] == {
         0: {"host": "10.0.0.2", "zmq_port": 50151},
     }
@@ -185,9 +191,10 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info():
 
 def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs():
     orchestrator = object.__new__(Orchestrator)
-    sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151})
     diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
     diffusion_stage.custom_process_input_func = lambda *_args, **_kwargs: []
+    sender_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151})
+    diffusion_pool = StagePool(1, diffusion_stage)
 
     class _AsyncQueue:
         def __init__(self):
@@ -197,10 +204,8 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs():
             self.items.append(item)
 
     orchestrator.num_stages = 2
-    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator.stage_pools = [sender_pool, diffusion_pool]
     orchestrator._cfg_tracker = CfgCompanionTracker()
-    orchestrator.stage_vllm_configs = [None, None]
-    orchestrator.output_processors = [None, None]
     orchestrator.output_async_queue = _AsyncQueue()
     orchestrator.request_states = {}
     orchestrator._pd_kv_params = {}
@@ -217,7 +222,6 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs():
     output = SimpleNamespace(request_id="req-empty", finished=True)
     asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-empty", 0, output, req_state))
 
-    assert sender_stage.engine_outputs == [output]
     assert diffusion_stage.calls == []
     assert len(orchestrator.output_async_queue.items) == 1
     terminal_msg = orchestrator.output_async_queue.items[0]
@@ -231,10 +235,11 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs():
 
 def test_prewarm_diffusion_attaches_kv_sender_info():
     orchestrator = object.__new__(Orchestrator)
-    sender_stage = _DummySenderStage({"host": "10.0.0.3", "zmq_port": 50151})
     diffusion_stage = _DummyDiffusionStage(engine_input_source=[0])
+    sender_pool = _build_sender_pool(0, {"host": "10.0.0.3", "zmq_port": 50151})
+    diffusion_pool = StagePool(1, diffusion_stage)
 
-    orchestrator.stage_clients = [sender_stage, diffusion_stage]
+    orchestrator.stage_pools = [sender_pool, diffusion_pool]
     orchestrator.num_stages = 2
 
     req_state = OrchestratorRequestState(

tests/entrypoints/test_serve.py

@@ -38,6 +38,7 @@ def _make_headless_args() -> argparse.Namespace:
     return argparse.Namespace(
         model="fake-model",
         stage_id=3,
+        replica_id=0,
         omni_master_address="127.0.0.1",
         omni_master_port=26000,
         api_server_count=0,
@@ -45,6 +46,7 @@ def _make_headless_args() -> argparse.Namespace:
         stage_configs_path=None,
         log_stats=False,
         disable_log_stats=False,
+        stage_init_timeout=600,
     )
 
 
@@ -59,6 +61,7 @@ def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker
         node_rank_within_dp=0,
     )
     vllm_config = mocker.Mock(parallel_config=parallel_config)
+    vllm_config.needs_dp_coordinator = False
     executor_class = mocker.Mock()
     engine_manager = mocker.Mock()
 
@@ -99,6 +102,7 @@ def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker
         omni_stage_id=3,
         omni_stage_config=stage_cfg,
         coordinator=None,
+        replica_id=0,
     )
     mock_manager_cls.assert_called_once()
     manager_kwargs = mock_manager_cls.call_args.kwargs
@@ -155,13 +159,56 @@ def test_run_headless_honors_explicit_log_stats_flag(mocker: MockerFixture) -> N
     assert manager_kwargs["log_stats"] is True
 
 
+def test_run_headless_registers_llm_replica_id(mocker: MockerFixture) -> None:
+    args = _make_headless_args()
+    args.replica_id = 2
+    stage_cfg = mocker.Mock(stage_id=3)
+    parallel_config = mocker.Mock(
+        data_parallel_size_local=1,
+        data_parallel_rank=0,
+        data_parallel_rank_local=0,
+        node_rank_within_dp=0,
+    )
+    vllm_config = mocker.Mock(parallel_config=parallel_config)
+    vllm_config.needs_dp_coordinator = False
+    engine_manager = mocker.Mock()
+
+    mocker.patch(
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
+        return_value=("/fake/stages.yaml", [stage_cfg]),
+    )
+    mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment")
+    mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=mocker.Mock())
+    mocker.patch("vllm_omni.engine.stage_init_utils.get_stage_connector_spec", return_value={})
+    mocker.patch("vllm_omni.engine.stage_init_utils.build_engine_args_dict", return_value={})
+    mocker.patch(
+        "vllm_omni.distributed.omni_connectors.utils.initialization.resolve_omni_kv_config_for_stage",
+        return_value=(None, None, None),
+    )
+    mocker.patch(
+        "vllm_omni.engine.stage_init_utils.build_vllm_config",
+        return_value=(vllm_config, mocker.Mock()),
+    )
+    mock_register = mocker.patch(
+        "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
+        return_value="tcp://127.0.0.1:26001",
+    )
+    mocker.patch("vllm.v1.engine.utils.CoreEngineProcManager", return_value=engine_manager)
+    mocker.patch("signal.signal")
+
+    run_headless(args)
+
+    assert mock_register.call_args.kwargs["replica_id"] == 2
+
+
 def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFixture) -> None:
     args = _make_headless_args()
+    args.replica_id = 1
     stage_cfg = mocker.Mock(stage_id=3, stage_type="diffusion")
     stage_cfg.engine_args = mocker.Mock()
     stage_cfg.engine_input_source = []
     stage_cfgs = [stage_cfg]
-    metadata = mocker.Mock(stage_id=3)
+    metadata = mocker.Mock(stage_id=3, stage_type="diffusion")
     od_config = mocker.Mock()
     proc = mocker.Mock()
     proc.exitcode = 0
@@ -199,6 +246,7 @@ def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFix
         omni_stage_id=3,
         omni_stage_config=stage_cfg,
         return_addresses=True,
+        replica_id=1,
     )
     mock_spawn.assert_called_once_with(
         "fake-model",
@@ -207,5 +255,5 @@ def test_run_headless_launches_diffusion_stage_via_omni_master(mocker: MockerFix
         request_address="tcp://127.0.0.1:26002",
         response_address="tcp://127.0.0.1:26003",
     )
-    mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001")
+    mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001", 600)
     proc.join.assert_called_once_with()

tests/entrypoints/test_utils.py

@@ -360,6 +360,47 @@ class TestLoadAndResolveStageConfigs:
         assert len(stage_configs) == 1
         assert "dtype" in stage_configs[0]["engine_args"]
 
+    def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas(
+        self, tmp_path, mocker: MockerFixture
+    ):
+        deploy_path = tmp_path / "qwen3_multi.yaml"
+        deploy_path.write_text(
+            'stages:\n  - stage_id: 0\n    devices: "0"\n  - stage_id: 1\n    devices: "1,2,3"\n    num_replicas: 3\n',
+            encoding="utf-8",
+        )
+
+        returned_stage_configs = [
+            create_config({"stage_id": 0, "runtime": {"devices": "0"}, "engine_args": {"model": "dummy"}}),
+            create_config(
+                {
+                    "stage_id": 1,
+                    "runtime": {"devices": "1,2,3", "num_replicas": 3},
+                    "engine_args": {"model": "dummy"},
+                }
+            ),
+        ]
+        load_stage_configs = mocker.patch(
+            "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+            return_value=returned_stage_configs,
+        )
+
+        config_path, stage_configs = load_and_resolve_stage_configs(
+            model="dummy-model",
+            stage_configs_path=str(deploy_path),
+            kwargs={},
+        )
+
+        load_stage_configs.assert_called_once_with(
+            "dummy-model",
+            base_engine_args={},
+            deploy_config_path=str(deploy_path),
+            stage_overrides=None,
+        )
+        assert config_path == str(deploy_path)
+        assert len(stage_configs) == 2
+        assert stage_configs[1].runtime.num_replicas == 3
+        assert stage_configs[1].runtime.devices == "1,2,3"
+
 
 class TestLoadStageConfigsFromYaml:
     """Regression tests for stage-config loading and merging."""

tests/helpers/runtime.py

@@ -385,7 +385,8 @@ class OmniServerStageCli(OmniServer):
         self.stage_ids = stage_ids or self._load_stage_ids(resolved_cfg)
         if 0 not in self.stage_ids:
             raise ValueError(f"Stage CLI test requires stage_id=0 in config: {stage_config_path}")
-        self.stage_procs: dict[int, subprocess.Popen] = {}
+        self.stage_replica_counts = self._load_stage_replica_counts(resolved_cfg)
+        self.stage_procs: dict[tuple[int, int], subprocess.Popen] = {}
         self.proc = None
 
     @staticmethod
@@ -415,6 +416,19 @@ class OmniServerStageCli(OmniServer):
                 runtime_devices[int(stage_id)] = str(devices)
         return runtime_devices
 
+    @staticmethod
+    def _load_stage_replica_counts(resolved_config: dict) -> dict[int, int]:
+        replica_counts: dict[int, int] = {}
+        for stage in OmniServerStageCli._stage_entries(resolved_config):
+            stage_id = stage.get("stage_id")
+            if stage_id is None:
+                continue
+            replica_counts[int(stage_id)] = max(
+                1,
+                int(stage.get("num_replicas") or stage.get("runtime", {}).get("num_replicas", 1)),
+            )
+        return replica_counts
+
     @classmethod
     def _parse_device_list(cls, devices: str | int) -> list[str]:
         if isinstance(devices, int):
@@ -452,13 +466,29 @@ class OmniServerStageCli(OmniServer):
 
         return ",".join(visible_device_list[idx] for idx in logical_ids)
 
-    def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str) -> None:
-        mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, devices)
+    def _devices_for_replica(self, stage_id: int, devices: str, replica_id: int) -> str:
+        replica_count = self.stage_replica_counts.get(stage_id, 1)
+        if replica_count == 1:
+            return devices
+
+        device_list = self._parse_device_list(devices)
+        if len(device_list) % replica_count != 0:
+            raise ValueError(
+                f"Stage {stage_id} has {len(device_list)} device(s) for {replica_count} replica(s); "
+                "device count must be divisible by replica count"
+            )
+        devices_per_replica = len(device_list) // replica_count
+        start = replica_id * devices_per_replica
+        return ",".join(device_list[start : start + devices_per_replica])
+
+    def _set_stage_device_env(self, stage_id: int, env: dict[str, str], devices: str, replica_id: int = 0) -> None:
+        replica_devices = self._devices_for_replica(stage_id, devices, replica_id)
+        mapped_devices = self._map_stage_devices(stage_id, self.visible_device_list, replica_devices)
         env_var = getattr(current_omni_platform, "device_control_env_var", None)
         if env_var:
             env[env_var] = mapped_devices
 
-    def _build_stage_cmd(self, stage_id: int, *, headless: bool) -> list[str]:
+    def _build_stage_cmd(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> list[str]:
         cmd = [
             sys.executable,
             "-m",
@@ -474,6 +504,8 @@ class OmniServerStageCli(OmniServer):
             self.host,
             "--omni-master-port",
             str(self.master_port),
+            "--replica-id",
+            str(replica_id),
         ]
 
         if headless:
@@ -484,7 +516,7 @@ class OmniServerStageCli(OmniServer):
         cmd += self.serve_args
         return cmd
 
-    def _launch_stage(self, stage_id: int, *, headless: bool) -> None:
+    def _launch_stage(self, stage_id: int, *, headless: bool, replica_id: int = 0) -> None:
         env = os.environ.copy()
         env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
         if self.env_dict is not None:
@@ -492,19 +524,20 @@ class OmniServerStageCli(OmniServer):
 
         devices = self.stage_runtime_devices.get(stage_id)
         if devices:
-            self._set_stage_device_env(stage_id, env, devices)
+            self._set_stage_device_env(stage_id, env, devices, replica_id=replica_id)
 
-        cmd = self._build_stage_cmd(stage_id, headless=headless)
-        print(f"Launching OmniServerStageCli stage {stage_id}: {' '.join(cmd)}")
+        cmd = self._build_stage_cmd(stage_id, headless=headless, replica_id=replica_id)
+        print(f"Launching OmniServerStageCli stage {stage_id} replica {replica_id}: {' '.join(cmd)}")
         # Capture each subprocess's stdout+stderr to a per-stage log file so
         # debugging "Stage N exited before API server ready" doesn't rely on
         # guessing; the file is surfaced in the RuntimeError message.
-        log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_{self.master_port}.log"
+        log_path = Path(tempfile.gettempdir()) / f"omni_stage_{stage_id}_replica_{replica_id}_{self.master_port}.log"
         self._stage_log_paths = getattr(self, "_stage_log_paths", {})
-        self._stage_log_paths[stage_id] = log_path
+        stage_key = (stage_id, replica_id)
+        self._stage_log_paths[stage_key] = log_path
         log_fh = open(log_path, "w", buffering=1)  # noqa: SIM115 - closed in __exit__
         self._stage_log_files = getattr(self, "_stage_log_files", {})
-        self._stage_log_files[stage_id] = log_fh
+        self._stage_log_files[stage_key] = log_fh
         proc = subprocess.Popen(
             cmd,
             env=env,
@@ -512,36 +545,39 @@ class OmniServerStageCli(OmniServer):
             stdout=log_fh,
             stderr=subprocess.STDOUT,
         )
-        self.stage_procs[stage_id] = proc
-        if stage_id == 0:
+        self.stage_procs[stage_key] = proc
+        if stage_id == 0 and replica_id == 0:
             self.proc = proc
 
     def _ensure_stage_processes_alive(self) -> None:
-        for stage_id, proc in self.stage_procs.items():
+        for (stage_id, replica_id), proc in self.stage_procs.items():
             ret = proc.poll()
             if ret is not None:
-                log_path = getattr(self, "_stage_log_paths", {}).get(stage_id)
+                log_path = getattr(self, "_stage_log_paths", {}).get((stage_id, replica_id))
                 tail = ""
                 if log_path and log_path.exists():
                     try:
                         with open(log_path, encoding="utf-8", errors="replace") as f:
                             lines = f.readlines()
-                        tail = "\n=== Last 60 lines of stage {} log ({}) ===\n{}".format(
-                            stage_id, log_path, "".join(lines[-60:]) or "<empty>"
+                        tail = "\n=== Last 60 lines of stage {} replica {} log ({}) ===\n{}".format(
+                            stage_id, replica_id, log_path, "".join(lines[-60:]) or "<empty>"
                         )
                     except Exception as exc:  # pragma: no cover - diagnostic only
                         tail = f"\n<failed to read stage log {log_path}: {exc}>"
-                raise RuntimeError(f"Stage {stage_id} exited with code {ret} before API server became ready.{tail}")
+                raise RuntimeError(
+                    f"Stage {stage_id} replica {replica_id} exited with code {ret} before API server became ready.{tail}"
+                )
 
     def _start_server(self) -> None:
         ordered_stage_ids = [0, *[stage_id for stage_id in self.stage_ids if stage_id != 0]]
 
-        self._launch_stage(0, headless=False)
+        self._launch_stage(0, headless=False, replica_id=0)
         time.sleep(2)
         self._ensure_stage_processes_alive()
 
         for stage_id in ordered_stage_ids[1:]:
-            self._launch_stage(stage_id, headless=True)
+            for replica_id in range(self.stage_replica_counts.get(stage_id, 1)):
+                self._launch_stage(stage_id, headless=True, replica_id=replica_id)
 
         max_wait = 1200
         start_time = time.time()
@@ -568,15 +604,15 @@ class OmniServerStageCli(OmniServer):
         whatever state the stage was in when it was torn down.
         """
         log_paths = getattr(self, "_stage_log_paths", {}) or {}
-        for stage_id in sorted(log_paths):
-            log_path = log_paths[stage_id]
+        for stage_id, replica_id in sorted(log_paths):
+            log_path = log_paths[(stage_id, replica_id)]
             if not log_path or not log_path.exists():
                 continue
             try:
                 with open(log_path, encoding="utf-8", errors="replace") as f:
                     lines = f.readlines()
             except Exception as exc:  # pragma: no cover - diagnostic only
-                print(f"[OmniServerStageCli] stage {stage_id} log read failed: {exc}", flush=True)
+                print(f"[OmniServerStageCli] stage {stage_id} replica {replica_id} log read failed: {exc}", flush=True)
                 continue
             total = len(lines)
             if total <= head_lines + tail_lines:
@@ -587,18 +623,18 @@ class OmniServerStageCli(OmniServer):
                 head_chunk = lines[:head_lines]
                 tail_chunk = lines[-tail_lines:]
                 elided = total - head_lines - tail_lines
-            print(f"\n=== stage {stage_id} log HEAD ({log_path}) ===", flush=True)
+            print(f"\n=== stage {stage_id} replica {replica_id} log HEAD ({log_path}) ===", flush=True)
             print("".join(head_chunk).rstrip("\n"), flush=True)
             if tail_chunk:
                 print(f"\n... [{elided} lines elided] ...", flush=True)
-                print(f"\n=== stage {stage_id} log TAIL ({log_path}) ===", flush=True)
+                print(f"\n=== stage {stage_id} replica {replica_id} log TAIL ({log_path}) ===", flush=True)
                 print("".join(tail_chunk).rstrip("\n"), flush=True)
-            print(f"=== end stage {stage_id} log ===\n", flush=True)
+            print(f"=== end stage {stage_id} replica {replica_id} log ===\n", flush=True)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self._dump_stage_logs_for_debug()
-        for stage_id in sorted(self.stage_procs, reverse=True):
-            proc = self.stage_procs[stage_id]
+        for stage_key in sorted(self.stage_procs, reverse=True):
+            proc = self.stage_procs[stage_key]
             if proc.poll() is None:
                 self._kill_process_tree(proc.pid)
         run_forced_gpu_cleanup_round()

tests/helpers/stage_config.py

@@ -387,6 +387,79 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = {
             },
         },
     },
+    "qwen3_omni_moe_multi_replicas_4gpu": {
+        "base_config": "qwen3_omni_moe.yaml",
+        "async_chunk": True,
+        "stages": [
+            {
+                "stage_id": 0,
+                "devices": "0",
+                "gpu_memory_utilization": 0.85,
+                "max_num_seqs": 6,
+                "max_model_len": 32768,
+                "mm_processor_cache_gb": 0,
+                "load_format": "dummy",
+                "default_sampling_params": {"max_tokens": 150, "ignore_eos": False},
+            },
+            {
+                "stage_id": 1,
+                "devices": "1,2,3",
+                "num_replicas": 3,
+                "gpu_memory_utilization": 0.6,
+                "max_num_seqs": 2,
+                "max_model_len": 32768,
+                "load_format": "dummy",
+                "default_sampling_params": {"max_tokens": 1000},
+            },
+            {
+                "stage_id": 2,
+                "devices": "1,2,3",
+                "num_replicas": 3,
+                "gpu_memory_utilization": 0.1,
+                "max_num_seqs": 2,
+                "max_num_batched_tokens": 65536,
+                "load_format": "dummy",
+                "default_sampling_params": {"max_tokens": 2000},
+            },
+        ],
+    },
+    "bagel_multi_replicas_4gpu": {
+        "base_config": "bagel.yaml",
+        "async_chunk": False,
+        "stages": [
+            {
+                "stage_id": 0,
+                "devices": "0",
+                "max_num_seqs": 6,
+                "max_num_batched_tokens": 16384,
+                "gpu_memory_utilization": 0.45,
+                "load_format": "dummy",
+                "default_sampling_params": {
+                    "temperature": 0.4,
+                    "top_p": 0.9,
+                    "top_k": 1,
+                    "max_tokens": 256,
+                    "detokenize": False,
+                },
+            },
+            {
+                "stage_id": 1,
+                "devices": "1,2,3",
+                "num_replicas": 3,
+                "max_num_seqs": 1,
+                "enforce_eager": True,
+                "gpu_memory_utilization": 0.7,
+                "load_format": "dummy",
+                "default_sampling_params": {
+                    "seed": 42,
+                    "num_inference_steps": 2,
+                    "guidance_scale": 0.0,
+                    "height": 512,
+                    "width": 512,
+                },
+            },
+        ],
+    },
     "bagel": {
         "base_config": "bagel.yaml",
         "stages": [

tests/helpers/test_runtime_stage_cli.py

@@ -0,0 +1,56 @@
+"""Tests for OmniServerStageCli process planning helpers."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from tests.helpers.runtime import OmniServerStageCli
+
+
+def _write_stage_config(tmp_path: Path) -> str:
+    path = tmp_path / "stages.yaml"
+    path.write_text(
+        """
+stages:
+  - stage_id: 0
+    devices: "0"
+  - stage_id: 1
+    devices: "1,2,3"
+    num_replicas: 3
+""".strip(),
+        encoding="utf-8",
+    )
+    return str(path)
+
+
+def test_stage_cli_builds_headless_replica_cmd(tmp_path: Path) -> None:
+    server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), ["--disable-log-stats"])
+
+    cmd = server._build_stage_cmd(1, headless=True, replica_id=2)
+
+    assert "--headless" in cmd
+    assert cmd[cmd.index("--stage-id") + 1] == "1"
+    assert cmd[cmd.index("--replica-id") + 1] == "2"
+    assert cmd[cmd.index("--stage-configs-path") + 1] == server.stage_config_path
+
+
+def test_stage_cli_splits_devices_per_replica(tmp_path: Path) -> None:
+    server = OmniServerStageCli("fake-model", _write_stage_config(tmp_path), [])
+
+    assert server._devices_for_replica(1, "1,2,3", 0) == "1"
+    assert server._devices_for_replica(1, "1,2,3", 1) == "2"
+    assert server._devices_for_replica(1, "1,2,3", 2) == "3"
+
+
+def test_stage_cli_maps_visible_devices_after_replica_split(tmp_path: Path) -> None:
+    server = OmniServerStageCli(
+        "fake-model",
+        _write_stage_config(tmp_path),
+        [],
+        env_dict={"CUDA_VISIBLE_DEVICES": "4,5,6,7"},
+    )
+    env: dict[str, str] = {}
+
+    server._set_stage_device_env(1, env, "1,2,3", replica_id=1)
+
+    assert env["CUDA_VISIBLE_DEVICES"] == "6"

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -44,16 +44,12 @@ def _transfer_manager(
 
 
 def test_text2flow_supports_batched_source_outputs():
-    stage_list = [
-        SimpleNamespace(
-            engine_outputs=[
-                _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[1, 2]])}),
-                _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[3, 4]])}),
-            ]
-        )
+    source_outputs = [
+        _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[1, 2]])}),
+        _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[3, 4]])}),
     ]
 
-    outputs = text2flow(stage_list=stage_list, engine_input_source=[0], prompt=None)
+    outputs = text2flow(source_outputs=source_outputs, prompt=None)
 
     assert len(outputs) == 2
     assert outputs[0]["prompt_token_ids"] == [1, 2, 3]

tests/model_executor/stage_input_processors/test_glm_image.py

@@ -32,11 +32,6 @@ def _source_output(token_ids: list[int], mm_output: dict | None = None):
     )
 
 
-def _stage_with_outputs(outputs):
-    """Create a stage list entry with engine_outputs."""
-    return SimpleNamespace(engine_outputs=outputs)
-
-
 # =============================================================================
 # Tests for _has_source_image
 # =============================================================================
@@ -285,11 +280,11 @@ class TestAr2Diffusion:
         """Test basic text-to-image pipeline: AR -> Diffusion."""
         # 1024x1024 t2i: small(256) + large(1024) + EOS
         token_ids = list(range(256)) + list(range(1024)) + [16385]
-        stage_list = [_stage_with_outputs([_source_output(token_ids)])]
+        source_outputs = [_source_output(token_ids)]
 
         prompt = {"prompt": "a cat", "mm_processor_kwargs": {"target_h": 1024, "target_w": 1024}}
 
-        result = ar2diffusion(stage_list, [0], prompt=[prompt])
+        result = ar2diffusion(source_outputs, prompt=[prompt])
         assert len(result) == 1
         assert result[0]["prompt"] == "a cat"
         assert result[0]["height"] == 1024
@@ -300,7 +295,7 @@ class TestAr2Diffusion:
         """Test image-to-image with prior_token_image_ids from AR model."""
         token_ids = list(range(1024)) + [16385]
         mm_output = {"ids": {"prior_image": torch.tensor([1, 2, 3])}}
-        stage_list = [_stage_with_outputs([_source_output(token_ids, mm_output)])]
+        source_outputs = [_source_output(token_ids, mm_output)]
 
         from PIL import Image
 
@@ -311,14 +306,14 @@ class TestAr2Diffusion:
             "multi_modal_data": {"image": img},
         }
 
-        result = ar2diffusion(stage_list, [0], prompt=[prompt])
+        result = ar2diffusion(source_outputs, prompt=[prompt])
         assert len(result) == 1
         assert result[0]["extra"]["prior_token_image_ids"] is not None
 
     def test_i2i_detected_via_modalities(self):
         """Test i2i mode detected via modalities field."""
         token_ids = list(range(1024)) + [16385]
-        stage_list = [_stage_with_outputs([_source_output(token_ids)])]
+        source_outputs = [_source_output(token_ids)]
 
         prompt = {
             "prompt": "edit this",
@@ -326,35 +321,26 @@ class TestAr2Diffusion:
             "modalities": ["img2img"],
         }
 
-        result = ar2diffusion(stage_list, [0], prompt=[prompt])
+        result = ar2diffusion(source_outputs, prompt=[prompt])
         assert len(result) == 1
 
-    def test_empty_engine_input_source_raises(self):
-        with pytest.raises(ValueError, match="engine_input_source cannot be empty"):
-            ar2diffusion([], [], prompt={})
-
-    def test_invalid_stage_id_raises(self):
-        with pytest.raises(IndexError, match="Invalid stage_id"):
-            ar2diffusion([_stage_with_outputs(None)], [5], prompt={})
-
-    def test_no_outputs_raises(self):
-        with pytest.raises(RuntimeError, match="has no outputs yet"):
-            ar2diffusion([SimpleNamespace(engine_outputs=None)], [0], prompt={})
+    def test_empty_source_outputs_returns_empty_list(self):
+        assert ar2diffusion([], prompt={}) == []
 
     def test_default_dimensions(self):
         """When no height/width in prompt, defaults to 1024x1024."""
         token_ids = list(range(256)) + list(range(1024)) + [16385]
-        stage_list = [_stage_with_outputs([_source_output(token_ids)])]
+        source_outputs = [_source_output(token_ids)]
 
         prompt = {"prompt": "test"}
-        result = ar2diffusion(stage_list, [0], prompt=[prompt])
+        result = ar2diffusion(source_outputs, prompt=[prompt])
         assert result[0]["height"] == 1024
         assert result[0]["width"] == 1024
 
     def test_requires_multimodal_data_with_pil_image(self):
         """Test that pil_image is included when requires_multimodal_data=True."""
         token_ids = list(range(256)) + list(range(1024)) + [16385]
-        stage_list = [_stage_with_outputs([_source_output(token_ids)])]
+        source_outputs = [_source_output(token_ids)]
 
         from PIL import Image
 
@@ -364,13 +350,13 @@ class TestAr2Diffusion:
             "multi_modal_data": {"image": img},
         }
 
-        result = ar2diffusion(stage_list, [0], prompt=[prompt], requires_multimodal_data=True)
+        result = ar2diffusion(source_outputs, prompt=[prompt], requires_multimodal_data=True)
         assert result[0]["pil_image"] is img
 
     def test_extra_params_passed_through(self):
         """Test that seed, num_inference_steps, guidance_scale, negative_prompt are passed."""
         token_ids = list(range(256)) + list(range(1024)) + [16385]
-        stage_list = [_stage_with_outputs([_source_output(token_ids)])]
+        source_outputs = [_source_output(token_ids)]
 
         prompt = {
             "prompt": "test",
@@ -380,7 +366,7 @@ class TestAr2Diffusion:
             "negative_prompt": "blurry",
         }
 
-        result = ar2diffusion(stage_list, [0], prompt=[prompt])
+        result = ar2diffusion(source_outputs, prompt=[prompt])
         assert result[0]["seed"] == 42
         assert result[0]["num_inference_steps"] == 50
         assert result[0]["guidance_scale"] == 7.5
@@ -390,14 +376,14 @@ class TestAr2Diffusion:
         """Test processing multiple requests in a batch."""
         tokens1 = list(range(256)) + list(range(1024)) + [16385]
         tokens2 = list(range(256)) + list(range(1024)) + [16385]
-        stage_list = [_stage_with_outputs([_source_output(tokens1), _source_output(tokens2)])]
+        source_outputs = [_source_output(tokens1), _source_output(tokens2)]
 
         prompts = [
             {"prompt": "first", "mm_processor_kwargs": {"target_h": 1024, "target_w": 1024}},
             {"prompt": "second", "mm_processor_kwargs": {"target_h": 512, "target_w": 512}},
         ]
 
-        result = ar2diffusion(stage_list, [0], prompt=prompts)
+        result = ar2diffusion(source_outputs, prompt=prompts)
         assert len(result) == 2
         assert result[0]["prompt"] == "first"
         assert result[1]["prompt"] == "second"

tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py

@@ -17,12 +17,11 @@ from vllm_omni.model_executor.stage_input_processors.mimo_audio import (
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _make_stage_list(codec_codes: torch.Tensor, request_id: str = "req-0"):
-    """Build a minimal stage_list[0] with a single talker output carrying codec_codes."""
+def _make_source_outputs(codec_codes: torch.Tensor, request_id: str = "req-0"):
     output = SimpleNamespace(multimodal_output={"codes": {"audio": codec_codes}})
+    """Build minimal source_outputs with one talker output carrying codec codes."""
     talker_output = SimpleNamespace(outputs=[output], request_id=request_id)
-    stage0 = SimpleNamespace(engine_outputs=[talker_output])
-    return [stage0]
+    return [talker_output]
 
 
 def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
@@ -33,7 +32,7 @@ def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
     frames = (MAX_CODE2WAV_TOKENS // 36) + 100
     codec_codes = torch.ones(frames, 1, 8, 4, dtype=torch.long)
 
-    stage_list = _make_stage_list(codec_codes, request_id="req-long")
+    source_outputs = _make_source_outputs(codec_codes, request_id="req-long")
 
     # Attach caplog's handler directly to the module logger so the warning is
     # captured regardless of propagation (vllm's logger configuration can
@@ -43,7 +42,7 @@ def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
     prev_level = target_logger.level
     target_logger.setLevel(logging.WARNING)
     try:
-        prompts = llm2code2wav(stage_list, engine_input_source=[0])
+        prompts = llm2code2wav(source_outputs)
     finally:
         target_logger.removeHandler(caplog.handler)
         target_logger.setLevel(prev_level)
@@ -58,9 +57,9 @@ def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog):
 def test_llm2code2wav_short_sequence_unchanged():
     """Short codec sequences are returned without truncation."""
     codec_codes = torch.ones(4, 1, 8, 4, dtype=torch.long)
-    stage_list = _make_stage_list(codec_codes, request_id="req-short")
+    source_outputs = _make_source_outputs(codec_codes, request_id="req-short")
 
-    prompts = llm2code2wav(stage_list, engine_input_source=[0])
+    prompts = llm2code2wav(source_outputs)
 
     assert len(prompts) == 1
     # 4 frames + 1 pad row, flattened col-major → well below the cap

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -342,7 +342,7 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
     )
 
-    prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])
+    prompts = talker2code2wav(stage.engine_outputs)
 
     assert len(prompts) == 1
     prompt = prompts[0]
@@ -388,7 +388,7 @@ def test_non_async_processor_filters_out_of_range_codec_values():
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
     )
 
-    prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])
+    prompts = talker2code2wav(stage.engine_outputs)
 
     assert len(prompts) == 1
     prompt = prompts[0]

tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py

@@ -38,7 +38,7 @@ class TestGenerator2Tokenizer:
         audio = torch.tensor([10, 20, 30, 40])
         stage = _make_stage([_make_output([audio])])
 
-        result = generator2tokenizer([stage], engine_input_source=[0])
+        result = generator2tokenizer(stage.engine_outputs)
 
         assert len(result) == 1
         assert result[0]["prompt_token_ids"] == [10, 20, 30, 40]
@@ -49,27 +49,15 @@ class TestGenerator2Tokenizer:
         audio = torch.tensor([[1, 2, 3], [4, 5, 6]])  # (2, 3)
         stage = _make_stage([_make_output([audio])])
 
-        result = generator2tokenizer([stage], engine_input_source=[0])
+        result = generator2tokenizer(stage.engine_outputs)
 
         assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6]
 
-    def test_empty_engine_input_source_raises(self):
-        """Empty engine_input_source raises ValueError."""
-        stage = _make_stage([_make_output([torch.tensor([1])])])
-        with pytest.raises(ValueError, match="cannot be empty"):
-            generator2tokenizer([stage], engine_input_source=[])
-
-    def test_invalid_stage_id_raises(self):
-        """Stage id beyond stage_list length raises IndexError."""
-        stage = _make_stage([_make_output([torch.tensor([1])])])
-        with pytest.raises(IndexError, match="Invalid stage_id"):
-            generator2tokenizer([stage], engine_input_source=[5])
-
     def test_no_outputs_yet_raises(self):
         """Stage with engine_outputs=None raises RuntimeError."""
         stage = _make_stage(engine_outputs=None)
-        with pytest.raises(RuntimeError, match="no outputs yet"):
-            generator2tokenizer([stage], engine_input_source=[0])
+        with pytest.raises(TypeError):
+            generator2tokenizer(stage.engine_outputs)
 
 
 # ---- Helpers for generator2tokenizer_async_chunk (streaming) ----

tests/test_config_factory.py

@@ -920,6 +920,27 @@ class TestDeployConfigLoading:
         assert s0.yaml_engine_args["engine_output_type"] == "latent"
         assert s0.yaml_extras["default_sampling_params"]["detokenize"] is True
 
+    def test_merge_pipeline_deploy_preserves_num_replicas(self, tmp_path):
+        from pathlib import Path
+
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
+        pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
+        base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        if not base.exists():
+            pytest.skip("Deploy config not found")
+
+        overlay = tmp_path / "multi_replicas.yaml"
+        overlay.write_text(f'base_config: {base}\nstages:\n  - stage_id: 1\n    devices: "1,2"\n    num_replicas: 2\n')
+
+        deploy = load_deploy_config(overlay)
+        assert deploy.stages[1].num_replicas == 2
+
+        stages = merge_pipeline_deploy(pipeline, deploy)
+        assert stages[1].yaml_runtime["devices"] == "1,2"
+        assert stages[1].yaml_runtime["num_replicas"] == 2
+
 
 class TestQwen3OmniPipeline:
     def test_registered(self):

tests/worker/test_omni_connector_mixin.py

@@ -354,15 +354,15 @@ class TestFinishedLoadReqsDrain(unittest.TestCase):
 
 
 class TestLoadCustomFuncSelection(unittest.TestCase):
-    def test_skips_legacy_stage_list_processors_for_full_payload_mode(self):
-        legacy_paths = [
+    def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self):
+        incompatible_paths = [
             "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav",
             "vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit",
             "vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow",
             "vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion",
         ]
 
-        for func_path in legacy_paths:
+        for func_path in incompatible_paths:
             selected_path, func = MixinHost._load_custom_func(
                 SimpleNamespace(
                     async_chunk=False,

vllm_omni/config/stage_config.py

@@ -407,6 +407,7 @@ class StageDeployConfig:
     # Stage identity and Omni runtime placement.
     stage_id: int
     devices: str | None = None
+    num_replicas: int = 1
 
     # Inter-stage connector wiring and request defaults.
     output_connectors: dict[str, str] | None = None
@@ -476,6 +477,7 @@ _STAGE_NON_ENGINE_KEYS = frozenset(
     {
         "stage_id",
         "devices",
+        "num_replicas",
         "output_connectors",
         "input_connectors",
         "default_sampling_params",
@@ -490,13 +492,20 @@ _STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name n
 def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
     """Parse a single stage entry from deploy YAML into StageDeployConfig."""
     if "engine_args" in stage_data:
+        runtime_cfg = dict(stage_data.get("runtime", {}))
         engine_args = dict(stage_data["engine_args"])
         devices = stage_data.get("runtime", {}).get("devices", stage_data.get("devices"))
+        num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1))
     else:
         engine_args = {k: v for k, v in stage_data.items() if k not in _STAGE_NON_ENGINE_KEYS and k != "stage_id"}
         devices = stage_data.get("devices")
+        num_replicas = stage_data.get("num_replicas", stage_data.get("runtime", {}).get("num_replicas", 1))
 
-    kwargs: dict[str, Any] = {"stage_id": stage_data["stage_id"], "devices": devices}
+    kwargs: dict[str, Any] = {
+        "stage_id": stage_data["stage_id"],
+        "devices": devices,
+        "num_replicas": int(num_replicas),
+    }
     for name, f in _STAGE_DEPLOY_FIELDS.items():
         if name in engine_args:
             kwargs[name] = engine_args.pop(name)
@@ -632,7 +641,11 @@ def _extract_platform_overrides(ps: dict[str, Any]) -> tuple[dict[str, Any], str
     the flat layout. ``devices`` is ``None`` when no override is set.
     """
     if "engine_args" in ps:
-        return dict(ps["engine_args"]), ps.get("runtime", {}).get("devices")
+        overrides = dict(ps["engine_args"])
+        runtime_cfg = ps.get("runtime", {})
+        if "num_replicas" in runtime_cfg:
+            overrides["num_replicas"] = runtime_cfg["num_replicas"]
+        return overrides, runtime_cfg.get("devices")
     overrides = {k: v for k, v in ps.items() if k not in ("stage_id", "devices")}
     return overrides, ps.get("devices")
 
@@ -837,8 +850,10 @@ def merge_pipeline_deploy(
             engine_args["async_scheduling"] = sched_cls is OmniARAsyncScheduler
         extras = _build_extras(ps, ds)
         runtime: dict[str, Any] = {"process": True}
-        if ds is not None and ds.devices is not None:
-            runtime["devices"] = ds.devices
+        if ds is not None:
+            if ds.devices is not None:
+                runtime["devices"] = ds.devices
+            runtime["num_replicas"] = ds.num_replicas
 
         result.append(
             StageConfig(

vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml

@@ -0,0 +1,101 @@
+# Qwen3-Omni-MoE multi-replica deploy, verified on 3x H20-96G
+# (stage 0 on cuda:0, stages 1+2 scaled out across cuda:1,2).
+#
+# Fields omitted from a stage fall back to StageDeployConfig defaults (see
+# vllm_omni/config/stage_config.py). Notable implicit defaults for this
+# layout:
+#   * Stages 0/1 (thinker, talker) do not set max_num_batched_tokens:
+#     chat-sized prefill fits in the 32768 default.
+#   * Stage 1 (talker) does not set enforce_eager: cudagraph runs by
+#     default (false).
+#   * Stage 2 (code2wav) sets enforce_eager=true because its generation
+#     loop is cudagraph-incompatible.
+#   * Stages 1/2 use num_replicas=2 and share devices "1,2". Replica
+#     routing is handled by StagePool with request-level affinity.
+async_chunk: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      codec_chunk_frames: 25
+      codec_left_context_frames: 25
+
+stages:
+  - stage_id: 0
+    gpu_memory_utilization: 0.9
+    devices: "0"
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    gpu_memory_utilization: 0.6
+    devices: "1,2"
+    num_replicas: 2
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      repetition_penalty: 1.05
+
+  - stage_id: 2
+    gpu_memory_utilization: 0.1
+    max_num_seqs: 1
+    enforce_eager: true
+    async_scheduling: false
+    max_num_batched_tokens: 51200
+    devices: "1,2"
+    num_replicas: 2
+    input_connectors:
+      from_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.1
+
+platforms:
+  npu:
+    stages:
+      - stage_id: 0
+        gpu_memory_utilization: 0.6
+        tensor_parallel_size: 2
+        devices: "0,1"
+      - stage_id: 1
+        gpu_memory_utilization: 0.6
+        enforce_eager: true
+        devices: "2"
+      - stage_id: 2
+        gpu_memory_utilization: 0.3
+        devices: "2"
+
+  rocm:
+    stages:
+      - stage_id: 0
+        enforce_eager: true
+
+  xpu:
+    stages:
+      - stage_id: 0
+        tensor_parallel_size: 4
+        enforce_eager: true
+        max_cudagraph_capture_size: 0
+        devices: "0,1,2,3"
+      - stage_id: 1
+        enforce_eager: true
+        max_cudagraph_capture_size: 0
+        devices: "4"
+      - stage_id: 2
+        gpu_memory_utilization: 0.3
+        max_cudagraph_capture_size: 0
+        devices: "4"

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -35,6 +35,7 @@ class InlineStageDiffusionClient:
     """Runs DiffusionEngine in a thread executor inside the Orchestrator."""
 
     stage_type: str = "diffusion"
+    replica_id: int = 0
 
     def __init__(
         self,
@@ -46,9 +47,11 @@ class InlineStageDiffusionClient:
         self.model = model
         self.od_config = od_config
         self.stage_id = metadata.stage_id
+        self.replica_id = metadata.replica_id
         self.final_output = metadata.final_output
         self.final_output_type = metadata.final_output_type
         self.default_sampling_params = metadata.default_sampling_params
+        self.requires_multimodal_data = metadata.requires_multimodal_data
         self.custom_process_input_func = metadata.custom_process_input_func
         self.engine_input_source = metadata.engine_input_source
         self.batch_size = batch_size
@@ -62,8 +65,9 @@ class InlineStageDiffusionClient:
         self._shutting_down = False
 
         logger.info(
-            "[InlineStageDiffusionClient] Stage-%s initialized inline (batch_size=%d)",
+            "[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (batch_size=%d)",
             self.stage_id,
+            self.replica_id,
             self.batch_size,
         )
 
@@ -82,6 +86,12 @@ class InlineStageDiffusionClient:
         sampling_params: OmniDiffusionSamplingParams,
         kv_sender_info: dict[int, dict[str, Any]] | None = None,
     ) -> None:
+        logger.info(
+            "[InlineStageDiffusionClient] stage-%s [rep-%s] add request: %s",
+            self.stage_id,
+            self.replica_id,
+            request_id,
+        )
         task = asyncio.create_task(
             self._dispatch_request(
                 request_id,
@@ -134,6 +144,13 @@ class InlineStageDiffusionClient:
         sampling_params: OmniDiffusionSamplingParams,
         kv_sender_info: dict[int, dict[str, Any]] | None = None,
     ) -> None:
+        logger.info(
+            "[InlineStageDiffusionClient] stage-%s [rep-%s] add batch request: %s (%d prompts)",
+            self.stage_id,
+            self.replica_id,
+            request_id,
+            len(prompts),
+        )
         task = asyncio.create_task(
             self._dispatch_batch(
                 request_id,
@@ -252,7 +269,7 @@ class InlineStageDiffusionClient:
             is_start = args[0] if args else True
             profile_prefix = args[1] if len(args) > 1 else None
             if is_start and profile_prefix is None:
-                profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}"
+                profile_prefix = f"stage_{self.stage_id}_rep_{self.replica_id}_diffusion_{int(time.time())}"
             return await loop.run_in_executor(
                 self._executor,
                 self._engine.profile,

vllm_omni/diffusion/stage_diffusion_client.py

@@ -67,6 +67,7 @@ class StageDiffusionClient:
     """
 
     stage_type: str = "diffusion"
+    replica_id: int = 0
 
     def __init__(
         self,
@@ -112,11 +113,13 @@ class StageDiffusionClient:
         batch_size: int,
     ) -> None:
         self.stage_id = metadata.stage_id
+        self.replica_id = metadata.replica_id
         self.final_output = metadata.final_output
         self.final_output_type = metadata.final_output_type
         self.default_sampling_params = metadata.default_sampling_params
-        self.custom_process_input_func = metadata.custom_process_input_func
-        self.engine_input_source = metadata.engine_input_source
+        self.requires_multimodal_data = getattr(metadata, "requires_multimodal_data", False)
+        self.custom_process_input_func = getattr(metadata, "custom_process_input_func", None)
+        self.engine_input_source = getattr(metadata, "engine_input_source", [])
         self._proc = proc
         self._owns_process = proc is not None
 
@@ -142,8 +145,9 @@ class StageDiffusionClient:
         self._start_proc_monitor()
 
         logger.info(
-            "[StageDiffusionClient] Stage-%s initialized (owns_process=%s, batch_size=%d)",
+            "[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s, batch_size=%d)",
             self.stage_id,
+            self.replica_id,
             self._owns_process,
             batch_size,
         )
@@ -172,8 +176,9 @@ class StageDiffusionClient:
                 return
             client._engine_dead = True
             logger.error(
-                "[StageDiffusionClient] Stage-%s StageDiffusionProc died unexpectedly (exit code %s).",
+                "[StageDiffusionClient] stage-%s [rep-%s] StageDiffusionProc died unexpectedly (exit code %s).",
                 client.stage_id,
+                client.replica_id,
                 proc.exitcode,
             )
 
@@ -195,8 +200,9 @@ class StageDiffusionClient:
             if raw == StageDiffusionProc.DIFFUSION_PROC_DEAD:
                 self._engine_dead = True
                 logger.error(
-                    "[StageDiffusionClient] Stage-%s received DIFFUSION_PROC_DEAD sentinel from subprocess.",
+                    "[StageDiffusionClient] stage-%s [rep-%s] received DIFFUSION_PROC_DEAD sentinel from subprocess.",
                     self.stage_id,
+                    self.replica_id,
                 )
                 break
 
@@ -212,8 +218,9 @@ class StageDiffusionClient:
                 rpc_id = msg.get("rpc_id")
                 error_msg = msg.get("error")
                 logger.error(
-                    "[StageDiffusionClient] Stage-%s subprocess error for %s: %s",
+                    "[StageDiffusionClient] stage-%s [rep-%s] subprocess error for %s: %s",
                     self.stage_id,
+                    self.replica_id,
                     rpc_id or req_id,
                     error_msg,
                 )
@@ -292,6 +299,12 @@ class StageDiffusionClient:
     ) -> None:
         if self._engine_dead:
             raise EngineDeadError()
+        logger.info(
+            "[StageDiffusionClient] stage-%s [rep-%s] add request: %s",
+            self.stage_id,
+            self.replica_id,
+            request_id,
+        )
         self._request_socket.send(
             self._encoder.encode(
                 {
@@ -321,6 +334,13 @@ class StageDiffusionClient:
         """
         if self._engine_dead:
             raise EngineDeadError()
+        logger.info(
+            "[StageDiffusionClient] stage-%s [rep-%s] add batch request: %s (%d prompts)",
+            self.stage_id,
+            self.replica_id,
+            request_id,
+            len(prompts),
+        )
         task = asyncio.create_task(
             self._run_batch(
                 request_id,
@@ -353,8 +373,9 @@ class StageDiffusionClient:
             )
         except Exception as e:
             logger.exception(
-                "[StageDiffusionClient] Stage-%s batch req=%s failed: %s",
+                "[StageDiffusionClient] stage-%s [rep-%s] batch req=%s failed: %s",
                 self.stage_id,
+                self.replica_id,
                 request_id,
                 e,
             )
@@ -414,7 +435,7 @@ class StageDiffusionClient:
             is_start = args_list[0] if args_list else True
             profile_prefix = args_list[1] if len(args_list) > 1 else None
             if is_start and profile_prefix is None:
-                profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}"
+                profile_prefix = f"stage_{self.stage_id}_rep_{self.replica_id}_diffusion_{int(time.time())}"
                 if len(args_list) > 1:
                     args_list[1] = profile_prefix
                 else:

vllm_omni/engine/__init__.py

@@ -76,6 +76,42 @@ class OmniEngineCoreRequest(EngineCoreRequest):
     # Optional additional information dictionary (serialized)
     additional_information: AdditionalInformationPayload | None = None
 
+    @classmethod
+    def from_request(
+        cls,
+        request: EngineCoreRequest,
+        *,
+        prompt_embeds: torch.Tensor | None = None,
+        additional_information: AdditionalInformationPayload | None = None,
+    ) -> "OmniEngineCoreRequest":
+        """Clone an EngineCoreRequest into an OmniEngineCoreRequest with optional payload overrides."""
+
+        if prompt_embeds is None:
+            prompt_embeds = request.prompt_embeds
+        if additional_information is None:
+            additional_information = getattr(request, "additional_information", None)
+
+        return cls(
+            request_id=request.request_id,
+            prompt_token_ids=request.prompt_token_ids,
+            mm_features=request.mm_features,
+            sampling_params=request.sampling_params,
+            pooling_params=request.pooling_params,
+            arrival_time=request.arrival_time,
+            lora_request=request.lora_request,
+            cache_salt=request.cache_salt,
+            data_parallel_rank=request.data_parallel_rank,
+            prompt_embeds=prompt_embeds,
+            client_index=request.client_index,
+            current_wave=request.current_wave,
+            priority=request.priority,
+            trace_headers=request.trace_headers,
+            resumable=request.resumable,
+            external_req_id=request.external_req_id,
+            reasoning_ended=request.reasoning_ended,
+            additional_information=additional_information,
+        )
+
 
 class OmniEngineCoreOutput(EngineCoreOutput):
     pooling_output: dict[str, torch.Tensor] | None = None

vllm_omni/engine/cfg_companion_tracker.py

@@ -10,8 +10,6 @@ from __future__ import annotations
 import logging
 from typing import Any
 
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
 logger = logging.getLogger(__name__)
 
 
@@ -31,6 +29,10 @@ class CfgCompanionTracker:
     def has_companions(self, parent_id: str) -> bool:
         return parent_id in self._companion_map
 
+    def get_parent_id(self, req_id: str) -> str | None:
+        """Return the parent request id for a companion, or None."""
+        return self._companion_to_parent.get(req_id)
+
     def all_companions_done(self, parent_id: str) -> bool:
         role_map = self._companion_map.get(parent_id, {})
         done_set = self._done.get(parent_id, set())
@@ -50,21 +52,6 @@ class CfgCompanionTracker:
         self._companion_ids.add(companion_id)
         self._companion_to_parent[companion_id] = parent_id
 
-    def attach_cfg_request_ids(self, parent_id: str, sampling_params: Any) -> Any:
-        cfg_ids = self.get_companion_request_ids(parent_id)
-        if not cfg_ids:
-            return sampling_params
-
-        if isinstance(sampling_params, OmniDiffusionSamplingParams):
-            sampling_params = sampling_params.clone()
-            sampling_params.cfg_kv_request_ids = cfg_ids
-            logger.info(
-                "Attaching cfg_kv_request_ids=%s to request %s",
-                cfg_ids,
-                parent_id,
-            )
-        return sampling_params
-
     def on_companion_completed(self, companion_id: str) -> str | None:
         """Mark done. Returns parent_id only if parent is pending and all companions finished."""
         parent_id = self._companion_to_parent.get(companion_id)

vllm_omni/engine/output_processor.py

@@ -359,6 +359,24 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             self.parent_requests[parent_req.request_id] = parent_req
         self.external_req_ids[req_state.external_req_id].append(request_id)
 
+    def remove_request(self, request_id: str) -> None:
+        """Rollback one previously registered request if it was never submitted."""
+        req_state = self.request_states.pop(request_id, None)
+        if req_state is None:
+            return
+
+        external_req_id = getattr(req_state, "external_req_id", None)
+        if external_req_id is not None:
+            request_ids = self.external_req_ids.get(external_req_id)
+            if request_ids is not None:
+                self.external_req_ids[external_req_id] = [rid for rid in request_ids if rid != request_id]
+                if not self.external_req_ids[external_req_id]:
+                    self.external_req_ids.pop(external_req_id, None)
+
+        parent_req = getattr(req_state, "parent_req", None)
+        if parent_req is not None:
+            self.parent_requests.pop(parent_req.request_id, None)
+
     def process_outputs(
         self,
         engine_core_outputs: list[EngineCoreOutput],

vllm_omni/engine/stage_engine_core_client.py

@@ -6,6 +6,7 @@ Directly inherits from vLLM's AsyncMPClient to reuse EngineCore architecture.
 
 from __future__ import annotations
 
+import inspect
 import multiprocessing.connection
 import socket
 import threading
@@ -34,6 +35,27 @@ logger = init_logger(__name__)
 SHUTDOWN_TIMEOUT_S = 5
 
 
+def _default_process_engine_inputs(
+    source_outputs: list[Any],
+    prompt: Any,
+    requires_multimodal_data: bool,
+) -> list[OmniTokensPrompt]:
+    from vllm_omni.inputs.data import OmniTokensPrompt
+
+    if not isinstance(prompt, list):
+        prompt = [prompt]
+
+    mm_data = {so.request_id: p.get("multi_modal_data") for so, p in zip(source_outputs, prompt)}
+
+    return [
+        OmniTokensPrompt(
+            prompt_token_ids=so.outputs[0].token_ids,
+            multi_modal_data=(mm_data[so.request_id] if requires_multimodal_data else None),
+        )
+        for so in source_outputs
+    ]
+
+
 class StageEngineCoreClientBase:
     """Shared stage-aware behavior for async EngineCore clients.
 
@@ -51,6 +73,8 @@ class StageEngineCoreClientBase:
     ``engine_manager`` / ``coordinator`` pair created elsewhere.
     """
 
+    replica_id: int = 0
+
     @staticmethod
     def make_async_mp_client(
         vllm_config: Any,
@@ -109,8 +133,10 @@ class StageEngineCoreClientBase:
         manage the process lifecycle on shutdown.
         """
         # -------- Stage metadata (public fields used at runtime) --------
+        self.replica_id = 0
         if metadata is not None:
             self.stage_id = metadata.stage_id
+            self.replica_id = getattr(metadata, "replica_id", 0)
             self.stage_type = metadata.stage_type
             self.engine_output_type = metadata.engine_output_type
             self.is_comprehension = metadata.is_comprehension
@@ -132,9 +158,10 @@ class StageEngineCoreClientBase:
 
         client_name = self.__class__.__name__
         logger.info(
-            "[%s] Stage-%s initializing EngineCore",
+            "[%s] stage-%s [rep-%s] initializing EngineCore",
             client_name,
             self.stage_id,
+            self.replica_id,
         )
         try:
             super().__init__(
@@ -151,17 +178,19 @@ class StageEngineCoreClientBase:
                 self.resources.coordinator = coordinator
         except Exception:
             logger.exception(
-                "[%s] Stage-%s EngineCore init failed",
+                "[%s] stage-%s [rep-%s] EngineCore init failed",
                 client_name,
                 self.stage_id,
+                self.replica_id,
             )
             try:
                 self.shutdown()
             except Exception as shutdown_error:
                 logger.warning(
-                    "[%s] Stage-%s cleanup after init failure failed: %s",
+                    "[%s] stage-%s [rep-%s] cleanup after init failure failed: %s",
                     client_name,
                     self.stage_id,
+                    self.replica_id,
                     shutdown_error,
                 )
             raise
@@ -172,9 +201,10 @@ class StageEngineCoreClientBase:
             self._start_proc_monitor()
 
         logger.info(
-            "[%s] Stage-%s EngineCore running",
+            "[%s] stage-%s [rep-%s] EngineCore running",
             client_name,
             self.stage_id,
+            self.replica_id,
         )
 
     def _start_proc_monitor(self) -> None:
@@ -188,6 +218,7 @@ class StageEngineCoreClientBase:
         proc = self._proc
         resources_ref = weakref.ref(self.resources)
         stage_id = self.stage_id
+        replica_id = self.replica_id
 
         def _monitor() -> None:
             try:
@@ -199,8 +230,9 @@ class StageEngineCoreClientBase:
                 return
             resources.engine_dead = True
             logger.error(
-                "[StageEngineCoreClient] Stage-%s subprocess died unexpectedly (exit code %s).",
+                "[StageEngineCoreClient] stage-%s [rep-%s] subprocess died unexpectedly (exit code %s).",
                 stage_id,
+                replica_id,
                 proc.exitcode,
             )
 
@@ -228,9 +260,10 @@ class StageEngineCoreClientBase:
     async def add_request_async(self, request: EngineCoreRequest) -> None:
         """Add request to the stage engine core."""
         logger.info(
-            "[%s] Stage-%s adding request: %s",
+            "[%s] stage-%s [rep-%s] add request: %s",
             self.__class__.__name__,
             self.stage_id,
+            self.replica_id,
             request.request_id,
         )
         await super().add_request_async(request)
@@ -311,9 +344,10 @@ class StageEngineCoreClientBase:
                 sender_port = int(base_port) + KV_TRANSFER_PORT_OFFSET + int(from_stage)
             except (TypeError, ValueError):
                 logger.warning(
-                    "[StageEngineCoreClient] Stage-%s could not resolve sender_zmq_port "
+                    "[StageEngineCoreClient] stage-%s [rep-%s] could not resolve sender_zmq_port "
                     "from base_port=%s and from_stage=%s",
                     self.stage_id,
+                    self.replica_id,
                     base_port,
                     from_stage,
                 )
@@ -358,48 +392,33 @@ class StageEngineCoreClientBase:
 
     def process_engine_inputs(
         self,
-        stage_list: list[Any],
-        prompt: OmniTokensPrompt | list[OmniTokensPrompt] | None = None,
+        source_outputs: list[Any],
+        prompt: Any = None,
         streaming_context: Any | None = None,
     ) -> list[OmniTokensPrompt]:
-        """Process inputs from upstream stages."""
-        from vllm_omni.inputs.data import OmniTokensPrompt
+        """Process inputs from upstream stages.
 
+        Transition planning is expressed in terms of the upstream outputs
+        and the original prompt.
+        """
         if self.custom_process_input_func is not None:
-            # Keep legacy arg call for non-streaming processors.
-            if bool(getattr(streaming_context, "enabled", False)):
+            signature = inspect.signature(self.custom_process_input_func)
+            if len(signature.parameters) >= 4:
                 return self.custom_process_input_func(
-                    stage_list,
-                    self.engine_input_source,
+                    source_outputs,
                     prompt,
                     self.requires_multimodal_data,
                     streaming_context,
                 )
             return self.custom_process_input_func(
-                stage_list,
-                self.engine_input_source,
+                source_outputs,
                 prompt,
                 self.requires_multimodal_data,
             )
 
         if not self.engine_input_source:
             raise ValueError(f"engine_input_source empty for stage {self.stage_id}")
-
-        source_id = self.engine_input_source[0]
-        source_outputs = stage_list[source_id].engine_outputs
-
-        if not isinstance(prompt, list):
-            prompt = [prompt]
-
-        mm_data = {so.request_id: p.get("multi_modal_data") for so, p in zip(source_outputs, prompt)}
-
-        return [
-            OmniTokensPrompt(
-                prompt_token_ids=so.outputs[0].cumulative_token_ids,
-                multi_modal_data=(mm_data[so.request_id] if self.requires_multimodal_data else None),
-            )
-            for so in source_outputs
-        ]
+        return _default_process_engine_inputs(source_outputs, prompt, self.requires_multimodal_data)
 
     async def collective_rpc_async(
         self,

vllm_omni/engine/stage_engine_startup.py

@@ -29,6 +29,8 @@ from vllm.v1.executor import Executor
 
 logger = init_logger(__name__)
 
+StageRoute = tuple[int, int]
+
 # Poll period (ms) used by the registration/handshake loop.
 _POLL_PERIOD_MS = 5_000
 # Default timeout (s) for a stage to send READY.
@@ -107,28 +109,33 @@ class OmniMasterServer:
         master_address: str,
         master_port: int,
         stage_ids: list[int],
+        stage_replica_counts: dict[int, int] | None = None,
     ) -> None:
         self._address = master_address
         self._port = master_port
-        self._allocations: dict[int, StageAllocation] = {}
-        self._stage_configs: dict[int, Any] = {}
-        self._stage_coordinator_addresses: dict[int, StageCoordinatorAddresses] = {}
-        self._stage_config_events: dict[int, threading.Event] = {}
+        self._stage_routes: dict[StageRoute, StageAllocation] = {}
+        self._stage_configs: dict[StageRoute, Any] = {}
+        self._stage_coordinator_addresses: dict[StageRoute, StageCoordinatorAddresses] = {}
+        self._stage_config_events: dict[StageRoute, threading.Event] = {}
         self._thread: threading.Thread | None = None
         self._stop_event = threading.Event()
+        stage_replica_counts = dict(stage_replica_counts or {})
 
         for sid in stage_ids:
-            self._stage_config_events[sid] = threading.Event()
-            self._stage_coordinator_addresses[sid] = StageCoordinatorAddresses()
-            hs_port, inp_port, out_port = get_open_ports_list(count=3)
-            self._allocations[sid] = StageAllocation(
-                handshake_bind_address=f"tcp://{master_address}:{hs_port}",
-                handshake_connect_address=f"tcp://{master_address}:{hs_port}",
-                input_bind_address=f"tcp://{master_address}:{inp_port}",
-                input_connect_address=f"tcp://{master_address}:{inp_port}",
-                output_bind_address=f"tcp://{master_address}:{out_port}",
-                output_connect_address=f"tcp://{master_address}:{out_port}",
-            )
+            replica_count = max(1, int(stage_replica_counts.get(sid, 1)))
+            for replica_id in range(replica_count):
+                route = (sid, replica_id)
+                self._stage_config_events[route] = threading.Event()
+                self._stage_coordinator_addresses[route] = StageCoordinatorAddresses()
+                hs_port, inp_port, out_port = get_open_ports_list(count=3)
+                self._stage_routes[route] = StageAllocation(
+                    handshake_bind_address=f"tcp://{master_address}:{hs_port}",
+                    handshake_connect_address=f"tcp://{master_address}:{hs_port}",
+                    input_bind_address=f"tcp://{master_address}:{inp_port}",
+                    input_connect_address=f"tcp://{master_address}:{inp_port}",
+                    output_bind_address=f"tcp://{master_address}:{out_port}",
+                    output_connect_address=f"tcp://{master_address}:{out_port}",
+                )
 
         logger.info(
             "[OmniMasterServer] Pre-allocated addresses for stages %s (master=%s:%d)",
@@ -150,71 +157,78 @@ class OmniMasterServer:
         """Return the registration port exposed to stage launchers."""
         return self._port
 
-    def get_allocation(self, stage_id: int) -> StageAllocation:
+    def get_allocation(self, stage_id: int, replica_id: int = 0) -> StageAllocation:
         """Return the full address allocation for *stage_id*."""
-        return self._allocations[stage_id]
+        return self._stage_routes[(stage_id, replica_id)]
 
     def register_stage_config(
         self,
         stage_id: int,
         stage_config: Any,
         coordinator_addresses: StageCoordinatorAddresses | None = None,
+        replica_id: int = 0,
     ) -> None:
         """Store the latest stage registration payload for *stage_id*."""
-        if stage_id not in self._allocations:
-            raise KeyError(stage_id)
-        self._stage_configs[stage_id] = stage_config
+        key = (stage_id, replica_id)
+        if key not in self._stage_routes:
+            raise KeyError(key)
+        self._stage_configs[key] = stage_config
         if coordinator_addresses is not None:
-            self._stage_coordinator_addresses[stage_id] = coordinator_addresses
-        self._stage_config_events[stage_id].set()
+            self._stage_coordinator_addresses[key] = coordinator_addresses
+        self._stage_config_events[key].set()
 
-    def get_stage_config(self, stage_id: int, timeout_s: float | None = None) -> Any:
+    def get_stage_config(self, stage_id: int, timeout_s: float | None = None, replica_id: int = 0) -> Any:
         """Return the stage config for *stage_id*, waiting if necessary."""
-        if stage_id not in self._allocations:
-            raise KeyError(stage_id)
+        key = (stage_id, replica_id)
+        if key not in self._stage_routes:
+            raise KeyError(key)
 
-        if stage_id in self._stage_configs:
-            return self._stage_configs[stage_id]
+        if key in self._stage_configs:
+            return self._stage_configs[key]
 
-        if not self._stage_config_events[stage_id].wait(timeout=timeout_s):
-            raise TimeoutError(f"Timed out waiting for stage config for stage {stage_id}.")
+        if not self._stage_config_events[key].wait(timeout=timeout_s):
+            raise TimeoutError(f"Timed out waiting for stage config for stage {stage_id} replica {replica_id}.")
 
-        return self._stage_configs[stage_id]
+        return self._stage_configs[key]
 
     def get_stage_coordinator_addresses(
         self,
         stage_id: int,
         timeout_s: float | None = None,
+        replica_id: int = 0,
     ) -> StageCoordinatorAddresses:
         """Return the registered coordinator addresses for *stage_id*."""
-        if stage_id not in self._allocations:
-            raise KeyError(stage_id)
+        key = (stage_id, replica_id)
+        if key not in self._stage_routes:
+            raise KeyError(key)
 
-        if not self._stage_config_events[stage_id].is_set():
-            if not self._stage_config_events[stage_id].wait(timeout=timeout_s):
-                raise TimeoutError(f"Timed out waiting for stage registration for stage {stage_id}.")
+        if not self._stage_config_events[key].is_set():
+            if not self._stage_config_events[key].wait(timeout=timeout_s):
+                raise TimeoutError(
+                    f"Timed out waiting for stage registration for stage {stage_id} replica {replica_id}."
+                )
 
-        return self._stage_coordinator_addresses[stage_id]
+        return self._stage_coordinator_addresses[key]
 
-    def get_client_addresses(self, stage_id: int) -> dict[str, str]:
+    def get_client_addresses(self, stage_id: int, replica_id: int = 0) -> dict[str, str]:
         """Return the addresses the client-side sockets should *bind* to."""
-        alloc = self._allocations[stage_id]
+        alloc = self.get_allocation(stage_id, replica_id)
         return {
             "input_address": alloc.input_bind_address,
             "output_address": alloc.output_bind_address,
         }
 
-    def get_zmq_addresses(self, stage_id: int) -> EngineZmqAddresses:
+    def get_zmq_addresses(self, stage_id: int, replica_id: int = 0) -> EngineZmqAddresses:
         """Return EngineZmqAddresses using the *bind* (client) side addresses."""
-        alloc = self._allocations[stage_id]
+        alloc = self.get_allocation(stage_id, replica_id)
         return EngineZmqAddresses(
             inputs=[alloc.input_bind_address],
             outputs=[alloc.output_bind_address],
         )
 
-    def get_engine_zmq_addresses(self, stage_id: int) -> EngineZmqAddresses:
+    def get_engine_zmq_addresses(self, stage_id: int, replica_id: int = 0) -> EngineZmqAddresses:
         """Return EngineZmqAddresses using the *connect* (engine) addresses."""
-        alloc = self._allocations[stage_id]
+        alloc = self.get_allocation(stage_id, replica_id)
         return EngineZmqAddresses(
             inputs=[alloc.input_connect_address],
             outputs=[alloc.output_connect_address],
@@ -266,7 +280,7 @@ class OmniMasterServer:
         poller = zmq.Poller()
         poller.register(reg_socket, zmq.POLLIN)
 
-        pending: set[int] = set(self._allocations.keys())
+        pending: set[StageRoute] = set(self._stage_routes.keys())
 
         while pending and not self._stop_event.is_set():
             events: list[tuple[zmq.Socket, int]] = poller.poll(_POLL_PERIOD_MS)  # type: ignore[assignment]
@@ -276,15 +290,15 @@ class OmniMasterServer:
 
             for sock, _ in events:
                 if sock is reg_socket:
-                    sid = self._handle_registration(reg_socket)
-                    if sid is not None:
-                        pending.discard(sid)
+                    route = self._handle_registration(reg_socket)
+                    if route is not None:
+                        pending.discard(route)
 
         # Cleanup
         reg_socket.close(linger=0)
         logger.info("[OmniMasterServer] All stages registered; server thread exiting.")
 
-    def _handle_registration(self, reg_socket: zmq.Socket) -> int | None:  # type: ignore[type-arg]
+    def _handle_registration(self, reg_socket: zmq.Socket) -> StageRoute | None:  # type: ignore[type-arg]
         """Receive a stage registration and reply with the handshake address.
 
         Returns the registered stage_id on success, or None on failure.
@@ -305,10 +319,13 @@ class OmniMasterServer:
             return None
 
         stage_id: int | None = msg.get("stage_id")
-        if stage_id not in self._allocations:
+        replica_id = int(msg.get("replica_id", 0) or 0)
+        key = (stage_id, replica_id)
+        if key not in self._stage_routes:
             logger.warning(
-                "[OmniMasterServer] Received registration for unknown stage_id=%s",
+                "[OmniMasterServer] Received registration for unknown stage_id=%s replica_id=%s",
                 stage_id,
+                replica_id,
             )
             return None
 
@@ -320,9 +337,10 @@ class OmniMasterServer:
                 coordinator_output=msg.get("coordinator_output"),
                 frontend_stats_publish_address=msg.get("frontend_stats_publish_address"),
             ),
+            replica_id=replica_id,
         )
 
-        alloc = self._allocations[stage_id]
+        alloc = self._stage_routes[key]
         response = msgspec.msgpack.encode(
             {
                 "handshake_address": alloc.handshake_connect_address,
@@ -333,11 +351,12 @@ class OmniMasterServer:
         # ROUTER-DEALER: reply is [identity, payload] (no empty delimiter).
         reg_socket.send_multipart([identity, response])
         logger.info(
-            "[OmniMasterServer] Stage %d registered; assigned handshake=%s",
+            "[OmniMasterServer] Stage %d replica %d registered; assigned handshake=%s",
             stage_id,
+            replica_id,
             alloc.handshake_connect_address,
         )
-        return stage_id
+        return key
 
 
 def register_stage_with_omni_master(
@@ -348,6 +367,7 @@ def register_stage_with_omni_master(
     omni_stage_config: Any = None,
     coordinator: DPCoordinator | None = None,
     return_addresses: bool = False,
+    replica_id: int = 0,
 ) -> str | tuple[str, str, str]:
     """Register a stage with the omni master server.
 
@@ -363,6 +383,7 @@ def register_stage_with_omni_master(
             reg_sock.connect(f"tcp://{omni_master_address}:{omni_master_port}")
             payload = {
                 "stage_id": omni_stage_id,
+                "replica_id": replica_id,
                 "stage_config": _serialize_stage_config(omni_stage_config),
             }
             if coordinator is not None:
@@ -469,9 +490,10 @@ def connect_remote_engine_cores(
     vllm_config: VllmConfig,
     omni_master_server: OmniMasterServer,
     stage_id: int,
+    replica_id: int = 0,
 ) -> Iterator[tuple[None, DPCoordinator | None, EngineZmqAddresses, None]]:
     """Wait for remote engine cores to connect through the omni handshake."""
-    addresses = omni_master_server.get_zmq_addresses(stage_id)
+    addresses = omni_master_server.get_zmq_addresses(stage_id, replica_id=replica_id)
     parallel_config = vllm_config.parallel_config
     # Mirror the engine-count logic from launch_omni_core_engines.
     remote_engine_count = (
@@ -482,7 +504,10 @@ def connect_remote_engine_cores(
     start_index = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0
     coordinator = None
 
-    registered_coordinator_addresses = omni_master_server.get_stage_coordinator_addresses(stage_id)
+    registered_coordinator_addresses = omni_master_server.get_stage_coordinator_addresses(
+        stage_id,
+        replica_id=replica_id,
+    )
     addresses.coordinator_input = registered_coordinator_addresses.coordinator_input
     addresses.coordinator_output = registered_coordinator_addresses.coordinator_output
     addresses.frontend_stats_publish_address = registered_coordinator_addresses.frontend_stats_publish_address
@@ -490,12 +515,13 @@ def connect_remote_engine_cores(
     engines_to_handshake = [CoreEngine(index=start_index + i, local=False) for i in range(remote_engine_count)]
 
     logger.info(
-        "Waiting for %d remote engine(s) for stage %d",
+        "Waiting for %d remote engine(s) for stage %d replica %d",
         remote_engine_count,
         stage_id,
+        replica_id,
     )
 
-    handshake_bind_address = omni_master_server.get_allocation(stage_id).handshake_bind_address
+    handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address
 
     with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
         yield None, coordinator, addresses, None
@@ -516,9 +542,10 @@ def launch_omni_core_engines(
     omni_master_server: OmniMasterServer,
     stage_id: int,
     stage_config: Any = None,
+    replica_id: int = 0,
 ) -> Iterator[tuple[CoreEngineProcManager, DPCoordinator | None, EngineZmqAddresses]]:
     """Launch local engine cores using the omni registration flow."""
-    addresses = omni_master_server.get_zmq_addresses(stage_id)
+    addresses = omni_master_server.get_zmq_addresses(stage_id, replica_id=replica_id)
     parallel_config = vllm_config.parallel_config
     # Determine the number of local engines and their ranks.
     local_engine_count = (
@@ -546,17 +573,19 @@ def launch_omni_core_engines(
         addresses.frontend_stats_publish_address = coordinator.get_stats_publish_address()
 
         logger.info(
-            "[omni] Started DP Coordinator process for stage %d (PID: %d)",
+            "[omni] Started DP Coordinator process for stage %d replica %d (PID: %d)",
             stage_id,
+            replica_id,
             coordinator.proc.pid,
         )
     else:
         coordinator = None
 
     logger.info(
-        "Starting %d local engine(s) for stage %d (dp_rank=%d)",
+        "Starting %d local engine(s) for stage %d replica %d (dp_rank=%d)",
         local_engine_count,
         stage_id,
+        replica_id,
         dp_rank,
     )
 
@@ -568,6 +597,7 @@ def launch_omni_core_engines(
         omni_stage_id=stage_id,
         omni_stage_config=stage_config,
         coordinator=coordinator,
+        replica_id=replica_id,
     )
 
     # One CoreEngine entry per local engine so wait_for_engine_startup can
@@ -575,7 +605,7 @@ def launch_omni_core_engines(
     engines_to_handshake = [CoreEngine(index=start_index + i, local=True) for i in range(local_engine_count)]
 
     # Bind the pre-allocated handshake socket for this stage.
-    handshake_bind_address = omni_master_server.get_allocation(stage_id).handshake_bind_address
+    handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address
 
     with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
         local_engine_manager = CoreEngineProcManager(

vllm_omni/engine/stage_init_utils.py

@@ -19,21 +19,48 @@ from typing import Any, Literal
 
 from vllm.logger import init_logger
 from vllm.sampling_params import SamplingParams
+from vllm.tokenizers import cached_tokenizer_from_config
 from vllm.usage.usage_lib import UsageContext
 from vllm.v1.engine.input_processor import InputProcessor
 from vllm.v1.executor import Executor
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices
 from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams
+from vllm_omni.inputs.preprocess import OmniInputPreprocessor
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.quantization.inc_config import OmniINCConfig
 
 logger = init_logger(__name__)
 
 
+@dataclass
+class ReplicaInitPlan:
+    """One concrete replica startup unit within a logical stage."""
+
+    replica_id: int
+    num_replicas: int
+    launch_mode: str
+    stage_cfg: Any
+    metadata: Any
+    stage_connector_spec: dict[str, Any]
+    omni_kv_connector: tuple[dict[str, Any] | None, str | None, str | None]
+    stage_vllm_config: Any | None = None
+    executor_class: type | None = None
+
+
+@dataclass
+class LogicalStageInitPlan:
+    """Startup plan for one logical stage."""
+
+    stage_idx: int
+    configured_stage_id: int
+    replicas: list[ReplicaInitPlan]
+
+
 def _resolve_model_to_local_path(model: str) -> str:
     """Resolve an HF Hub model ID to a local cache path."""
     if os.path.isdir(model):
@@ -85,6 +112,14 @@ def terminate_alive_proc(proc, timeout=5):
             proc.kill()
 
 
+def patch_generation_config_if_needed(model_config: Any) -> None:
+    """Guard InputProcessor init for models whose config lacks model_type."""
+    try:
+        model_config.try_get_generation_config()
+    except Exception:
+        model_config.try_get_generation_config = lambda: {}
+
+
 def resolve_worker_cls(engine_args: dict[str, Any]) -> None:
     """Resolve worker_cls from worker_type for non-diffusion stages."""
     worker_type = engine_args.get("worker_type", None)
@@ -259,20 +294,9 @@ class StageMetadata:
     runtime_cfg: Any
     prompt_expand_func: Callable | None = None
     cfg_kv_collect_func: Callable | None = None
-
-
-@dataclass
-class StartedLlmStage:
-    """Resources for an LLM stage that has completed startup."""
-
-    stage_id: int
-    metadata: Any
-    vllm_config: Any
-    executor_class: type
-    addresses: Any
-    proc: Any = None
-    engine_manager: Any = None
-    coordinator: Any = None
+    # Multi-replica: replica_id distinguishes replicas of the same stage.
+    # For single-replica stages this defaults to 0.
+    replica_id: int = 0
 
 
 def extract_stage_metadata(stage_config: Any) -> StageMetadata:
@@ -375,6 +399,119 @@ def prepare_engine_environment() -> None:
         pass
 
 
+def split_devices_for_replicas(
+    devices_str: str | None,
+    num_replicas: int,
+    tp_size: int,
+    stage_id: int,
+) -> list[str]:
+    """Split a devices string into per-replica subsets.
+
+    When ``num_replicas`` is 1, returns ``[devices_str]`` unchanged.
+    Otherwise, the total number of device IDs must equal
+    ``num_replicas * tp_size``; each replica gets ``tp_size`` consecutive
+    device IDs.
+
+    Example::
+
+        split_devices_for_replicas("1,2,3,4", num_replicas=2, tp_size=2, stage_id=1)
+        # → ["1,2", "3,4"]
+    """
+    if num_replicas <= 1 or devices_str is None:
+        return [devices_str] if devices_str is not None else [devices_str]
+
+    device_list = [d.strip() for d in devices_str.split(",") if d.strip()]
+    required = num_replicas * tp_size
+    if len(device_list) != required:
+        raise ValueError(
+            f"Stage {stage_id}: num_replicas={num_replicas}, "
+            f"tensor_parallel_size={tp_size} requires "
+            f"{required} devices, got {len(device_list)}: {devices_str}"
+        )
+
+    result: list[str] = []
+    for r in range(num_replicas):
+        chunk = device_list[r * tp_size : (r + 1) * tp_size]
+        result.append(",".join(chunk))
+    return result
+
+
+def get_stage_tp_size(stage_cfg: Any) -> int:
+    """Extract tensor_parallel_size from a stage config object."""
+    engine_args = getattr(stage_cfg, "engine_args", {})
+    if hasattr(engine_args, "get"):
+        return int(engine_args.get("tensor_parallel_size", 1) or 1)
+    return int(getattr(engine_args, "tensor_parallel_size", 1) or 1)
+
+
+def get_stage_devices_per_replica(stage_cfg: Any) -> int:
+    """Return the number of devices consumed by one replica of *stage_cfg*."""
+    if getattr(stage_cfg, "stage_type", "llm") != "diffusion":
+        return get_stage_tp_size(stage_cfg)
+
+    parallel_config = _get_attr_or_item(getattr(stage_cfg, "engine_args", {}), "parallel_config")
+    if parallel_config is None:
+        return 1
+
+    world_size = _get_attr_or_item(parallel_config, "world_size")
+    if world_size is not None:
+        return max(1, int(world_size))
+
+    try:
+        from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+        return max(1, int(DiffusionParallelConfig.from_dict(_to_dict(parallel_config)).world_size))
+    except Exception:
+        return 1
+
+
+def compute_replica_layout(
+    stage_configs: Sequence[Any],
+) -> tuple[list[int], dict[int, list[str]]]:
+    """Compute per-stage replica counts and device assignments.
+
+    Returns:
+        replicas_per_stage: num_replicas per logical stage.
+        replica_devices_map: stage_idx -> per-replica device strings
+            (only for stages with num_replicas > 1).
+    """
+    replicas_per_stage: list[int] = []
+    for stage_cfg in stage_configs:
+        runtime_cfg = getattr(stage_cfg, "runtime", {})
+        num_replicas = int(
+            runtime_cfg.get("num_replicas", 1)
+            if hasattr(runtime_cfg, "get")
+            else getattr(runtime_cfg, "num_replicas", 1)
+        )
+        replicas_per_stage.append(max(1, num_replicas))
+
+    replica_devices_map: dict[int, list[str]] = {}
+    for stage_id, stage_cfg in enumerate(stage_configs):
+        num_replicas = replicas_per_stage[stage_id]
+        if num_replicas <= 1:
+            continue
+        runtime_cfg = getattr(stage_cfg, "runtime", {})
+        devices_str = (
+            runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None)
+        )
+        devices_per_replica = get_stage_devices_per_replica(stage_cfg)
+        replica_devices_map[stage_id] = split_devices_for_replicas(
+            devices_str,
+            num_replicas,
+            devices_per_replica,
+            stage_id,
+        )
+        logger.info(
+            "[stage_init] Stage %s: %d replicas, devices_per_replica=%d, devices split: %s",
+            stage_id,
+            num_replicas,
+            devices_per_replica,
+            replica_devices_map[stage_id],
+        )
+
+    return replicas_per_stage, replica_devices_map
+
+
 def setup_stage_devices(stage_id: int, runtime_cfg: Any) -> None:
     """Device mapping via set_stage_devices for a single stage."""
     physical_devices = set_stage_devices(
@@ -419,6 +556,13 @@ def build_engine_args_dict(
     if stage_type != "diffusion":
         resolve_worker_cls(engine_args_dict)
 
+    if engine_args_dict.get("worker_type") == "generation":
+        # Non-AR generation stages (e.g. code2wav) do not benefit from
+        # prefix caching and can expose hybrid KV-cache layouts that vLLM's
+        # prefix-cache coordinator does not handle.
+        engine_args_dict.setdefault("disable_hybrid_kv_cache_manager", True)
+        engine_args_dict.setdefault("enable_prefix_caching", False)
+
     # Check whether the stage's default_sampling_params defines extra_args.
     default_sp = _to_dict(getattr(stage_config, "default_sampling_params", {}))
     engine_args_dict["has_sampling_extra_args"] = bool(default_sp.get("extra_args"))
@@ -475,6 +619,35 @@ def build_vllm_config(
     return vllm_config, executor_class
 
 
+def build_llm_stage_output_processor(plan: LogicalStageInitPlan, stage_vllm_config: Any) -> Any | None:
+    """Build one output processor per logical LLM stage."""
+
+    metadata = plan.replicas[0].metadata
+    if stage_vllm_config.model_config.skip_tokenizer_init:
+        tokenizer = None
+    else:
+        tokenizer = cached_tokenizer_from_config(
+            model_config=stage_vllm_config.model_config,
+        )
+    return MultimodalOutputProcessor(
+        tokenizer=tokenizer,
+        log_stats=False,
+        engine_core_output_type=metadata.engine_output_type,
+    )
+
+
+def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor:
+    """Build the shared stage-0 input processor."""
+
+    patch_generation_config_if_needed(stage_vllm_config.model_config)
+    input_processor = InputProcessor(vllm_config=stage_vllm_config)
+    input_processor.input_preprocessor = OmniInputPreprocessor(
+        vllm_config=stage_vllm_config,
+        renderer=input_processor.renderer,
+    )
+    return input_processor
+
+
 def acquire_device_locks(
     stage_id: int,
     engine_args_dict: dict[str, Any],
@@ -678,7 +851,6 @@ def build_diffusion_config(
     metadata: StageMetadata,
 ) -> Any:
     """Build diffusion config for a stage."""
-    from vllm_omni.diffusion.data import OmniDiffusionConfig
 
     engine_args_dict = build_engine_args_dict(stage_cfg, model)
     od_config = OmniDiffusionConfig.from_kwargs(**engine_args_dict)
@@ -726,108 +898,5 @@ def initialize_diffusion_stage(
     """
     from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client
 
-    engine_args = _to_dict(stage_cfg.engine_args)
-    engine_args.pop("stage_id", None)
-    od_config = OmniDiffusionConfig.from_kwargs(
-        stage_id=stage_id,
-        model=model,
-        **engine_args,
-    )
-    if metadata.cfg_kv_collect_func is not None:
-        od_config.cfg_kv_collect_func = metadata.cfg_kv_collect_func
     od_config = build_diffusion_config(model, stage_cfg, metadata)
     return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline)
-
-
-def _shutdown_or_close_resource(resource: Any, resource_name: str, stage_id: int) -> None:
-    """vLLM CoreEngineProcManager / coordinators use ``shutdown()``, not ``close()``."""
-    if resource is None:
-        return
-    shutdown = getattr(resource, "shutdown", None)
-    if callable(shutdown):
-        try:
-            shutdown()
-        except Exception as cleanup_error:
-            logger.warning(
-                "[stage_init] Failed to shutdown launched %s for stage %s: %s",
-                resource_name,
-                stage_id,
-                cleanup_error,
-            )
-        return
-    close = getattr(resource, "close", None)
-    if callable(close):
-        try:
-            close()
-        except Exception as cleanup_error:
-            logger.warning(
-                "[stage_init] Failed to close launched %s for stage %s: %s",
-                resource_name,
-                stage_id,
-                cleanup_error,
-            )
-
-
-def close_started_llm_stage(started: StartedLlmStage) -> None:
-    """Release resources owned by a launched stage that never attached."""
-    if started.proc is not None:
-        try:
-            terminate_alive_proc(started.proc)
-        except Exception as cleanup_error:
-            logger.warning(
-                "[stage_init] Failed to terminate process for stage %s: %s",
-                started.stage_id,
-                cleanup_error,
-            )
-    _shutdown_or_close_resource(started.engine_manager, "engine manager", started.stage_id)
-    _shutdown_or_close_resource(started.coordinator, "coordinator", started.stage_id)
-
-
-def finalize_initialized_stages(
-    stage_clients: list[Any | None],
-    input_processor: InputProcessor | None,
-) -> tuple[list[Any], list[Any], list[dict[str, Any]]]:
-    """Validate successful init and build runtime metadata lists."""
-    if any(stage_client is None for stage_client in stage_clients):
-        raise RuntimeError("Stage initialization completed with missing stage clients")
-
-    initialized_stage_clients = [stage_client for stage_client in stage_clients if stage_client is not None]
-    default_sampling_params_list = [stage_client.default_sampling_params for stage_client in initialized_stage_clients]
-    stage_metadata = [
-        {
-            "final_output": stage_client.final_output,
-            "final_output_type": stage_client.final_output_type,
-            "stage_type": stage_client.stage_type,
-        }
-        for stage_client in initialized_stage_clients
-    ]
-
-    if not isinstance(input_processor, InputProcessor):
-        has_llm_stage = any(metadata.get("stage_type") != "diffusion" for metadata in stage_metadata)
-        if has_llm_stage:
-            raise RuntimeError("Failed to initialize stage-0 InputProcessor for LLM pipeline")
-
-    return initialized_stage_clients, default_sampling_params_list, stage_metadata
-
-
-def cleanup_failed_stage_initialization(
-    stage_clients: list[Any | None],
-    started_llm_stages: list[StartedLlmStage],
-) -> None:
-    """Shutdown attached stages and close any launched-but-unattached engines."""
-    for cleanup_stage_id, stage_client in reversed(list(enumerate(stage_clients))):
-        if stage_client is None:
-            continue
-        try:
-            stage_client.shutdown()
-        except Exception as cleanup_error:
-            logger.warning(
-                "[stage_init] Failed to shutdown initialized stage %s after init failure: %s",
-                cleanup_stage_id,
-                cleanup_error,
-            )
-
-    for started in reversed(started_llm_stages):
-        if stage_clients[started.stage_id] is not None:
-            continue
-        close_started_llm_stage(started)

vllm_omni/engine/stage_pool.py

@@ -0,0 +1,403 @@
+"""Unified stage-local runtime abstraction for vLLM-Omni."""
+
+from __future__ import annotations
+
+import asyncio
+import time as _time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from vllm.logger import init_logger
+from vllm.v1.engine import EngineCoreOutputs
+
+from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
+from vllm_omni.metrics.stats import StageStats
+from vllm_omni.metrics.utils import count_tokens_from_outputs
+
+if TYPE_CHECKING:
+    from vllm_omni.engine.orchestrator import OrchestratorRequestState
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class _ReplicaMetrics:
+    """Per-replica metrics accumulators owned by a stage pool."""
+
+    batch_seq: int = 0
+    agg_total_tokens: int = 0
+    agg_total_gen_time_ms: float = 0.0
+
+
+class StagePool:
+    """Replicas of one logical stage with RR + affinity selection."""
+
+    def __init__(
+        self,
+        stage_id: int,
+        clients: Any | list[Any],
+        *,
+        output_processor: Any = None,
+        stage_vllm_config: Any = None,
+    ) -> None:
+        if isinstance(clients, list):
+            normalized_clients = list(clients)
+        else:
+            normalized_clients = [clients]
+
+        if not normalized_clients:
+            raise ValueError(f"StagePool for stage {stage_id} has no replicas")
+        self.stage_id = stage_id
+        self.clients: list[Any] = normalized_clients
+        self._output_processor = output_processor
+        self._stage_vllm_config = stage_vllm_config
+        self._next_replica_id = 0
+        self._request_bindings: dict[str, int] = {}
+        self._replica_metrics: list[_ReplicaMetrics] = [_ReplicaMetrics() for _ in self.clients]
+
+    # ---- Stage-level properties ----
+
+    @property
+    def num_replicas(self) -> int:
+        return len(self.clients)
+
+    @property
+    def stage_type(self) -> str | None:
+        return getattr(self.stage_client, "stage_type", None)
+
+    @property
+    def final_output(self) -> bool:
+        return bool(getattr(self.clients[0], "final_output", False))
+
+    @property
+    def stage_client(self) -> Any:
+        return self.clients[0]
+
+    @property
+    def stage_vllm_config(self) -> Any:
+        return self._stage_vllm_config
+
+    @property
+    def output_processor(self) -> Any:
+        return self._output_processor
+
+    # ---- Route binding lifecycle ----
+
+    def get_bound_replica_id(self, request_id: str) -> int | None:
+        """Return the currently bound replica id for *request_id* if present."""
+        return self._request_bindings.get(request_id)
+
+    def get_bound_client(self, request_id: str) -> Any | None:
+        """Return the currently bound client for *request_id* if present."""
+        replica_id = self.get_bound_replica_id(request_id)
+        if replica_id is None:
+            return None
+        return self.clients[replica_id]
+
+    def release_binding(self, request_id: str) -> None:
+        """Drop the route binding for *request_id* in this stage."""
+        self._request_bindings.pop(request_id, None)
+
+    def release_bindings(self, request_ids: list[str]) -> None:
+        """Drop route bindings for the given request ids in this stage."""
+        for request_id in request_ids:
+            self.release_binding(request_id)
+
+    def select_replica_id(
+        self,
+        request_id: str,
+        *,
+        affinity_request_id: str | None = None,
+    ) -> int:
+        """Pick a replica id for *request_id* and cache the choice."""
+        cached = self.get_bound_replica_id(request_id)
+        if cached is not None:
+            return cached
+
+        chosen = self.get_bound_replica_id(affinity_request_id) if affinity_request_id is not None else None
+        if chosen is None:
+            if self.num_replicas == 1:
+                chosen = 0
+            else:
+                chosen = self._next_replica_id
+                self._next_replica_id = (self._next_replica_id + 1) % self.num_replicas
+
+        self._request_bindings[request_id] = chosen
+        return chosen
+
+    # ---- Metrics ----
+
+    def build_stage_metrics(
+        self,
+        request_outputs: list[Any],
+        *,
+        submit_ts: float,
+        replica_id: int,
+    ) -> StageRequestMetrics:
+        """Build stage metrics for outputs produced on one replica."""
+        now = _time.time()
+        stage_gen_time_ms = (now - submit_ts) * 1000.0
+
+        num_tokens_out = count_tokens_from_outputs(request_outputs)
+        num_tokens_in = 0
+        if self.stage_id == 0:
+            for ro in request_outputs:
+                ptids = getattr(ro, "prompt_token_ids", None)
+                if ptids is not None:
+                    num_tokens_in += len(ptids)
+
+        metrics = self._replica_metrics[replica_id]
+        metrics.batch_seq += 1
+        batch_id = metrics.batch_seq
+        metrics.agg_total_tokens += num_tokens_out
+        metrics.agg_total_gen_time_ms += stage_gen_time_ms
+
+        return StageRequestMetrics(
+            num_tokens_in=num_tokens_in,
+            num_tokens_out=num_tokens_out,
+            stage_gen_time_ms=stage_gen_time_ms,
+            batch_id=batch_id,
+            batch_size=1,
+            rx_decode_time_ms=0.0,
+            rx_transfer_bytes=0,
+            rx_in_flight_time_ms=0.0,
+            stage_stats=StageStats(
+                total_token=metrics.agg_total_tokens,
+                total_gen_time_ms=metrics.agg_total_gen_time_ms,
+            ),
+        )
+
+    # ---- Stage-local admission ----
+
+    async def submit_initial(
+        self,
+        request_id: str,
+        req_state: OrchestratorRequestState,
+        request: Any,
+        *,
+        prompt_text: Any = None,
+        affinity_request_id: str | None = None,
+        submit_kwargs: dict[str, Any] | None = None,
+        params_override: Any = None,
+    ) -> int:
+        """Submit a stage-entry request into this pool."""
+        params = params_override if params_override is not None else req_state.sampling_params_list[self.stage_id]
+        submit_kwargs = dict(submit_kwargs or {})
+        if self.stage_type == "diffusion":
+            replica_id = self.select_replica_id(
+                request_id,
+                affinity_request_id=affinity_request_id,
+            )
+            client = self.clients[replica_id]
+            if isinstance(request, list):
+                await client.add_batch_request_async(request_id, request, params, **submit_kwargs)
+            else:
+                await client.add_request_async(request_id, request, params, **submit_kwargs)
+            return replica_id
+
+        replica_id = self.select_replica_id(
+            request_id,
+            affinity_request_id=affinity_request_id,
+        )
+        try:
+            self.output_processor.add_request(
+                request=request,
+                prompt=prompt_text,
+                parent_req=None,
+                request_index=0,
+                queue=None,
+            )
+        except Exception:
+            self.release_binding(request_id)
+            raise
+
+        try:
+            await self.clients[replica_id].add_request_async(request, **submit_kwargs)
+        except Exception:
+            self.release_binding(request_id)
+            rollback = getattr(self.output_processor, "remove_request", None)
+            if callable(rollback):
+                try:
+                    rollback(request_id)
+                except Exception as rollback_error:
+                    logger.warning(
+                        "[StagePool] Failed to rollback output processor state for req=%s stage-%s: %s",
+                        request_id,
+                        self.stage_id,
+                        rollback_error,
+                    )
+            raise
+        return replica_id
+
+    async def submit_update(
+        self,
+        request_id: str,
+        req_state: OrchestratorRequestState,
+        request: Any,
+        *,
+        prompt_text: Any = None,
+    ) -> int:
+        """Submit a streaming update to an already admitted request."""
+        params = req_state.sampling_params_list[self.stage_id]
+        replica_id = self.get_bound_replica_id(request_id)
+        if replica_id is None:
+            replica_id = self.select_replica_id(request_id)
+
+        if self.stage_type == "diffusion":
+            await self.clients[replica_id].add_request_async(request_id, request, params)
+        else:
+            # Refresh the shared output-processor state before yielding to the
+            # stage client so streaming segments are merged against the latest
+            # prompt/token metadata.
+            self.output_processor.add_request(
+                request=request,
+                prompt=prompt_text,
+                parent_req=None,
+                request_index=0,
+                queue=None,
+            )
+            await self.clients[replica_id].add_request_async(request)
+        return replica_id
+
+    # ---- Stage-local polling ----
+
+    async def _poll_stage_raw(self, client: Any) -> EngineCoreOutputs | None:
+        """Pull raw EngineCoreOutputs from a stage replica without processing."""
+        outputs = await client.get_output_async()
+        if not outputs.outputs:
+            return None
+        return outputs
+
+    async def process_llm_raw_outputs(
+        self,
+        replica_id: int,
+        raw_outputs: EngineCoreOutputs,
+    ) -> list[Any]:
+        """Run the shared LLM output processor on one raw poll result."""
+        client = self.clients[replica_id]
+        processor = self.output_processor
+        processed = processor.process_outputs(
+            raw_outputs.outputs,
+            raw_outputs.timestamp,
+            None,
+        )
+
+        if processed.reqs_to_abort:
+            await client.abort_requests_async(processed.reqs_to_abort)
+
+        if raw_outputs.scheduler_stats is not None:
+            processor.update_scheduler_stats(raw_outputs.scheduler_stats)
+
+        return processed.request_outputs
+
+    async def poll_llm_raw_output(
+        self,
+        replica_id: int,
+        *,
+        timeout_s: float = 0.001,
+    ) -> EngineCoreOutputs | None:
+        """Poll raw EngineCore outputs from one LLM replica once."""
+        client = self.clients[replica_id]
+        try:
+            return await asyncio.wait_for(
+                self._poll_stage_raw(client),
+                timeout=timeout_s,
+            )
+        except asyncio.TimeoutError:
+            return None
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logger.exception(
+                "[StagePool] _poll_stage_raw failed for stage-%s replica-%s",
+                self.stage_id,
+                replica_id,
+            )
+            raise
+
+    def poll_diffusion_output(self, replica_id: int) -> Any | None:
+        """Drain one ready diffusion output from the given replica if present."""
+        return self.clients[replica_id].get_diffusion_output_nowait()
+
+    # ---- Stage-local control plane ----
+
+    async def abort_requests(self, request_ids: list[str]) -> None:
+        """Abort the given requests in this stage pool.
+
+        Request-bound abort routing stays inside the pool because route affinity
+        (`request_id -> replica_id`) is pool-owned.
+        """
+        if not request_ids:
+            return
+
+        request_ids_by_replica: dict[int, list[str]] = {}
+        for request_id in request_ids:
+            replica_id = self.get_bound_replica_id(request_id)
+            if replica_id is None:
+                logger.debug("[StagePool] abort: no binding for req=%s in stage-%s", request_id, self.stage_id)
+                continue
+            request_ids_by_replica.setdefault(replica_id, []).append(request_id)
+
+        for replica_id, replica_request_ids in request_ids_by_replica.items():
+            await self.clients[replica_id].abort_requests_async(replica_request_ids)
+
+        # Clean up OutputProcessor state (e.g. mm_accumulated tensors) that
+        # would otherwise leak — aborted requests never produce a final
+        # EngineCoreOutput, so process_outputs() never fires its cleanup path.
+        all_aborted = [rid for ids in request_ids_by_replica.values() for rid in ids]
+        if all_aborted and self._output_processor is not None:
+            self._output_processor.abort_requests(all_aborted, internal=True)
+
+    async def collective_rpc(
+        self,
+        replica_id: int,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        """Dispatch a stage-scoped control-plane RPC to one physical route."""
+        kwargs = dict(kwargs or {})
+        client = self.clients[replica_id]
+        try:
+            if hasattr(client, "collective_rpc_async"):
+                return await client.collective_rpc_async(
+                    method=method,
+                    timeout=timeout,
+                    args=args,
+                    kwargs=kwargs,
+                )
+            return {
+                "supported": False,
+                "todo": True,
+                "reason": f"{client.__class__.__name__}.collective_rpc_async is not implemented yet",
+            }
+        except Exception as exc:
+            logger.exception(
+                "[StagePool] collective_rpc failed: stage=%s replica=%s method=%s",
+                self.stage_id,
+                replica_id,
+                method,
+            )
+            return {
+                "supported": False,
+                "error": str(exc),
+            }
+
+    def shutdown_replica(self, replica_id: int) -> None:
+        """Shutdown one backend handle in this stage pool."""
+        client = self.clients[replica_id]
+        try:
+            client.shutdown()
+            logger.info(
+                "[StagePool] Stage %d replica %d shut down",
+                self.stage_id,
+                replica_id,
+            )
+        except Exception as e:
+            logger.warning(
+                "[StagePool] Failed to shutdown stage %d replica %d: %s",
+                self.stage_id,
+                replica_id,
+                e,
+            )

vllm_omni/entrypoints/cli/serve.py

@@ -187,6 +187,12 @@ class OmniServeCommand(CLISubcommand):
             default=None,
             help="Select and launch a single stage by stage_id.",
         )
+        omni_config_group.add_argument(
+            "--replica-id",
+            type=int,
+            default=0,
+            help="Replica id to register when launching a single headless stage.",
+        )
         omni_config_group.add_argument(
             "--stage-init-timeout",
             type=int,
@@ -529,14 +535,18 @@ def run_headless(args: argparse.Namespace) -> None:
 
     model = args.model
     stage_id: int | None = args.stage_id
+    replica_id: int = args.replica_id
     omni_master_address: str | None = args.omni_master_address
     omni_master_port: int | None = args.omni_master_port
 
     if stage_id is None:
         raise ValueError("--stage-id is required in headless mode")
+    if replica_id < 0:
+        raise ValueError("--replica-id must be >= 0 in headless mode")
     if omni_master_address is None or omni_master_port is None:
         raise ValueError("--omni-master-address and --omni-master-port are required in headless mode")
-    if getattr(args, "api_server_count", 0) and args.api_server_count > 1:
+    api_server_count = args.api_server_count or 0
+    if api_server_count > 1:
         raise ValueError("api_server_count can't be set in headless mode")
     if args.worker_backend != "multi_process":
         raise ValueError("headless mode requires worker_backend=multi_process")
@@ -552,20 +562,19 @@ def run_headless(args: argparse.Namespace) -> None:
     # Locate the stage config that matches stage_id.
     stage_cfg = None
     for cfg in stage_configs:
-        if getattr(cfg, "stage_id", None) == stage_id:
+        if cfg.stage_id == stage_id:
             stage_cfg = cfg
             break
     if stage_cfg is None:
         raise ValueError(
-            f"No stage config found for stage_id={stage_id}. "
-            f"Available stage ids: {[getattr(c, 'stage_id', None) for c in stage_configs]}"
+            f"No stage config found for stage_id={stage_id}. Available stage ids: {[c.stage_id for c in stage_configs]}"
         )
 
     prepare_engine_environment()
     omni_transfer_config = load_omni_transfer_config_for_model(model, config_path)
     omni_conn_cfg, omni_from, omni_to = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id)
 
-    if getattr(stage_cfg, "stage_type", "llm") == "diffusion":
+    if stage_cfg.stage_type == "diffusion":
         metadata = extract_stage_metadata(stage_cfg)
         if omni_conn_cfg:
             inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to)
@@ -573,8 +582,9 @@ def run_headless(args: argparse.Namespace) -> None:
         od_config = build_diffusion_config(model, stage_cfg, metadata)
 
         logger.info(
-            "[Headless] Launching diffusion stage %d via OmniMasterServer at %s:%d",
+            "[Headless] Launching diffusion stage %d replica %d via OmniMasterServer at %s:%d",
             stage_id,
+            replica_id,
             omni_master_address,
             omni_master_port,
         )
@@ -587,6 +597,7 @@ def run_headless(args: argparse.Namespace) -> None:
                 omni_stage_id=stage_id,
                 omni_stage_config=stage_cfg,
                 return_addresses=True,
+                replica_id=replica_id,
             )
             proc, _, _, _ = spawn_diffusion_proc(
                 model,
@@ -595,13 +606,13 @@ def run_headless(args: argparse.Namespace) -> None:
                 request_address=request_address,
                 response_address=response_address,
             )
-            complete_diffusion_handshake(proc, handshake_address)
+            complete_diffusion_handshake(proc, handshake_address, args.stage_init_timeout)
             proc.join()
             if proc.exitcode not in (None, 0):
-                raise RuntimeError(f"Diffusion stage {stage_id} exited with code {proc.exitcode}")
+                raise RuntimeError(f"Diffusion stage {stage_id} replica {replica_id} exited with code {proc.exitcode}")
             return
         finally:
-            logger.info("[Headless] Shutting down stage %d.", stage_id)
+            logger.info("[Headless] Shutting down stage %d replica %d.", stage_id, replica_id)
             if proc is not None and proc.is_alive():
                 terminate_alive_proc(proc)
 
@@ -677,15 +688,17 @@ def run_headless(args: argparse.Namespace) -> None:
             enable_wave_coordination=vllm_config.model_config.is_moe,
         )
         logger.info(
-            "[Headless] Started DP Coordinator process for stage %d (PID: %d)",
+            "[Headless] Started DP Coordinator process for stage %d replica %d (PID: %d)",
             stage_id,
+            replica_id,
             coordinator.proc.pid,
         )
 
     logger.info(
-        "[Headless] Launching %d engine core(s) for stage %d via OmniMasterServer at %s:%d",
+        "[Headless] Launching %d engine core(s) for stage %d replica %d via OmniMasterServer at %s:%d",
         local_engine_count,
         stage_id,
+        replica_id,
         omni_master_address,
         omni_master_port,
     )
@@ -700,11 +713,12 @@ def run_headless(args: argparse.Namespace) -> None:
         omni_stage_id=stage_id,
         omni_stage_config=stage_cfg,
         coordinator=coordinator,
+        replica_id=replica_id,
     )
 
     engine_manager = None
-    log_stats = bool(getattr(args, "log_stats", False))
-    if getattr(args, "disable_log_stats", False):
+    log_stats = bool(args.log_stats)
+    if args.disable_log_stats:
         log_stats = False
 
     try:

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -63,15 +63,11 @@ def _decode_additional_information(raw_info: Any) -> dict[str, Any]:
 
 
 def text2flow(
-    stage_list: list[Any],
-    engine_input_source: list[int],
+    source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt = None,
-    requires_multimodal_data: bool = True,
+    _requires_multimodal_data: bool = True,
 ):
     """Build stage-1 inputs by prefixing stage-0 prompt ids to its outputs."""
-    source_stage_id = engine_input_source[0]
-    source_outputs = stage_list[source_stage_id].engine_outputs
-
     engine_inputs: list[OmniTokensPrompt] = []
     for source_output in source_outputs:
         output = source_output.outputs[0]

vllm_omni/model_executor/stage_input_processors/dynin_omni.py

@@ -87,22 +87,10 @@ def _decode_runtime_bridge_info(value: Any) -> dict[str, Any]:
 
 
 def _bridge_tokens(
-    stage_list,
-    engine_input_source,
+    source_outputs,
     prompt: OmniTokensPrompt | TextPrompt = None,
     requires_multimodal_data: bool = False,
 ):
-    if not engine_input_source:
-        raise ValueError("engine_input_source cannot be empty")
-
-    source_stage_id = engine_input_source[0]
-    if source_stage_id >= len(stage_list):
-        raise IndexError(f"Invalid stage_id: {source_stage_id}")
-
-    if stage_list[source_stage_id].engine_outputs is None:
-        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
-
-    source_outputs = stage_list[source_stage_id].engine_outputs
     next_inputs = []
     if not isinstance(prompt, list):
         prompt = [prompt]
@@ -119,9 +107,7 @@ def _bridge_tokens(
         if not token_ids:
             token_ids = list(output.cumulative_token_ids or [])
         if not token_ids:
-            raise RuntimeError(
-                f"Stage {source_stage_id} output for request {source_output.request_id} has no token_ids"
-            )
+            raise RuntimeError(f"Stage output for request {source_output.request_id} has no token_ids")
 
         detok_id = _to_int(mm_out.get("detok_id"), default=0)
         src_prompt = prompt_meta_by_reqid.get(source_output.request_id, {})
@@ -147,18 +133,16 @@ def _bridge_tokens(
 
 
 def token2text_to_token2image(
-    stage_list,
-    engine_input_source,
+    source_outputs,
     prompt: OmniTokensPrompt | TextPrompt = None,
     requires_multimodal_data: bool = False,
 ):
-    return _bridge_tokens(stage_list, engine_input_source, prompt, requires_multimodal_data)
+    return _bridge_tokens(source_outputs, prompt, requires_multimodal_data)
 
 
 def token2image_to_token2audio(
-    stage_list,
-    engine_input_source,
+    source_outputs,
     prompt: OmniTokensPrompt | TextPrompt = None,
     requires_multimodal_data: bool = False,
 ):
-    return _bridge_tokens(stage_list, engine_input_source, prompt, requires_multimodal_data)
+    return _bridge_tokens(source_outputs, prompt, requires_multimodal_data)

vllm_omni/model_executor/stage_input_processors/fish_speech.py

@@ -24,16 +24,14 @@ def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None:
 
 
 def slow_ar_to_dac_decoder(
-    stage_list: list[Any],
-    engine_input_source: list[int],
-    prompt: Any = None,
-    requires_multimodal_data: bool = False,
+    source_outputs: list[Any],
+    _prompt: Any = None,
+    _requires_multimodal_data: bool = False,
 ) -> list[Any]:
     """Non-async processor: wait for Slow AR to finish, then pass all codes to DAC decoder."""
     from vllm_omni.inputs.data import OmniTokensPrompt
-    from vllm_omni.model_executor.stage_input_processors.qwen3_omni import _validate_stage_inputs
 
-    slow_ar_outputs = _validate_stage_inputs(stage_list, engine_input_source)
+    slow_ar_outputs = source_outputs
     dac_inputs: list[OmniTokensPrompt] = []
 
     for output in slow_ar_outputs:

vllm_omni/model_executor/stage_input_processors/glm_image.py

@@ -212,25 +212,20 @@ def _parse_generated_tokens(
 
 
 def ar2diffusion(
-    stage_list: list[Any],
-    engine_input_source: list[int],
+    source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt | list | None = None,
     requires_multimodal_data: bool = False,
+    streaming_context: Any | None = None,
 ) -> list[dict[str, Any]]:
-    """Process AR stage outputs to create Diffusion stage inputs."""
+    """Process AR stage outputs to create Diffusion stage inputs.
+
+    This processor accepts the stage-pool transition interface:
+    ``ar2diffusion(source_outputs, prompt, requires_multimodal_data)``.
+    """
+    del streaming_context
+
     _t_total = time.perf_counter()
-
-    if not engine_input_source:
-        raise ValueError("engine_input_source cannot be empty")
-
-    source_stage_id = engine_input_source[0]
-    if source_stage_id >= len(stage_list):
-        raise IndexError(f"Invalid stage_id: {source_stage_id}")
-
-    if stage_list[source_stage_id].engine_outputs is None:
-        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
-
-    ar_outputs = stage_list[source_stage_id].engine_outputs
+    ar_outputs = source_outputs
     diffusion_inputs = []
 
     # Normalize prompt to list

vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py

@@ -22,33 +22,20 @@ logger = init_logger(__name__)
 
 
 def ar2diffusion(
-    stage_list: list[Any],
-    engine_input_source: list[int],
+    source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt | list | None = None,
     requires_multimodal_data: bool = False,
 ) -> list[dict[str, Any]]:
     """Process AR stage outputs to create Diffusion stage inputs.
 
     Args:
-        stage_list: List of stage clients (set by orchestrator).
-        engine_input_source: List of source stage IDs (from YAML).
         prompt: Original user prompt (may contain multimodal data).
         requires_multimodal_data: Whether to forward multimodal data.
 
     Returns:
         List of dicts, each consumable by the HunyuanImage3 diffusion pipeline.
     """
-    if not engine_input_source:
-        raise ValueError("engine_input_source cannot be empty")
-
-    source_stage_id = engine_input_source[0]
-    if source_stage_id >= len(stage_list):
-        raise IndexError(f"Invalid source stage_id: {source_stage_id}")
-
-    if stage_list[source_stage_id].engine_outputs is None:
-        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
-
-    ar_outputs = stage_list[source_stage_id].engine_outputs
+    ar_outputs = source_outputs
     diffusion_inputs = []
 
     # Normalize prompt to list

vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py

@@ -9,15 +9,12 @@ from vllm_omni.inputs.data import OmniTokensPrompt
 
 
 def ar2dit(
-    stage_list: list[Any],
-    engine_input_source: list[int],
+    source_outputs: list[Any],
     prompts: OmniTokensPrompt | TextPrompt | None = None,
-    requires_multimodal_data: bool = False,  # noqa: ARG001 — interface param, unused for ar2dit
+    _requires_multimodal_data: bool = False,
 ) -> list[OmniTokensPrompt]:
     """Convert AR stage outputs to DiT stage inputs."""
-
-    source_stage_id = engine_input_source[0]
-    ar_outputs = stage_list[source_stage_id].engine_outputs
+    ar_outputs = source_outputs
 
     dit_inputs: list[OmniTokensPrompt] = []
     for ar_output, prompt in zip(ar_outputs, prompts):

vllm_omni/model_executor/stage_input_processors/mimo_audio.py

@@ -182,8 +182,7 @@ def llm2code2wav_async_chunk(
 
 
 def llm2code2wav(
-    stage_list: list[Any],
-    engine_input_source: list[int],
+    source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt | None = None,
     requires_multimodal_data: bool = False,
 ) -> list[OmniTokensPrompt]:
@@ -196,25 +195,13 @@ def llm2code2wav(
     3. Package for code2wav stage
 
     Args:
-        stage_list: List of stage objects
-        engine_input_source: Source stage IDs (typically [1] for talker)
         prompt: Original prompt data
         requires_multimodal_data: Whether multimodal data is required
 
     Returns:
         List of OmniTokensPrompt for code2wav stage
     """
-    if not engine_input_source:
-        raise ValueError("engine_input_source cannot be empty")
-
-    source_stage_id = engine_input_source[0]
-    if source_stage_id >= len(stage_list):
-        raise IndexError(f"Invalid stage_id: {source_stage_id}")
-
-    if stage_list[source_stage_id].engine_outputs is None:
-        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
-
-    talker_outputs = stage_list[source_stage_id].engine_outputs
+    talker_outputs = source_outputs
     code2wav_inputs = []
 
     # Process each talker output

vllm_omni/model_executor/stage_input_processors/omnivoice.py

@@ -10,22 +10,15 @@ from vllm_omni.inputs.data import OmniTokensPrompt
 
 
 def tokens2audio(
-    stage_list: list[Any],
-    engine_input_source: list[int],
-    prompt: OmniTokensPrompt | TextPrompt = None,
-    requires_multimodal_data: bool = True,
+    source_outputs: list[Any],
+    _prompt: OmniTokensPrompt | TextPrompt = None,
+    _requires_multimodal_data: bool = True,
 ):
     """Build stage-1 (decoder) inputs from stage-0 (generator) outputs.
 
     Takes the 8-codebook audio tokens from the generator and packages
     them for the HiggsAudioV2 decoder.
     """
-    source_stage_id = engine_input_source[0]
-    source_outputs = stage_list[source_stage_id].engine_outputs
-
-    if not isinstance(prompt, list):
-        prompt = [prompt]
-
     source_output = source_outputs[0]
     output = source_output.outputs[0]
 

vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py

@@ -10,19 +10,11 @@ TALKER_CODEC_END_TOKEN_ID = 8294
 
 
 def thinker2talker(
-    stage_list,
-    engine_input_source,
+    source_outputs,
     prompt: OmniTokensPrompt | TextPrompt = None,
     requires_multimodal_data: bool = False,
 ):
-    if not engine_input_source:
-        raise ValueError("engine_input_source cannot be empty")
-    source_stage_id = engine_input_source[0]
-    if source_stage_id >= len(stage_list):
-        raise IndexError(f"Invalid stage_id: {source_stage_id}")
-    if stage_list[source_stage_id].engine_outputs is None:
-        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
-    thinker_outputs = stage_list[source_stage_id].engine_outputs
+    thinker_outputs = source_outputs
     talker_inputs = []
     if not isinstance(prompt, list):
         prompt = [prompt]

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -90,36 +90,23 @@ def _ensure_list(x):
     return list(x)
 
 
-def _validate_stage_inputs(stage_list, engine_input_source):
-    if not engine_input_source:
-        raise ValueError("engine_input_source cannot be empty")
-
-    stage_id = engine_input_source[0]
-    if stage_id >= len(stage_list):
-        raise IndexError(f"Invalid stage_id: {stage_id}")
-
-    stage = stage_list[stage_id]
-    if stage.engine_outputs is None:
-        raise RuntimeError(f"Stage {stage_id} has no outputs yet")
-
-    return stage.engine_outputs
-
-
 # =========================
 # PD disaggregation helpers
 # =========================
 
 
-def _get_prefill_stage(stage_list: list[Any], source_stage_id: int) -> Any | None:
-    if source_stage_id <= 0:
+def _get_prefill_multimodal_output(
+    request_id: str,
+    streaming_context: Any | None,
+) -> dict[str, Any] | None:
+    bridge_states = getattr(streaming_context, "bridge_states", None)
+    if not isinstance(bridge_states, dict):
         return None
-    source_stage = stage_list[source_stage_id]
-    if not getattr(source_stage, "is_decode_only", False):
+    by_req = bridge_states.get("pd_prefill_multimodal_output_by_req")
+    if not isinstance(by_req, dict):
         return None
-    prev_stage = stage_list[source_stage_id - 1]
-    if getattr(prev_stage, "is_prefill_only", False) and prev_stage.engine_outputs is not None:
-        return prev_stage
-    return None
+    prefill_mm = by_req.get(request_id)
+    return prefill_mm if isinstance(prefill_mm, dict) else None
 
 
 def _merge_pd_embeddings(
@@ -167,16 +154,6 @@ def _merge_pd_embeddings(
     return merged_emb, merged_hid
 
 
-def _get_prefill_multimodal_output(prefill_stage: Any, output_index: int) -> dict[str, Any] | None:
-    """Return multimodal_output dict from the PD prefill stage for a given batch index."""
-    try:
-        prefill_eos = prefill_stage.engine_outputs
-        prefill_eo = prefill_eos[min(output_index, len(prefill_eos) - 1)]
-        return prefill_eo.outputs[0].multimodal_output
-    except Exception:
-        return None
-
-
 def _resolve_tts_token_embedding(
     key: str,
     *,
@@ -407,8 +384,7 @@ def thinker2talker_async_chunk(
 
 
 def thinker2talker(
-    stage_list: list[Any],
-    engine_input_source: list[int],
+    source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt | None = None,
     requires_multimodal_data: bool = False,
     streaming_context: Any | None = None,
@@ -425,23 +401,17 @@ def thinker2talker(
     decode-stage generated embeddings before handing off to the talker.
 
     Args:
-        stage_list: List of stage objects
-        engine_input_source: Source stage IDs (typically [0] for thinker)
         prompt: Original prompt data
         requires_multimodal_data: Whether multimodal data is required
 
     Returns:
         List of OmniTokensPrompt for talker stage
     """
-    thinker_outputs = _validate_stage_inputs(stage_list, engine_input_source)
+    thinker_outputs = source_outputs
     talker_inputs: list[OmniTokensPrompt] = []
 
     device = torch.device(current_platform.device_type)
 
-    # PD disaggregation: look up the preceding prefill stage (if any)
-    source_stage_id = engine_input_source[0]
-    prefill_stage = _get_prefill_stage(stage_list, source_stage_id)
-
     # Process each thinker output
     for i, thinker_output in enumerate(thinker_outputs):
         output = thinker_output.outputs[0]
@@ -478,8 +448,7 @@ def thinker2talker(
         thinker_hid = hid_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:]
 
         prefill_mm: dict[str, Any] | None = None
-        if prefill_stage is not None:
-            prefill_mm = _get_prefill_multimodal_output(prefill_stage, i)
+        prefill_mm = _get_prefill_multimodal_output(req_id, streaming_context)
 
         if prefill_mm is not None:
             expected_total = len(prompt_token_ids) + len(output_ids)
@@ -610,10 +579,9 @@ def talker2code2wav_async_chunk(
 
 
 def talker2code2wav(
-    stage_list: list[Any],
-    engine_input_source: list[int],
-    prompt: OmniTokensPrompt | TextPrompt | None = None,
-    requires_multimodal_data: bool = False,
+    source_outputs: list[Any],
+    _prompt: OmniTokensPrompt | TextPrompt | None = None,
+    _requires_multimodal_data: bool = False,
     streaming_context: Any | None = None,
 ) -> list[OmniTokensPrompt]:
     """
@@ -625,15 +593,10 @@ def talker2code2wav(
     3. Package for code2wav stage
 
     Args:
-        stage_list: List of stage objects
-        engine_input_source: Source stage IDs (typically [1] for talker)
-        prompt: Original prompt data
-        requires_multimodal_data: Whether multimodal data is required
-
     Returns:
         List of OmniTokensPrompt for code2wav stage
     """
-    talker_outputs = _validate_stage_inputs(stage_list, engine_input_source)
+    talker_outputs = source_outputs
     code2wav_inputs: list[OmniTokensPrompt] = []
     # Process each talker output
     for i, talker_output in enumerate(talker_outputs):

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -20,16 +20,14 @@ logger = init_logger(__name__)
 
 
 def talker2code2wav(
-    stage_list: list[Any],
-    engine_input_source: list[int],
+    source_outputs: list[Any],
     prompt: Any = None,
-    requires_multimodal_data: bool = False,
+    _requires_multimodal_data: bool = False,
 ) -> list[Any]:
     """Non-async: collect all talker codes, then pass to code2wav at once."""
     from vllm_omni.inputs.data import OmniTokensPrompt
-    from vllm_omni.model_executor.stage_input_processors.qwen3_omni import _validate_stage_inputs
 
-    talker_outputs = _validate_stage_inputs(stage_list, engine_input_source)
+    talker_outputs = source_outputs
     code2wav_inputs: list[OmniTokensPrompt] = []
     for i, talker_output in enumerate(talker_outputs):
         if not talker_output.finished:

vllm_omni/model_executor/stage_input_processors/voxcpm.py

@@ -42,24 +42,10 @@ def _coerce_finished_flag(value: Any) -> bool:
 
 
 def latent2vae(
-    stage_list: list[Any],
-    engine_input_source: list[int],
-    prompt: OmniTokensPrompt | TextPrompt | None = None,
-    requires_multimodal_data: bool = False,
+    source_outputs: list[Any],
+    _prompt: OmniTokensPrompt | TextPrompt | None = None,
+    _requires_multimodal_data: bool = False,
 ) -> list[OmniTokensPrompt]:
-    del prompt, requires_multimodal_data
-
-    if not engine_input_source:
-        raise ValueError("engine_input_source cannot be empty")
-
-    source_stage_id = engine_input_source[0]
-    if source_stage_id >= len(stage_list):
-        raise IndexError(f"Invalid stage_id: {source_stage_id}")
-
-    source_outputs = stage_list[source_stage_id].engine_outputs
-    if source_outputs is None:
-        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
-
     vae_inputs: list[OmniTokensPrompt] = []
     for source_output in source_outputs:
         output = source_output.outputs[0]

vllm_omni/model_executor/stage_input_processors/voxtral_tts.py

@@ -10,22 +10,12 @@ logger = init_logger(__name__)
 
 
 def generator2tokenizer(
-    stage_list,
-    engine_input_source,
-    prompt: OmniTokensPrompt | TextPrompt = None,
-    requires_multimodal_data: bool = False,
+    source_outputs: list[Any],
+    _prompt: OmniTokensPrompt | TextPrompt = None,
+    _requires_multimodal_data: bool = False,
 ):
-    if not engine_input_source:
-        raise ValueError("engine_input_source cannot be empty")
-    source_stage_id = engine_input_source[0]
-    if source_stage_id >= len(stage_list):
-        raise IndexError(f"Invalid stage_id: {source_stage_id}")
-    if stage_list[source_stage_id].engine_outputs is None:
-        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
-
-    generator_outputs = stage_list[source_stage_id].engine_outputs
     tokenizer_inputs = []
-    for generator_output in generator_outputs:
+    for generator_output in source_outputs:
         output = generator_output.outputs[0]
         audio_tokens = torch.cat(output.multimodal_output["audio"], dim=-1).flatten().detach().cpu().tolist()
         tokenizer_inputs.append(