[Bugfix] Fix Qwen3-TTS word timestamps with async chunking

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 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 @@ async def run():
 
     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 @@ def apply_sampling_overrides(self, sampling_params_list, request, prompt, reques
             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)