[Pinoeer-kingxi]

Purpose

Add PipelineConfig.stage_transitions for the routing work in #8264 Item 14 / RFC #8740. None preserves sequential execution; explicit transitions support compatible skip and reorder chains, such as 0 → 2 and 0 → 2 → 1.

• Use one route for startup, forwarding, prewarming, worker endpoints, output selection and completion. Keep sampling arrays indexed by stage ID; append ForcedAligner after the actual tail.
• Bind producers before prewarming receivers. If a bound sender is lost, fail its request instead of switching replicas.
• Stop admission after cancellation, retain failed-abort state while stages drain, and discard late bridge validation errors after ID reuse. Wait for selected outputs and all submitted non-streaming async stages.

Static turn-based chains only; skipped stages still initialize. Branches, joins, cycles and non-sequential duplex are unsupported. Deployment YAML owns placement and backends. Resource release remains best-effort; abort acknowledgement is not physical buffer release (#8744).

Configuration guide.

Test Plan

Head 539648e7; vLLM 0.31.0, Python 3.12.4, PyTorch 2.13.0+cu130.

CPU command, including the known failing input-cache cases:

CUDA_VISIBLE_DEVICES='' OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 \
HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 PYTHONDONTWRITEBYTECODE=1 \
python -m pytest \
  $(git diff --name-only 60984a2a 539648e794954acfe549ef26bbefe32f4d960ac1 -- 'tests/**/*.py') \
  tests/engine/test_orchestrator_input_cache.py \
  -m 'core_model and cpu' --run-level=core_model -q -o addopts='' -p no:cacheprovider

Test Result

• 137 CPU cases passed at 539648e7: failed/in-flight aborts, late bridge errors and ID reuse. An earlier candidate integrated with abff9664 passed 977 cases, covering default/skip/reorder routes, both loops, headless/multi-API projection and real SHM transfer. Counts overlap; this is not a full current-head suite pass.
• Ruff/format, SPDX, import/CUDA rules, test marks, Markdown and diff checks passed. mypy had no new diagnostics versus the tested base; existing diagnostics remain.
• Earlier vLLM 0.31 A6000 smokes: Qwen3-TTS-12Hz-0.6B-Base + ForcedAligner-0.6B, BF16/eager, seed 42, max 96 tokens. V1 sync 0 → 1 → 2 and MRv2 async 0 → 2 → 3 produced finite audio and nine ordered word timestamps, completed audio-only prefixes, and passed admission-time cancellation, ID reuse and binding release.

Known gaps

• Input-cache forwarding tests fail because their fixture omits the required next_logical argument. The call-site fix is not in this head and is needed before merge.
• Current-main vLLM 0.32 GPU validation is pending. Earlier smokes do not establish alignment accuracy, full WER, performance or RDMA.

Reproduce the two model route smokes

Use one idle 48 GiB CUDA GPU and cached weights. Save the Python block as validate_route.py at the repo root. Set ROUTING_ALIGNER_MODEL to the local Qwen3-ForcedAligner-0.6B snapshot. The script pins the TTS revision and uses a repository WAV. These smokes cancel after admission, not during a native PROC READ.

export CUDA_VISIBLE_DEVICES=0  # Select one idle physical GPU.
export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1
export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 PYTHONDONTWRITEBYTECODE=1
export ROUTING_REPO="$PWD" ROUTING_ALIGNER_MODEL=/path/to/Qwen3-ForcedAligner-0.6B
export ROUTING_GPU_WORK=$(mktemp -d /tmp/omni-routing.XXXXXX)
export TMPDIR="$ROUTING_GPU_WORK"
export TORCHINDUCTOR_CACHE_DIR="$ROUTING_GPU_WORK/inductor"
export TRITON_CACHE_DIR="$ROUTING_GPU_WORK/triton"
export FLASHINFER_WORKSPACE_BASE="$ROUTING_GPU_WORK"
python validate_route.py --route default --runner v1
python validate_route.py --route skip --runner v2 --async-chunk

"""Real TTS -> injected aligner, asserting route, finite audio and word timestamps."""

import argparse, asyncio, copy, json, os, time
from dataclasses import replace
from pathlib import Path
import torch, yaml
from vllm.sampling_params import RequestOutputKind, SamplingParams
from tests.e2e.offline_inference.test_qwen3_tts_base import (
    MODEL,
    REF_AUDIO_URL,
    REF_TEXT,
    get_prompt,
)
from tests.helpers.runtime import OmniRunner
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES
from vllm_omni.config.stage_routing import StageRouting
from vllm_omni.entrypoints.async_omni import AsyncOmni
from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE
from vllm_omni.utils.forced_aligner import extract_word_timestamps
from vllm_omni.utils.qwen3_force_align_processor import segment_words


async def main(args):
    repo = Path(os.environ["ROUTING_REPO"])
    work = Path(os.environ["ROUTING_GPU_WORK"])
    pipeline = QWEN3_TTS_PIPELINE
    deploy = yaml.safe_load((repo / "vllm_omni/deploy/qwen3_tts.yaml").read_text())
    deploy.update(model_runner=args.runner, async_chunk=args.async_chunk)
    if args.route == "skip":
        decoder = pipeline.stages[1]
        pipeline = replace(
            pipeline,
            stages=(
                pipeline.stages[0],
                replace(decoder, final_output=False),
                replace(decoder, stage_id=2),
            ),
            stage_transitions=((0, 2),),
        )
        deploy["stages"][0]["output_connectors"] = {
            "to_stage_2": "connector_of_shared_memory"
        }
        active = copy.deepcopy(deploy["stages"][1])
        active["stage_id"] = 2
        deploy["stages"].append(active)
        deploy["stages"][1].pop("input_connectors")
    OMNI_PIPELINES["qwen3_tts"] = pipeline
    for stage in deploy["stages"]:
        stage.update(
            devices="0",
            enforce_eager=True,
            async_scheduling=False,
            max_num_seqs=2,
            gpu_memory_utilization=0.15,
            kv_cache_memory_bytes=1073741824,
        )
    config = work / f"tts-aligner-{args.route}-{args.runner}.yaml"
    config.write_text(yaml.safe_dump(deploy, sort_keys=False))
    align_config = work / "forced-aligner-round25.yaml"
    align_config.write_text(
        yaml.safe_dump(
            {
                "forced_aligner": {
                    "model": os.environ["ROUTING_ALIGNER_MODEL"],
                    "architecture": "Qwen3ASRForcedAlignerForTokenClassification",
                    "runner": "pooling",
                    "pooling_task": "token_classify",
                    "dtype": "bfloat16",
                    "gpu_memory_utilization": 0.15,
                    "max_model_len": 4096,
                    "trust_remote_code": True,
                }
            }
        )
    )
    engine = AsyncOmni(
        model=MODEL,
        deploy_config=str(config),
        forced_aligner_config=str(align_config),
        forced_aligner_device="0",
        revision="5d83992436eae1d760afd27aff78a71d676296fc",
        stage_init_timeout=600,
        init_timeout=900,
        enforce_eager=True,
        async_scheduling=False,
        max_num_seqs=2,
        kv_cache_memory_bytes=1073741824,
    )
    try:
        actual = engine.engine.pipeline_config
        order = StageRouting.from_transitions(
            len(actual.stages), actual.stage_transitions
        ).stage_order
        expected = (0, 2, 3) if args.route == "skip" else (0, 1, 2)
        assert order == expected, (order, expected)
        assert actual.stages[-1].input_sources == (expected[-2],)
        builder = object.__new__(OmniRunner)
        builder.model_name = MODEL
        builder._prompt_len_estimate_cache = {}
        text = get_prompt()
        prompt = builder.get_omni_inputs(
            text,
            modalities=["audio"],
            mm_processor_kwargs={
                "task_type": "Base",
                "ref_audio": REF_AUDIO_URL,
                "ref_text": REF_TEXT,
                "max_new_tokens": 96,
            },
        )[0]
        params = copy.deepcopy(engine.default_sampling_params_list)
        for item in params:
            if isinstance(item, SamplingParams):
                item.output_kind = RequestOutputKind.CUMULATIVE
        params[0].max_tokens = 96
        params[0].seed = 42
        outputs = []
        started = time.perf_counter()
        async for output in engine.generate(
            prompt,
            request_id="item14-aligner",
            sampling_params_list=params,
            output_modalities=["audio", "timestamps"],
        ):
            assert not output.error, output
            outputs.append(output)
        audio = [output for output in outputs if output.final_output_type == "audio"]
        align = [
            output for output in outputs if output.final_output_type == "timestamps"
        ]
        assert audio and align and outputs[-1].finished, (
            len(audio),
            len(align),
            outputs,
        )
        chunk = audio[-1].outputs[0].multimodal_output["audio"]
        if isinstance(chunk, list):
            chunk = chunk[0]
        assert chunk.numel() > 0 and torch.isfinite(chunk).all()
        timestamps = extract_word_timestamps(align[-1], text)
        assert timestamps and len(timestamps) == len(segment_words(text)), timestamps
        prev = 0
        for item in timestamps:
            assert prev <= item["start_ms"] <= item["end_ms"], timestamps
            prev = item["end_ms"]
        assert prev > 0, timestamps
        assert not engine.request_states
        assert all(
            not pool._request_bindings and not pool._affinity
            for pool in engine.engine.stage_pools
        )
        alignment_elapsed = time.perf_counter() - started
        long_params = copy.deepcopy(params)
        long_params[0].max_tokens = 1024
        long_params[0].ignore_eos = True
        long_params[0].stop_token_ids = []

        async def cancellable():
            async for output in engine.generate(
                prompt,
                request_id="item14-aligner-cancel",
                sampling_params_list=long_params,
                output_modalities=["audio", "timestamps"],
            ):
                assert not output.error, output

        task = asyncio.create_task(cancellable())
        deadline = time.monotonic() + 30
        while not engine.engine.stage_pools[0]._request_bindings:
            assert not task.done() and time.monotonic() < deadline
            await asyncio.sleep(0.05)
        await engine.abort("item14-aligner-cancel", timeout=45)
        await asyncio.wait_for(task, 30)
        assert not engine.request_states
        assert all(
            not pool._request_bindings and not pool._affinity
            for pool in engine.engine.stage_pools
        )
        audio_only = []
        async for output in engine.generate(
            prompt,
            request_id="item14-aligner-cancel",
            sampling_params_list=params,
            output_modalities=["audio"],
        ):
            assert not output.error and output.final_output_type == "audio", output
            audio_only.append(output)
        assert audio_only and audio_only[-1].finished and not engine.request_states
        assert all(
            not pool._request_bindings and not pool._affinity
            for pool in engine.engine.stage_pools
        )
        result = {
            "route": order,
            "runner": args.runner,
            "async_chunk": args.async_chunk,
            "samples": chunk.numel(),
            "timestamps": timestamps,
            "elapsed_seconds": alignment_elapsed,
            "audio_only_prefix": True,
            "cancel_and_id_reuse": True,
            "all_bindings_released": True,
            "finite_audio": True,
            "model": "Qwen3-ForcedAligner-0.6B",
        }
        print("REAL ALIGNER ROUTE PASS", json.dumps(result), flush=True)
    finally:
        engine.shutdown()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--route", choices=("default", "skip"), required=True)
    parser.add_argument("--runner", choices=("v1", "v2"), default="v1")
    parser.add_argument("--async-chunk", action="store_true")
    asyncio.run(main(parser.parse_args()))

Contributor intent

1. What outcome do you want to achieve?

Let model pipelines declare a static execution chain, including compatible skip and reorder paths, while preserving sequential routing by default.

2. What is the highest-leverage variable for that outcome?

Configuration, initialization, prewarming, forwarding and completion must follow the same validated route.

3. How will this work move that variable?

Add optional PipelineConfig.stage_transitions and pass the resolved topology through the existing runtime. Attach ForcedAligner after the actual route tail. Limit this change to static chains, without dynamic routing or a generic cleanup framework.

4. Whose customer or user needs does this address, and what are those needs?

Model integration developers need to adjust compatible execution chains without editing stage-number assumptions throughout the runtime, while keeping deployment, sampling-array and request-lifecycle contracts intact.

AI assistance: Codex helped with code, tests and writing.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	⚠️ Failed 2026-10-19T04:39:50.770777Z	539648e	Draft marked ready

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/observability.md.

Module owners: @lishunyang12 @vraiti @fake0fan

Routing: @lishunyang12 via module of the changed files, CODEOWNERS; @vraiti via module of the changed files; @fake0fan via CODEOWNERS

@Pinoeer-kingxi, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 539648e79495 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.