[Bugfix] Fix multi-replica stage identity and clean up redundant deploy config (#4410)

Signed-off-by: ZhengWG <zwg0606@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

examples/online_serving/qwen3_omni/README.md

@@ -24,12 +24,15 @@ 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:
+For a 3x-GPU multi-replica layout (talker/code2wav scale-out on cuda:1,2),
+use `--stage-overrides` on top of the default config:
 
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
-    --deploy-config vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml
+    --stage-overrides '{
+        "1": {"num_replicas": 2, "devices": "1,2"},
+        "2": {"num_replicas": 2, "devices": "1,2"}
+    }'
 ```
 
 ### Launch individual stages (stage-based CLI)

vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml

@@ -1,109 +0,0 @@
-# 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:
-      initial_codec_chunk_frames: 4
-      codec_chunk_frames: 25
-      codec_left_context_frames: 25
-
-stages:
-  - stage_id: 0
-    max_num_batched_tokens: 32768
-    max_num_seqs: 64
-    gpu_memory_utilization: 0.9
-    trust_remote_code: true
-    enable_prefix_caching: false
-    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
-    max_num_batched_tokens: 32768
-    max_num_seqs: 64
-    gpu_memory_utilization: 0.6
-    trust_remote_code: true
-    enable_prefix_caching: false
-    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
-    max_num_batched_tokens: 65536
-    max_num_seqs: 64
-    gpu_memory_utilization: 0.1
-    enforce_eager: true
-    trust_remote_code: true
-    enable_prefix_caching: false
-    async_scheduling: false
-    enable_chunked_prefill: false
-    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
-        max_num_batched_tokens: 8192
-        devices: "0,1"
-      - stage_id: 1
-        gpu_memory_utilization: 0.6
-        max_num_batched_tokens: 8192
-        devices: "2"
-      - stage_id: 2
-        gpu_memory_utilization: 0.3
-        devices: "2"
-
-  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/engine/async_omni_engine.py

@@ -921,6 +921,8 @@ class AsyncOmniEngine:
                                         vllm_config=vllm_config,
                                         executor_class=executor_class,
                                         log_stats=self._log_stats,
+                                        omni_stage_id=plan.metadata.stage_id,
+                                        omni_replica_id=plan.replica_id,
                                     )
                                 logger.info(
                                     "[AsyncOmniEngine] Stage %s engine launch started",

vllm_omni/engine/stage_engine_core_proc.py

@@ -218,6 +218,8 @@ def spawn_stage_core(
     vllm_config: VllmConfig,
     executor_class: type[Executor],
     log_stats: bool = False,
+    omni_stage_id: int | None = None,
+    omni_replica_id: int = 0,
 ) -> tuple[EngineZmqAddresses, BaseProcess, str]:
     """Spawn a *StageEngineCoreProc* subprocess without performing the handshake.
 
@@ -241,6 +243,8 @@ def spawn_stage_core(
             "log_stats": log_stats,
             "dp_rank": 0,
             "local_dp_rank": 0,
+            "omni_stage_id": omni_stage_id,
+            "omni_replica_id": omni_replica_id,
         },
     )
     proc.start()