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

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

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]