[Bugfix] Fix Qwen3-TTS word timestamps with async chunking (#7544)

Signed-off-by: bezdarnost <amanurumbekov@gmail.com>

tests/config/test_forced_aligner_injection.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import pytest
 
 import vllm_omni.config.pipeline_registry  # noqa: F401  (populate registry)
@@ -69,3 +72,25 @@ def test_inject_noop_without_forced_aligner():
 
     assert ext_pipeline is pipeline
     assert len(ext_deploy.stages) == 0
+
+
+@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):
+    from vllm_omni.config.config_factory import StageConfigFactory
+
+    pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
+    deploy = DeployConfig(async_chunk=async_chunk)
+    resolution = StageConfigFactory._resolve_legacy_from_registry(
+        pipeline,
+        {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B", "async_chunk": cli_async_chunk},
+        user_deploy_config=deploy,
+    )
+    stages = [stage.to_omegaconf() for stage in resolution.stage_configs]
+    expected_async = async_chunk if cli_async_chunk is None else cli_async_chunk
+    assert len(stages) == 3
+    assert [stage.engine_args.async_chunk for stage in stages] == [expected_async, expected_async, False]
+    assert stages[2].custom_process_input_func.endswith("code2wav2aligner")
+    assert stages[2].engine_args.runner == "pooling"
+    assert len(pipeline.stages) == 2
+    assert deploy.stages == []

tests/config/test_omni_config.py

@@ -1705,3 +1705,37 @@ def test_compact_offload_config_is_validated_during_projection():
                 "components": ["dit"],
             }
         )
+
+
+@pytest.mark.parametrize("pipeline_async", [True, False])
+@pytest.mark.parametrize("stage_async", [None, False, True])
+def test_stage_async_chunk_opt_out_matches_legacy_config(pipeline_async, stage_async):
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    deploy = DeployConfig(
+        async_chunk=pipeline_async,
+        stages=[StageDeployConfig(stage_id=i, async_chunk=stage_async) for i in (0, 1)],
+    )
+    config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
+    legacy = merge_pipeline_deploy(pipeline, deploy)
+    expected = pipeline_async and stage_async is not False
+    assert config.stage_by_id(0).connector_config.async_chunk is expected
+    assert config.stage_by_id(1).connector_config.async_chunk is expected
+    assert legacy[1].yaml_engine_args["async_chunk"] is expected
+    assert config.stage_by_id(1).custom_process_input_func == legacy[1].custom_process_input_func
+    if not expected:
+        assert config.stage_by_id(1).custom_process_input_func.endswith("talker2code2wav_token_only")
+
+
+@pytest.mark.parametrize("disabled_stage", [0, 1])
+@pytest.mark.parametrize("builder", [merge_pipeline_deploy, VllmOmniConfig.from_pipeline_config])
+def test_async_chunk_rejects_mismatched_connector_edge(disabled_stage, builder):
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    deploy = DeployConfig(
+        async_chunk=True,
+        stages=[StageDeployConfig(stage_id=disabled_stage, async_chunk=False)],
+    )
+    with pytest.raises(ValueError, match="incompatible async_chunk settings on connector edge 0 -> 1"):
+        if builder is merge_pipeline_deploy:
+            builder(pipeline, deploy)
+        else:
+            builder(pipeline, user_deploy_config=deploy)

tests/e2e/online_serving/test_qwen3_tts_word_timestamps.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import base64
+import json
+
+import pytest
+from websockets.sync.client import connect
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
+
+_DEPLOY = modify_stage_config(
+    get_deploy_config_path("qwen3_tts.yaml"),
+    updates={
+        "stages": {
+            0: {
+                "max_num_seqs": 4,
+                "max_num_batched_tokens": 2048,
+                "max_model_len": 2048,
+                "gpu_memory_utilization": 0.45,
+                "compilation_config.cudagraph_capture_sizes": [1, 2, 4],
+                "default_sampling_params.max_tokens": 256,
+            },
+            1: {
+                "max_num_seqs": 4,
+                "max_num_batched_tokens": 4096,
+                "max_model_len": 8192,
+                "gpu_memory_utilization": 0.15,
+            },
+        }
+    },
+)
+_PARAMS = [
+    OmniServerParams(
+        model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+        stage_config_path=_DEPLOY,
+        server_args=[
+            "--trust-remote-code",
+            "--forced-aligner",
+            "Qwen/Qwen3-ForcedAligner-0.6B",
+            "--kv-cache-memory-bytes",
+            str(512 * 1024**2),
+            "--max-num-seqs",
+            "4",
+            "--max-num-batched-tokens",
+            "2048",
+        ],
+    )
+]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_async_chunk_word_timestamps(omni_server, online_client):
+    online_client.send_audio_speech_request(
+        {
+            "model": omni_server.model,
+            "input": "The weather is nice today, perfect for a walk in the park.",
+            "voice": "vivian",
+            "language": "English",
+            "word_timestamps": True,
+            "response_format": "wav",
+            "min_audio_bytes": 40_000,
+            "timeout": 120,
+        },
+        request_num=2,
+    )
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_async_chunk_speech_without_timestamps(omni_server, online_client):
+    online_client.send_audio_speech_request(
+        {
+            "model": omni_server.model,
+            "input": "The weather is nice today, perfect for a walk in the park.",
+            "voice": "vivian",
+            "language": "English",
+            "stream": True,
+            "stream_format": "audio",
+            "response_format": "wav",
+            "min_audio_bytes": 40_000,
+            "timeout": 120,
+        }
+    )
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
+def test_async_chunk_streaming_word_timestamps(omni_server):
+    text = "The weather is nice today, perfect for a walk in the park."
+    chunks = []
+    timestamps = None
+    with connect(f"ws://{omni_server.host}:{omni_server.port}/v1/audio/speech/stream") as ws:
+        ws.send(
+            json.dumps(
+                {
+                    "type": "session.config",
+                    "voice": "vivian",
+                    "language": "English",
+                    "stream_audio": True,
+                    "response_format": "pcm",
+                    "word_timestamps": True,
+                }
+            )
+        )
+        ws.send(json.dumps({"type": "input.text", "text": text}))
+        ws.send(json.dumps({"type": "input.done"}))
+        while True:
+            message = json.loads(ws.recv(timeout=120))
+            assert message["type"] != "error", message
+            if message["type"] == "audio.chunk":
+                if message["audio_b64"]:
+                    chunks.append(base64.b64decode(message["audio_b64"]))
+                if message["timestamps"] is not None:
+                    timestamps = message["timestamps"]
+                    sample_rate = message["sample_rate"]
+            elif message["type"] == "audio.done":
+                assert not message["error"], message
+                assert message["total_bytes"] == sum(map(len, chunks))
+            elif message["type"] == "session.done":
+                break
+
+    assert len(chunks) > 1
+    assert timestamps
+    assert len(timestamps) == len(text.split())
+    duration_ms = sum(map(len, chunks)) / (2 * sample_rate) * 1000
+    final_chunk_ms = len(chunks[-1]) / (2 * sample_rate) * 1000
+    assert timestamps[-1]["end_ms"] > max(duration_ms / 2, final_chunk_ms)
+    previous_end = 0
+    for timestamp in timestamps:
+        assert previous_end <= timestamp["start_ms"] <= timestamp["end_ms"] <= duration_ms + 1
+        previous_end = timestamp["end_ms"]

tests/engine/test_async_omni_engine_stage_init.py

@@ -10,6 +10,7 @@ import time
 import types
 
 import pytest
+from omegaconf import OmegaConf
 
 from vllm_omni.diffusion.data import AttentionConfig
 from vllm_omni.engine import omni_engine_base as async_omni_engine_module
@@ -1829,3 +1830,25 @@ def test_port_from_zmq_address_parsing():
     assert _port_from_zmq_address(None) is None
     assert _port_from_zmq_address("ipc:///tmp/sock") is None
     assert _port_from_zmq_address("tcp://host:not-a-port") is None
+
+
+@pytest.mark.parametrize(
+    "stage_modes,expected", [([False, True, True], True), ([False, False], False), ([True, False], True)]
+)
+def test_engine_async_chunk_includes_downstream_stages(monkeypatch, stage_modes, expected):
+    engine = object.__new__(AsyncOmniEngine)
+    stages = [OmegaConf.create({"engine_args": {"async_chunk": mode}}) for mode in stage_modes]
+    monkeypatch.setattr(async_omni_engine_module.StageConfigFactory, "get_pipeline_config", lambda *a, **k: None)
+    monkeypatch.setattr(engine, "_resolve_stage_configs", lambda *a, **k: (None, stages))
+    monkeypatch.setattr(engine, "_set_pipeline_runtime_config", lambda *a: None)
+
+    class ConfigResolvedError(Exception):
+        pass
+
+    def stop_before_queues(*args, **kwargs):
+        raise ConfigResolvedError
+
+    monkeypatch.setattr(async_omni_engine_module.janus, "Queue", stop_before_queues)
+    with pytest.raises(ConfigResolvedError):
+        engine.__init__("test-model")
+    assert engine.async_chunk is expected

tests/engine/test_orchestrator_kv_sender_info.py

@@ -414,3 +414,67 @@ def test_prewarm_submits_bound_payload_endpoint_for_concurrent_replicas():
 
     assert all(isinstance(request, OmniEngineCoreRequest) for request in submitted.values())
     assert {key: request.payload_sender_info for key, request in submitted.items()} == endpoints
+
+
+@pytest.mark.asyncio
+async def test_async_pipeline_defers_sync_stage_until_audio_finishes(mocker):
+    from vllm import PoolingParams, RequestOutput
+    from vllm.config import VllmConfig
+
+    from vllm_omni.config.model import OmniModelConfig
+    from vllm_omni.engine import OmniEngineCoreRequest
+
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.async_chunk = True
+    orchestrator._pd_pair = None
+    orchestrator._cfg_tracker = CfgCompanionTracker()
+    orchestrator.output_async_queue = asyncio.Queue()
+    orchestrator._duplex_output_decision = lambda *args: None
+    orchestrator._is_duplex_session_request = lambda state: False
+    orchestrator._forward_to_next_stage = mocker.AsyncMock()
+    orchestrator._cleanup_request_ids = mocker.AsyncMock()
+    # A sender stage stands in for Code2Wav, whose audio is also a final output.
+    audio_pool = _build_sender_pool(0, {})
+    audio_pool.stage_client.final_output = True
+    aligner = _DummyDiffusionStage()
+    aligner_pool = StagePool(
+        1,
+        aligner,
+        stage_vllm_config=mocker.Mock(
+            spec=VllmConfig, model_config=mocker.Mock(spec=OmniModelConfig, async_chunk=False)
+        ),
+    )
+    orchestrator.stage_pools = [audio_pool, aligner_pool]
+    state = OrchestratorRequestState(
+        request_id="aligned",
+        prompt={"additional_information": {"text": ["Hello world"]}},
+        sampling_params_list=[SamplingParams(), PoolingParams(task="token_classify")],
+        final_stage_id=1,
+        final_output_stage_ids={0, 1},
+    )
+
+    await orchestrator._prewarm_async_chunk_stages(
+        "aligned", mocker.Mock(spec=OmniEngineCoreRequest, prompt_token_ids=[1, 2]), state
+    )
+    assert aligner.calls == []
+    assert 1 not in state.stage_submit_ts
+
+    output = RequestOutput("aligned", None, [1, 2], None, [], False)
+    await orchestrator._route_output(0, 0, output, state, None)
+    orchestrator._forward_to_next_stage.assert_not_awaited()
+    output.finished = True
+    await orchestrator._route_output(0, 0, output, state, None)
+    orchestrator._forward_to_next_stage.assert_awaited_once_with(
+        "aligned",
+        0,
+        output,
+        state,
+        src_replica_id=0,
+        is_streaming_session=False,
+        is_final_update=False,
+    )
+    orchestrator._cleanup_request_ids.assert_not_awaited()
+    await orchestrator._route_output(1, 0, output, state, None)
+    orchestrator._cleanup_request_ids.assert_awaited_once()
+    messages = [orchestrator.output_async_queue.get_nowait() for _ in range(3)]
+    assert [message.finished for message in messages] == [False, False, True]

tests/engine/test_output_processor.py

@@ -888,3 +888,24 @@ def test_abort_snapshot_leaves_state_until_commit():
     assert "0_parent" not in processor.request_states
     assert "parent" not in processor.parent_requests
     assert "parent" not in processor.external_req_ids
+
+
+def test_cumulative_audio_terminal_output_supplies_full_aligner_waveform():
+    from vllm_omni.model_executor.stage_input_processors.forced_aligner import code2wav2aligner
+
+    state = _make_state(RequestOutputKind.CUMULATIVE)
+    for index, count in enumerate((800, 1200, 400)):
+        state.add_multimodal_tensor(torch.full((1, count), float(index)), mm_type=AUDIO)
+        output = state.make_request_output([index], None, FinishReason.STOP if index == 2 else None, None)
+        assert output is not None
+        if index < 2:
+            assert code2wav2aligner([output], {}) == []
+    prompt = {"additional_information": {"text": ["Hello world"], "language": ["English"]}}
+    aligned = code2wav2aligner([output], prompt)[0]
+    waveform, sr = aligned["multi_modal_data"]["audio"]
+    assert sr == 24000
+    assert waveform.shape == (2400,)
+    assert (waveform[:800] == 0).all()
+    assert (waveform[800:2000] == 1).all()
+    assert (waveform[2000:] == 2).all()
+    assert aligned["additional_information"]["aligner_audio_duration_ms"] == [100.0]

tests/entrypoints/openai_api/test_serving_speech.py

@@ -5212,6 +5212,52 @@ class TestTTSAsyncOffloading:
             for call in log_info.call_args_list
         )
 
+    @pytest.mark.parametrize("word_timestamps", [False, True])
+    def test_streaming_timestamps_retain_complete_audio(self, qwen3_tts_server, mocker, word_timestamps):
+        from vllm import SamplingParams
+        from vllm.sampling_params import RequestOutputKind
+
+        defaults = [SamplingParams(), SamplingParams()]
+        qwen3_tts_server.engine_client.default_sampling_params_list = defaults
+        qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None)
+        qwen3_tts_server._adapter._build_tts_params = mocker.MagicMock(
+            return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]}
+        )
+        qwen3_tts_server._adapter._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
+        request = OpenAICreateSpeechRequest(
+            input="hello", stream=True, response_format="pcm", word_timestamps=word_timestamps
+        )
+        asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
+
+        params = qwen3_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"]
+        expected = RequestOutputKind.CUMULATIVE if word_timestamps else RequestOutputKind.DELTA
+        assert all(sp.output_kind == expected for sp in params)
+        assert all(sp.output_kind == RequestOutputKind.CUMULATIVE for sp in defaults)
+        assert all(sp is not original for sp, original in zip(params, defaults))
+
+    @pytest.mark.asyncio
+    async def test_cumulative_timestamp_audio_emits_each_sample_once(self, qwen3_tts_server):
+        audio = torch.linspace(-0.5, 0.5, 2400)
+
+        async def cumulative_generator():
+            for end in (800, 1600, 2400, 2400):
+                yield OmniRequestOutput(
+                    request_id="req-cumulative",
+                    final_output_type="audio",
+                    _multimodal_output={"audio": audio[:end], "sr": 24000},
+                )
+
+        chunks = [
+            chunk
+            async for chunk in qwen3_tts_server._generate_pcm_chunks(
+                cumulative_generator(), "req-cumulative", cumulative_audio=True
+            )
+        ]
+        assert len(chunks) == 3
+        assert len(b"".join(chunks)) == 2400 * 2
+        decoded = np.frombuffer(b"".join(chunks), dtype=np.int16)
+        assert np.all(np.diff(decoded.astype(np.int32)) >= 0)
+
     def test_prepare_speech_generation_treats_sse_as_streaming(self, qwen3_tts_server, mocker: MockerFixture):
         """stream_format=sse should request delta-style multimodal outputs."""
         qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None)

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -56,6 +56,7 @@ def _build_test_app(
             include_sample_rate=False,
             tts_params=None,
             collect=None,
+            cumulative_audio=False,
         ):
             for chunk in (b"\x01\x02", b"\x03\x04\x05"):
                 yield (chunk, 24000) if include_sample_rate else chunk
@@ -386,9 +387,11 @@ class TestStreamingSpeechWebSocket:
             include_sample_rate=False,
             tts_params=None,
             collect=None,
+            cumulative_audio=False,
         ):
             assert request_start_s is not None
             assert request_arrival_ts is not None
+            assert cumulative_audio is True
             for chunk in (first_chunk, second_chunk):
                 yield (chunk, 1000) if include_sample_rate else chunk
             if collect is not None:
@@ -480,6 +483,7 @@ class TestStreamingSpeechWebSocket:
             include_sample_rate=False,
             tts_params=None,
             collect=None,
+            cumulative_audio=False,
         ):
             chunk = b"\x01" * 1000
             yield (chunk, 1000) if include_sample_rate else chunk

tests/helpers/assertions.py

@@ -955,6 +955,17 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
     """
     assert response.success, "The request failed."
 
+    if request_config.get("word_timestamps"):
+        timestamps = response.word_timestamps
+        assert isinstance(timestamps, list) and timestamps, "Expected nonempty X-Word-Timestamps"
+        previous_start = 0
+        for timestamp in timestamps:
+            assert isinstance(timestamp["word"], str) and timestamp["word"]
+            start, end = timestamp["start_ms"], timestamp["end_ms"]
+            assert isinstance(start, int) and isinstance(end, int)
+            assert previous_start <= start <= end
+            previous_start = start
+
     # Optional floor on decoded audio size (models with very short clips may use a lower value).
     min_audio = request_config.get("min_audio_bytes")
     if min_audio is not None:

tests/helpers/client.py

@@ -102,6 +102,7 @@ class OmniResponse:
     audio_content: str | None = None
     audio_format: str | None = None
     audio_bytes: bytes | None = None
+    word_timestamps: list[dict[str, str | int]] | None = None
     #: End-to-end wall time in **seconds** (``perf_counter`` delta), from just before the
     #: OpenAI client call through response parsing and local post-process (e.g. audio decode).
     e2e_latency: float | None = None
@@ -1419,6 +1420,9 @@ class OnlineOmniClient:
             result.success = True
             result.audio_format = getattr(response, "response", None)
             if result.audio_format is not None:
+                timestamps = result.audio_format.headers.get("X-Word-Timestamps")
+                if timestamps is not None:
+                    result.word_timestamps = json.loads(timestamps)
                 result.audio_format = result.audio_format.headers.get("content-type", "")
 
         except Exception as e:
@@ -1472,6 +1476,7 @@ class OnlineOmniClient:
             "extra_params",
             "stream_format",
             "x_vector_only_mode",
+            "word_timestamps",
         ):
             if key in request_config:
                 extra_body[key] = request_config[key]

tests/helpers/tests/test_assertions.py

@@ -173,3 +173,26 @@ def test_bounded_tail_allows_ratio_bound_for_long_answers():
     # Twenty expected words allow up to max(2, ceil(0.2 * 20)) = 4 tail words.
     assert assertions._transcript_has_bounded_tail(four_word_tail, expected)
     assert not assertions._transcript_has_bounded_tail(five_word_tail, expected)
+
+
+def test_speech_timestamp_header_is_required_and_decoded():
+    import time
+
+    import httpx
+    from openai._legacy_response import HttpxBinaryResponseContent
+
+    from tests.helpers.client import OnlineOmniClient
+
+    client = object.__new__(OnlineOmniClient)
+    request = {"word_timestamps": True}
+    for headers in ({}, {"X-Word-Timestamps": '[{"word":"Hello","start_ms":0,"end_ms":300}]'}):
+        response = client._process_non_stream_audio_speech_response(
+            HttpxBinaryResponseContent(httpx.Response(200, content=b"audio", headers=headers)),
+            wall_start=time.perf_counter(),
+        )
+        if headers:
+            assert_audio_speech_response(response, request)
+            assert response.word_timestamps == [{"word": "Hello", "start_ms": 0, "end_ms": 300}]
+        else:
+            with pytest.raises(AssertionError, match="X-Word-Timestamps"):
+                assert_audio_speech_response(response, request)

vllm_omni/config/omni_config.py

@@ -50,6 +50,8 @@ from vllm_omni.config.stage_config import (
     merge_sampling_constraints,
     normalize_pipeline_cli_overrides,
     reconcile_diffusion_attention_overrides,
+    resolve_stage_async_chunk,
+    validate_stage_async_chunk_edges,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 
@@ -1520,7 +1522,7 @@ def _build_common_stage_config_kwargs(
     *,
     model: str | None,
 ) -> tuple[dict[str, Any], str | None, str | None]:
-    input_proc, next_stage_proc = _select_processor_funcs(topology, bool(deploy.async_chunk))
+    input_proc, next_stage_proc = _select_processor_funcs(topology, resolve_stage_async_chunk(deploy, stage_deploy))
     quantization_config = _build_quantization_config(deploy, engine.quantization)
     parallel_config = _build_parallel_config(deploy, engine.parallel, parallel_config_cls)
 
@@ -1835,7 +1837,7 @@ def _build_connector_config(
     output_connectors = stage_deploy.output_connectors if stage_deploy is not None else None
     input_connectors = stage_deploy.input_connectors if stage_deploy is not None else None
     return cast(Any, OmniStageConnectorConfig)(
-        async_chunk=bool(deploy.async_chunk),
+        async_chunk=resolve_stage_async_chunk(deploy, stage_deploy),
         omni_kv_config=_copy_value(engine.get("omni_kv_config")),
         output_connectors=_copy_value(output_connectors) if output_connectors else None,
         input_connectors=_copy_value(input_connectors) if input_connectors else None,
@@ -1966,6 +1968,7 @@ class VllmOmniConfig:
         if len(pipeline_cfg.stages) <= 1:
             deploy.async_chunk = False
         _validate_async_chunk_support(pipeline_cfg, deploy)
+        validate_stage_async_chunk_edges(pipeline_cfg, deploy)
         deploy_by_id = {stage.stage_id: stage for stage in deploy.stages}
         model = cli_overrides.get("model")
 

vllm_omni/config/stage_config.py

@@ -390,6 +390,9 @@ class StageDeployConfig:
     num_replicas: int = 1
     env: dict[str, Any] | None = None
 
+    # False opts this stage out of pipeline-wide async chunking.
+    async_chunk: bool | None = None
+
     # Inter-stage connector wiring and request defaults.
     output_connectors: dict[str, str] | None = None
     input_connectors: dict[str, str] | None = None
@@ -587,6 +590,7 @@ class DeployConfig:
 
 _STAGE_RESERVED_KEYS = frozenset(
     {
+        "async_chunk",
         "stage_id",
         "devices",
         "num_replicas",
@@ -645,6 +649,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
         if name in flat_args:
             kwargs[name] = flat_args.pop(name)
 
+    kwargs["async_chunk"] = stage_data.get("async_chunk")
     kwargs["output_connectors"] = stage_data.get("output_connectors")
     kwargs["input_connectors"] = stage_data.get("input_connectors")
     kwargs["default_sampling_params"] = stage_data.get("default_sampling_params")
@@ -888,6 +893,28 @@ def _resolve_execution_mode(
     return _EXECUTION_TYPE_TO_STAGE_WORKER.get(execution_type, (StageType.LLM, None))
 
 
+def resolve_stage_async_chunk(deploy: DeployConfig, stage: StageDeployConfig | None) -> bool:
+    return bool(deploy.async_chunk and (stage is None or stage.async_chunk is not False))
+
+
+def validate_stage_async_chunk_edges(pipeline: PipelineConfig, deploy: DeployConfig) -> None:
+    stages = {stage.stage_id: stage for stage in pipeline.stages}
+    deploy_by_id = {stage.stage_id: stage for stage in deploy.stages}
+    for consumer in pipeline.stages:
+        for source_id in consumer.input_sources:
+            producer = stages[source_id]
+            if not producer.async_chunk_process_next_stage_input_func:
+                continue
+            producer_async = resolve_stage_async_chunk(deploy, deploy_by_id.get(source_id))
+            consumer_async = resolve_stage_async_chunk(deploy, deploy_by_id.get(consumer.stage_id))
+            if producer_async != consumer_async:
+                raise ValueError(
+                    f"Pipeline {pipeline.model_type!r} has incompatible async_chunk settings on "
+                    f"connector edge {source_id} -> {consumer.stage_id}. "
+                    "Set the same async_chunk mode on both stages or disable pipeline-wide async_chunk."
+                )
+
+
 def _select_processor_funcs(
     ps: StagePipelineConfig,
     async_chunk: bool,
@@ -964,9 +991,7 @@ def _build_engine_args(
                 continue
             engine_args[k] = v
         engine_args.update(ds.engine_extras)
-    # Materialize the resolved pipeline-wide async_chunk value into every
-    # stage so explicit False overrides do not get lost downstream.
-    engine_args["async_chunk"] = bool(deploy.async_chunk)
+    engine_args["async_chunk"] = resolve_stage_async_chunk(deploy, ds)
     engine_args["session_mode"] = deploy.session_mode
     if deploy.session_mode == "duplex":
         # The engine admission limit is also the authoritative capacity for
@@ -1060,11 +1085,12 @@ def merge_pipeline_deploy(
             "Either set async_chunk=False or implement an async-chunk producer on the pipeline."
         )
 
+    validate_stage_async_chunk_edges(pipeline, deploy)
     result: list[StageConfig] = []
     for ps in pipeline.stages:
         ds = deploy_by_id.get(ps.stage_id)
         stage_type, worker_type = _resolve_execution_mode(ps.execution_type)
-        input_proc, next_stage_proc = _select_processor_funcs(ps, deploy.async_chunk)
+        input_proc, next_stage_proc = _select_processor_funcs(ps, resolve_stage_async_chunk(deploy, ds))
         engine_args = _build_engine_args(ps, ds, pipeline, deploy, next_stage_proc)
         # Downstream stages may share a multimodal wrapper class without owning
         # an encoder. Do not make vLLM profile dummy multimodal inputs for them.

vllm_omni/distributed/omni_connectors/utils/config.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from dataclasses import dataclass, field
 from typing import Any
@@ -33,7 +33,7 @@ def get_stage_connector_role(model_config: Any) -> str | None:
 
 def stage_receives_chunks(model_config: Any) -> bool:
     """Whether connector chunks, rather than the orchestrator, feed a stage."""
-    return get_stage_connector_role(model_config) != "sender"
+    return bool(getattr(model_config, "async_chunk", True)) and get_stage_connector_role(model_config) != "sender"
 
 
 def stage_sends_async_output(model_config: Any) -> bool:

vllm_omni/engine/omni_engine_base.py

@@ -271,8 +271,9 @@ class OmniEngineBase:
             )
 
         self.num_stages = len(self.stage_configs)
-        stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None
-        self.async_chunk = bool(getattr(stage0_args, "async_chunk", False))
+        self.async_chunk = any(
+            bool(getattr(getattr(stage, "engine_args", None), "async_chunk", False)) for stage in self.stage_configs
+        )
         self.stage_pools: list[StagePool] = []
         self.stage_clients: list[StageClient] = []  # logical-stage view for external readers
         self.input_processor: InputProcessor | None = None

vllm_omni/entrypoints/openai/serving_speech.py

@@ -32,6 +32,7 @@ from vllm.entrypoints.launchers.launcher import terminate_if_errored
 from vllm.entrypoints.serve.engine.protocol import ErrorResponse
 from vllm.logger import init_logger
 from vllm.multimodal.media import MediaConnector
+from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.utils import random_uuid
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
@@ -1444,14 +1445,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         tts_params: dict[str, Any] | None = None,
         collect: dict | None = None,
         target_sample_rate: int | None = None,
+        cumulative_audio: bool = False,
     ):
         """Generate audio chunks for streaming response.
 
         Handles two audio output modes from the engine:
         - Cumulative mode (list): Engine returns growing list of chunks;
         we emit only the new tail on each iteration.
-        - Per-step mode (tensor): Engine returns single tensor per iteration;
-        we emit it directly.
+        - Tensor mode: Emit each tensor directly, or only its new samples when
+        cumulative_audio is set for requests that retain audio for alignment.
 
         Args:
             generator: Async generator from the engine
@@ -1462,6 +1464,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             Raw audio bytes for each chunk (with WAV header for first chunk if wav format)
         """
         prev_count = 0
+        emitted_samples = 0
         sample_rate_val = 24000
         first_chunk = True
         first_audio_chunk_s: float | None = None
@@ -1563,6 +1566,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     )
                     if chunk_np.ndim > 1:
                         chunk_np = chunk_np.squeeze()
+                    if cumulative_audio and not isinstance(audio_val, list):
+                        total_samples = chunk_np.shape[-1]
+                        chunk_np = chunk_np[..., emitted_samples:]
+                        emitted_samples = total_samples
+                        if chunk_np.size == 0:
+                            continue
                     if resampler is not None:
                         chunk_np = resampler.process(chunk_np)
                         if chunk_np.size == 0:
@@ -1957,6 +1966,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             from vllm_omni.model_executor.stage_input_processors.forced_aligner import TIMESTAMPS_MODALITY
 
             output_modalities.append(TIMESTAMPS_MODALITY)
+            if is_streaming_request:
+                # The aligner consumes the terminal waveform, so retain earlier
+                # audio chunks instead of draining them after each DELTA output.
+                sampling_params_list = [
+                    sp.clone() if isinstance(sp, SamplingParams) else sp for sp in sampling_params_list
+                ]
+                for sp in sampling_params_list:
+                    if isinstance(sp, SamplingParams):
+                        sp.output_kind = RequestOutputKind.CUMULATIVE
 
         generator = self.engine_client.generate(
             prompt=prompt,
@@ -1990,6 +2008,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         tts_params: dict[str, Any] | None = None,
         collect: dict | None = None,
         target_sample_rate: int | None = None,
+        cumulative_audio: bool = False,
     ):
         """Yield raw PCM byte chunks from the engine generator.
 
@@ -2009,6 +2028,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 tts_params=tts_params,
                 collect=collect,
                 target_sample_rate=target_sample_rate,
+                cumulative_audio=cumulative_audio,
             )
         ) as chunks:
             async for chunk in chunks:
@@ -2024,6 +2044,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     request_id,
                     tts_params=tts_params,
                     target_sample_rate=request.sample_rate,
+                    cumulative_audio=request.word_timestamps,
                 )
             ) as chunks:
                 async for chunk in chunks:

vllm_omni/entrypoints/openai/serving_speech_stream.py

@@ -483,6 +483,7 @@ class OmniStreamingSpeechHandler:
                 include_sample_rate=True,
                 tts_params=tts_params,
                 collect=collect,
+                cumulative_audio=True,
             )
         ) as stream:
             async for chunk, chunk_sample_rate in stream:

vllm_omni/utils/forced_aligner.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Forced-aligner config + word-timestamp decoding for TTS.
 
 The aligner runs as a pooling stage appended to the pipeline by
@@ -153,6 +153,7 @@ def inject_forced_aligner_stage(
         # Pooling is one-shot -> sync AR scheduler (the async one assumes
         # token-by-token generation).
         async_scheduling=False,
+        async_chunk=False,
     )
     deploy = replace(deploy, stages=list(deploy.stages) + [aligner_ds])
     logger.info("[stage_init] Injected forced-aligner stage %d (model=%s)", new_id, fa.model)