[Core] Route stages through configured pipeline transitions

docs/configuration/README.md

@@ -14,6 +14,8 @@ For a specific example, see the [Qwen2.5-Omni deploy config](gh-file:vllm_omni/d
 
 For an introduction, see [Pipeline and deploy configurations](./stage_configs.md).
 
+For developer-owned control flow, see [Static stage transitions](./stage_routing.md).
+
 ## Memory Configuration
 
 - **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance

docs/configuration/stage_routing.md

@@ -0,0 +1,97 @@
+# Static stage transitions
+
+Model developers can set `PipelineConfig.stage_transitions` to an explicit tuple
+of `(source_stage_id, target_stage_id)` pairs. This is part of the registered,
+developer-owned pipeline topology. Deployment YAML continues to select hardware,
+replicas and connector backends.
+
+With `stage_transitions=None`, requests follow the existing numeric sequence
+`0 -> 1 -> ...`. Explicit transitions replace that sequence completely; omitted
+edges do not gain a sequential fallback.
+
+For example, a pipeline with three registered stages can run `0 -> 2`:
+
+```python
+PipelineConfig(
+    model_type="example",
+    stages=(
+        StagePipelineConfig(stage_id=0, model_stage="encoder"),
+        StagePipelineConfig(stage_id=1, model_stage="unused"),
+        StagePipelineConfig(
+            stage_id=2,
+            model_stage="decoder",
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="audio",
+        ),
+    ),
+    stage_transitions=((0, 2),),
+)
+```
+
+The explicit route starts at stage 0 and has at most one successor and one
+predecessor per stage. Stage IDs must remain contiguous array indices, even when
+the route skips stages or orders them differently. Duplicate edges, cycles,
+self edges, unknown IDs and disconnected edges are rejected. The last stage of
+the route must declare `final_output=True`. An empty tuple explicitly selects
+stage 0 alone, so stage 0 must declare a final output in that case.
+
+`input_sources` still describes data dependencies. Every input source of an
+active stage must precede that stage on the route. Changing a transition cannot
+make a processor accept an incompatible payload or omit a required dependency.
+
+Configuration extensions use the actual route tail. For example,
+`--forced-aligner` appends its pooling stage after that tail, uses it as the
+audio input source and extends explicit transitions. A default `None` route
+remains `None`. Typed and compatibility launch configurations use the same
+extended pipeline.
+
+Requests execute a prefix of the configured route ending at `final_stage_id`.
+If the model bypasses its initial stages for a request, execution starts at that
+request's admitted entry stage within the prefix.
+Requested output stages must occur in that prefix and declare final outputs;
+they must include the endpoint. Streaming updates retain the admitted endpoint
+and output stages. Public sampling arrays retain one slot per registered stage,
+including inactive stages; IPC admission checks coverage through the highest
+stage ID in the prefix.
+Output modalities select the last matching stage in route order. For example,
+under `0 -> 2 -> 1`, stage 1 can finish a request after stage 2. Stage sampling
+parameters remain indexed by stage ID.
+Worker connector endpoints remain fixed for the full route. A request ending
+early can leave unread outgoing payloads, which follow the backend cleanup
+contract below; this feature does not prune all data-plane sends per request.
+An explicitly requested registered output modality is rejected if all stages
+that produce it are inactive. Omitting modalities selects from active outputs.
+
+Async-chunk prewarming follows the same prefix and binds each upstream sender
+through its StagePool before advertising its address to a receiver. The later
+bridge submission reuses that binding. A lost bound replica fails its requests
+instead of moving their senders after the receivers have been configured.
+Connector payload send/receive endpoints use the projected transition,
+including on headless workers. Backend projection respects `default_connector`;
+request keys and chunk counters retain their existing per-request semantics.
+Stages fed by an orchestrator bridge receive their real payload before dispatch.
+Control edges and backend edges serve different purposes: a bridge-fed stage can
+have an outgoing connector without an incoming one. A middle stage with both
+connector directions requires the same backend on both edges, because one
+worker owns one connector. Absent payload backends retain the existing local
+IPC behavior; choosing a backend suitable for the deployment remains necessary.
+
+Omitted stages receive no request submissions or prewarming. They are still
+initialized, so this feature does not reduce startup memory or process counts.
+Normal and empty terminal outputs share the same completion condition: all
+selected outputs have finished, and non-streaming async requests also wait for
+every submitted stage. The endpoint's terminal output is retained until then,
+so delayed upstream usage and stop reasons arrive before completion.
+Cleanup retains the existing per-replica batched, best-effort resource release.
+Abort acknowledgement does not certify physical release of transport buffers.
+Changing a route does not strengthen a backend's release guarantees. Legacy V1
+full-payload and native KV transfer retain their existing cleanup paths; NIXL
+ownership changes are reviewed separately.
+
+This first implementation supports static routes for turn-based requests.
+Branches, joins, loops and token-dependent decisions need a separate execution
+contract. Non-sequential duplex deployments are rejected before plugin loading;
+their model-specific session contracts need a separate extension. Existing
+native AR-to-DiT KV deployment restrictions still apply.
+Streaming-input requests with a native MRv2 downstream receiver remain unsupported.

docs/design/module/engine_orchestration.md

@@ -90,6 +90,11 @@ This document owns `AsyncOmniEngine`, `Orchestrator`, request-state creation,
 cross-stage routing, output ordering, companion tracking, control/RPC
 correlation, cancellation propagation, and terminal-state convergence.
 
+After configuration resolution, the engine retains the effective `PipelineConfig`,
+including injected stages such as ForcedAligner. Stage startup and orchestration
+read transitions from that same configuration; the registered pipeline is not
+modified.
+
 It does not own stage placement, replica selection, stage-process startup,
 payload schemas, public protocol rendering, connector transport, or semantic
 error classification. Those responsibilities belong to the stage runtime,
@@ -132,8 +137,14 @@ child must drop `parent_requests` once no children remain. EngineCore control
 RPCs (`sleep` / `wake_up` / `pause_scheduler` / `resume_scheduler`) must
 propagate worker exceptions rather than returning `{"supported": False}`.
 When a request has multiple final output stages, only the last abort message
-is request-terminal. Output-processor abort state is committed only after the
-physical EngineCore abort succeeds.
+is request-terminal. Output-processor abort state is committed after the stage
+client's abort call succeeds or a natural terminal confirms stage completion.
+For AR clients, a successful call confirms message submission, not EngineCore
+completion or physical release of transport buffers.
+An abort failure still attempts the other stages and replicas. The request
+remains closed to admission while stages with failed abort calls drain; their natural
+terminals permit cleanup without publishing a normal completion. Correlated
+abort failure also reports a request-scoped error to unblock its output stream.
 Changes that cross into stage runtime or public error behavior require review
 from that contract's owners.
 

tests/config/test_forced_aligner_injection.py

@@ -9,8 +9,11 @@
 from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY
 from vllm_omni.config.stage_config import (
     DeployConfig,
+    PipelineConfig,
     StageExecutionType,
+    StagePipelineConfig,
 )
+from vllm_omni.config.stage_routing import StageRouting
 from vllm_omni.engine.stage_init_utils import (
     build_engine_args_dict_from_omni_stage_config,
     extract_stage_metadata_from_omni_stage_config,
@@ -80,6 +83,51 @@ def test_inject_noop_without_forced_aligner():
     assert len(ext_deploy.stages) == 0
 
 
+@pytest.mark.parametrize(
+    ("transitions", "order"),
+    [
+        (None, (0, 1, 2)),
+        (((0, 1), (1, 2)), (0, 1, 2)),
+        (((0, 2),), (0, 2)),
+        (((2, 1), (0, 2)), (0, 2, 1)),
+        ((), (0,)),
+    ],
+)
+def test_injected_aligner_extends_actual_route_tail(transitions, order):
+    from vllm_omni.config.config_factory import StageConfigFactory
+    from vllm_omni.engine.omni_engine_base import OmniEngineBase
+
+    pipeline = PipelineConfig(
+        model_type="test",
+        stages=tuple(StagePipelineConfig(i, "audio", final_output=True, final_output_type="audio") for i in range(3)),
+        stage_transitions=transitions,
+    )
+    deploy = DeployConfig()
+    overrides = {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"}
+    extended, extended_deploy = inject_forced_aligner_stage(pipeline, deploy, overrides)
+
+    assert extended.stages[-1].input_sources == (order[-1],)
+    route = StageRouting.from_transitions(len(extended.stages), extended.stage_transitions)
+    assert route.stage_order == (*order, 3)
+    assert pipeline.stage_transitions == transitions
+    assert len(pipeline.stages) == 3 and deploy.stages == []
+    assert (extended.stage_transitions is None) == (transitions is None)
+
+    typed = VllmOmniConfig.from_pipeline_config(extended, user_deploy_config=extended_deploy)
+    legacy = StageConfigFactory._resolve_legacy_from_registry(pipeline, overrides, user_deploy_config=deploy)
+    assert typed.pipeline_config == legacy.pipeline_config == extended
+    assert tuple(typed.stage_configs[-1].input_sources) == (order[-1],)
+    assert tuple(legacy.stage_configs[-1].input_sources) == (order[-1],)
+
+    engine = object.__new__(OmniEngineBase)
+    engine.pipeline_config = pipeline
+    engine._set_pipeline_runtime_config(legacy.pipeline_config, None)
+    assert engine.pipeline_config is legacy.pipeline_config
+    assert StageRouting.from_transitions(
+        len(legacy.stage_configs), engine.pipeline_config.stage_transitions
+    ).stage_order == (*order, 3)
+
+
 @pytest.mark.parametrize("async_chunk", [True, False])
 @pytest.mark.parametrize("cli_async_chunk", [None, True, False])
 def test_aligner_uses_completed_audio_without_disabling_upstream_chunks(async_chunk, cli_async_chunk):

tests/config/test_stage_routing.py

@@ -0,0 +1,181 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig
+from vllm_omni.config.stage_routing import StageRouting
+from vllm_omni.distributed.omni_connectors.utils.config import (
+    ConnectorSpec,
+    OmniTransferConfig,
+    get_stage_connector_peer,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("num_stages", [1, 2, 5])
+def test_default_route_preserves_sequential_prefixes(num_stages):
+    route = StageRouting.from_transitions(num_stages)
+    assert route.stage_order == tuple(range(num_stages))
+    for final in range(num_stages):
+        assert route.path_to(final) == tuple(range(final + 1))
+        assert route.next_stage(final, final) is None
+
+
+@pytest.mark.parametrize(
+    ("transitions", "order"),
+    [(((0, 2),), (0, 2)), (((2, 1), (0, 2)), (0, 2, 1)), ((), (0,))],
+)
+def test_explicit_route_has_no_implicit_edges(transitions, order):
+    route = StageRouting.from_transitions(3, transitions)
+    assert route.stage_order == order
+    assert route.next_stage(order[-1], order[-1]) is None
+
+
+@pytest.mark.parametrize(
+    ("transitions", "error"),
+    [
+        (((0, 3),), "invalid stage"),
+        (((-1, 1),), "invalid stage"),
+        (((0, True),), "invalid stage"),
+        (((0, 0),), "Self transition"),
+        (((0, 1), (1, 0)), "cycle"),
+        (((0, 1), (0, 2)), "branching"),
+        (((0, 1), (0, 1)), "duplicate"),
+        (((0, 2), (1, 2)), "multiple incoming"),
+        (((1, 2),), "unreachable"),
+    ],
+)
+def test_invalid_topology_is_rejected(transitions, error):
+    with pytest.raises(ValueError, match=error):
+        StageRouting.from_transitions(3, transitions)
+
+
+def test_unreachable_request_endpoint_is_rejected():
+    route = StageRouting.from_transitions(3, ((0, 2),))
+    with pytest.raises(ValueError, match="not reachable"):
+        route.path_to(1)
+    assert route.next_stage(1, 2) is None
+
+
+@pytest.mark.parametrize("transitions", [[(0, 2)], ([0, 2],), ((0,),), ((0, 1, 2),)])
+def test_mutable_or_malformed_transition_configuration_is_rejected(transitions):
+    with pytest.raises(ValueError, match="immutable tuple"):
+        StageRouting.from_transitions(3, transitions)
+
+
+def test_pipeline_accepts_skipped_stage_with_satisfied_dependencies():
+    pipeline = PipelineConfig(
+        model_type="test",
+        stages=(
+            StagePipelineConfig(0, "entry"),
+            StagePipelineConfig(1, "inactive", input_sources=(0,)),
+            StagePipelineConfig(2, "output", input_sources=(0,), final_output=True),
+        ),
+        stage_transitions=((0, 2),),
+    )
+    assert pipeline.get_validation_errors() == []
+
+
+@pytest.mark.parametrize("sources", [(1,), (2,)])
+def test_pipeline_rejects_unsatisfied_route_dependency(sources):
+    with pytest.raises(ValueError, match="must precede"):
+        PipelineConfig(
+            model_type="test",
+            stages=(
+                StagePipelineConfig(0, "entry"),
+                StagePipelineConfig(1, "inactive"),
+                StagePipelineConfig(2, "output", input_sources=sources, final_output=True),
+            ),
+            stage_transitions=((0, 2),),
+        )
+
+
+def test_pipeline_rejects_route_without_terminal_output():
+    with pytest.raises(ValueError, match="last stage"):
+        PipelineConfig(
+            model_type="test",
+            stages=(StagePipelineConfig(0, "entry"), StagePipelineConfig(1, "output", final_output=True)),
+            stage_transitions=(),
+        )
+
+
+@pytest.mark.parametrize("stage_id", [3, 1.0, True])
+def test_pipeline_rejects_ids_that_cannot_index_stage_arrays(stage_id):
+    with pytest.raises(ValueError, match="contiguous IDs"):
+        PipelineConfig(
+            model_type="test",
+            stages=(StagePipelineConfig(0, "entry"), StagePipelineConfig(stage_id, "output", final_output=True)),
+            stage_transitions=((0, 1),),
+        )
+
+
+@pytest.mark.parametrize(("stage_id", "source", "target"), [(0, None, 2), (2, 0, 1), (1, 2, None)])
+def test_connector_projection_preserves_both_route_directions(stage_id, source, target):
+    from vllm_omni.engine.stage_init_utils import get_stage_connector_spec
+
+    route = StageRouting.from_transitions(3, ((0, 2), (2, 1)))
+    spec = get_stage_connector_spec(None, stage_id, True, route)
+    config = SimpleNamespace(stage_connector_config=spec)
+    assert get_stage_connector_peer(config, stage_id, "from_stage") == source
+    assert get_stage_connector_peer(config, stage_id, "to_stage") == target
+
+
+def test_legacy_connector_endpoints_keep_numeric_defaults():
+    config = SimpleNamespace(stage_connector_config={})
+    assert get_stage_connector_peer(config, 1, "from_stage") == 0
+    assert get_stage_connector_peer(config, 1, "to_stage") == 2
+
+
+@pytest.mark.parametrize("value", ["", " ", 1.5])
+def test_legacy_unset_connector_target_keeps_sequential_default(value):
+    config = SimpleNamespace(stage_connector_config={"to_stage": value})
+    assert get_stage_connector_peer(config, 1, "to_stage") == 2
+
+
+def test_connector_projection_honors_default_without_mutating_deployment():
+    from vllm_omni.engine.stage_init_utils import get_stage_connector_spec
+
+    backend = ConnectorSpec("NixlConnector", {"zmq_port": 50051})
+    config = OmniTransferConfig({("0", "1"): ConnectorSpec("SharedMemoryConnector")}, backend)
+    original = dict(config.connectors)
+    route = StageRouting.from_transitions(3, ((0, 2), (2, 1)))
+    spec = get_stage_connector_spec(config, 2, True, route)
+    assert spec["name"] == "NixlConnector"
+    assert (spec["from_stage"], spec["to_stage"]) == (0, 1)
+    assert spec["extra"]["outgoing"]["from_stage"] == 2
+    assert spec["extra"]["outgoing"]["to_stage"] == 1
+    assert config.connectors == original and config.default_connector is backend
+    assert backend.extra == {"zmq_port": 50051}
+
+
+def test_connector_projection_preserves_bridge_fed_stage():
+    from vllm_omni.engine.stage_init_utils import get_stage_connector_spec
+
+    config = OmniTransferConfig({("2", "1"): ConnectorSpec("SharedMemoryConnector")})
+    route = StageRouting.from_transitions(3, ((0, 2), (2, 1)))
+    spec = get_stage_connector_spec(config, 2, True, route)
+    assert spec["extra"]["role"] == "sender"
+    assert (spec["from_stage"], spec["to_stage"]) == (0, 1)
+
+
+@pytest.mark.parametrize(
+    "backends", [("SharedMemoryConnector", "NixlConnector"), ("NixlConnector", "SharedMemoryConnector")]
+)
+def test_middle_stage_rejects_incompatible_connector_backends(backends):
+    from vllm_omni.engine.stage_init_utils import get_stage_connector_spec
+
+    config = OmniTransferConfig({("0", "2"): ConnectorSpec(backends[0]), ("2", "1"): ConnectorSpec(backends[1])})
+    route = StageRouting.from_transitions(3, ((0, 2), (2, 1)))
+    with pytest.raises(ValueError, match="same connector backend"):
+        get_stage_connector_spec(config, 2, True, route)
+
+
+def test_inactive_stage_has_no_payload_peers():
+    from vllm_omni.engine.stage_init_utils import get_stage_connector_spec
+
+    spec = get_stage_connector_spec(None, 1, True, StageRouting.from_transitions(3, ((0, 2),)))
+    assert spec == {"from_stage": None, "to_stage": None, "extra": {"role": "sender"}}

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -116,12 +116,14 @@ def test_chunk_registration_ready_and_terminal_lifecycle():
     coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
     req = _make_request("internal", status=RequestStatus.WAITING)
     req.external_req_id = "external"
+    req.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 51424}
     waiting = MockQueue([req])
     coord.process_pending_chunks(waiting, [], set(), set())
     assert req.status == RequestStatus.WAITING_FOR_CHUNK
     [handle] = coord.pending_chunk_registrations
     assert isinstance(handle, OmniChunkRecvHandle)
     assert (handle.request_id, handle.external_req_id) == ("internal", "external")
+    assert handle.payload_sender_info == req.payload_sender_info
     coord.restore_queues(waiting, [])
     coord.process_pending_chunks(waiting, [], {"internal"}, set())
     assert req.status == RequestStatus.WAITING