[Feature] timestamps for TTS via a forced-aligner pooling (#4795)

Signed-off-by: wjinxu <1299461899@qq.com>
Co-authored-by: xurui <xurui1@modelbest.onaliyun.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

tests/config/test_forced_aligner_injection.py

@@ -0,0 +1,71 @@
+import pytest
+
+import vllm_omni.config.pipeline_registry  # noqa: F401  (populate registry)
+from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY
+from vllm_omni.config.stage_config import (
+    DeployConfig,
+    StageExecutionType,
+)
+from vllm_omni.model_executor.stage_input_processors.forced_aligner import (
+    POOLING_OUTPUT_DECODER_PATH,
+    TIMESTAMPS_MODALITY,
+)
+from vllm_omni.utils.forced_aligner import _DEFAULT_CONFIG_PATH, inject_forced_aligner_stage
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_inject_appends_pooling_aligner_stage():
+    pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
+    deploy = DeployConfig()
+    n0 = len(pipeline.stages)
+
+    ext_pipeline, ext_deploy = inject_forced_aligner_stage(
+        pipeline, deploy, {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"}
+    )
+
+    # A new pooling stage is appended; the original pipeline is unchanged.
+    assert len(pipeline.stages) == n0
+    assert len(ext_pipeline.stages) == n0 + 1
+    # The original deploy config is not mutated either (double-injection safe).
+    assert len(deploy.stages) == 0
+    assert len(ext_deploy.stages) == 1
+    aligner = ext_pipeline.stages[-1]
+    # Pooling stage reuses LLM_AR (+ runner=pooling), no dedicated LLM_POOLING.
+    assert aligner.execution_type == StageExecutionType.LLM_AR
+    assert aligner.final_output is True
+    assert aligner.final_output_type == TIMESTAMPS_MODALITY
+    assert aligner.input_sources == (n0 - 1,)
+    assert aligner.requires_multimodal_data is True
+    assert aligner.owns_tokenizer is True
+
+    # The deploy stage carries the separate aligner checkpoint + decoder hook.
+    ds = ext_deploy.stages[-1]
+    assert ds.engine_extras["model"] == "/models/Qwen3-ForcedAligner-0.6B"
+    assert ds.engine_extras["runner"] == "pooling"
+    assert ds.engine_extras["pooling_output_decoder"] == POOLING_OUTPUT_DECODER_PATH
+
+
+def test_inject_from_config_only():
+    # --forced-aligner-config alone (model in the YAML) must also inject.
+    pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
+    deploy = DeployConfig()
+    n0 = len(pipeline.stages)
+
+    ext_pipeline, ext_deploy = inject_forced_aligner_stage(
+        pipeline, deploy, {"forced_aligner": None, "forced_aligner_config": str(_DEFAULT_CONFIG_PATH)}
+    )
+
+    assert len(ext_pipeline.stages) == n0 + 1
+    assert ext_pipeline.stages[-1].final_output_type == TIMESTAMPS_MODALITY
+    assert ext_deploy.stages[-1].engine_extras["model"] == "Qwen/Qwen3-ForcedAligner-0.6B"
+
+
+def test_inject_noop_without_forced_aligner():
+    pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
+    deploy = DeployConfig()
+
+    ext_pipeline, ext_deploy = inject_forced_aligner_stage(pipeline, deploy, {})
+
+    assert ext_pipeline is pipeline
+    assert len(ext_deploy.stages) == 0

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -27,6 +27,7 @@ _MIXIN_UPDATE_HELPERS = (
     "_attach_finished_request_sets",
     "_attach_scheduler_stats",
     "_capture_omni_connector_output",
+    "_maybe_decode_pooling_output",
 )
 
 
@@ -138,6 +139,7 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
         pending_stop_after_extraction=set(),
         waiting_for_transfer_free=set(),
         _new_prompt_len_snapshot={},
+        _pooling_output_decoder=None,
         finished_req_ids=set(),
         finished_req_ids_dict=defaultdict(set),
         kv_cache_manager=SimpleNamespace(take_events=lambda: None),
@@ -267,3 +269,72 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
     assert output_by_id["bad"].new_token_ids == []
     assert output_by_id["good"].new_token_ids == [8]
     np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8])
+
+
+def _pooling_model_runner_output(pooler_tensor) -> SimpleNamespace:
+    return SimpleNamespace(
+        sampled_token_ids=[[]],
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=[pooler_tensor],
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={"req": 0},
+        routed_experts=None,
+    )
+
+
+def _make_pooling_stub(decoder) -> tuple[_Request, SimpleNamespace]:
+    request = _Request("req")
+    request.sampling_params = SimpleNamespace(num_logprobs=None)
+    request.pooling_params = SimpleNamespace()
+    scheduler = _make_scheduler_stub([request])
+    scheduler._pooling_output_decoder = decoder
+    scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(hf_config=None))
+    _bind_request_lifecycle(scheduler, update_request=lambda req, tokens: (tokens, False))
+    return request, scheduler
+
+
+def test_pooling_decode_success_emits_decoded_payload() -> None:
+    import torch
+
+    decoded = torch.tensor([[1, 2]], dtype=torch.int32)
+    request, scheduler = _make_pooling_stub(lambda logits, req, hf: decoded)
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"req": 1},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    outputs = OmniARScheduler.update_from_output(
+        scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4))
+    )
+    (output,) = outputs[0].outputs
+
+    assert output.finish_reason is FinishReason.STOP
+    assert output.pooling_output is decoded
+
+
+def test_pooling_decode_failure_finishes_request_with_error() -> None:
+    import torch
+
+    def raising_decoder(logits, req, hf):
+        raise ValueError("boom")
+
+    request, scheduler = _make_pooling_stub(raising_decoder)
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"req": 1},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    outputs = OmniARScheduler.update_from_output(
+        scheduler, scheduler_output, _pooling_model_runner_output(torch.zeros(2, 4))
+    )
+    (output,) = outputs[0].outputs
+
+    # Decoder failure is a request error (500), not an empty success.
+    assert output.finish_reason is FinishReason.ERROR
+    assert output.pooling_output is None
+    assert "pooling output decode failed" in str(request.stop_reason)

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -3,6 +3,7 @@ import base64
 from types import SimpleNamespace
 
 import pytest
+import torch
 from fastapi import FastAPI, WebSocket
 from pytest_mock import MockerFixture
 from starlette.testclient import TestClient
@@ -11,11 +12,24 @@ from starlette.websockets import WebSocketDisconnect
 from vllm_omni.entrypoints.openai import serving_speech_stream as streaming_speech_module
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
-from vllm_omni.utils.forced_aligner import WordTimestamp
+from vllm_omni.model_executor.stage_input_processors.forced_aligner import ALIGNER_WORDS_KEY
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def _fake_aligner_res(pairs, words):
+    """Build a stand-in for the forced-aligner stage's pooling output.
+
+    Mirrors what rides the generator in production: ``res.outputs.data`` is a
+    ``[n_words, 2]`` int32 tensor of ``[start_ms, end_ms]`` and the word strings
+    travel in ``additional_information``. Decoded by ``extract_word_timestamps``.
+    """
+    return SimpleNamespace(
+        outputs=SimpleNamespace(data=torch.tensor(pairs, dtype=torch.int32), additional_information=None),
+        additional_information={ALIGNER_WORDS_KEY: list(words)},
+    )
+
+
 def _build_test_app(
     speech_service=None,
     *,
@@ -28,9 +42,9 @@ def _build_test_app(
         speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
         speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav"))
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
-        speech_service.forced_aligner_config = None
+        speech_service.forced_aligner_enabled = False
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
             for chunk in (b"\x01\x02", b"\x03\x04\x05"):
                 yield (chunk, 24000) if include_sample_rate else chunk
 
@@ -235,7 +249,7 @@ class TestStreamingSpeechWebSocket:
         speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
-        speech_service.forced_aligner_config = None
+        speech_service.forced_aligner_enabled = False
 
         async def mock_prepare_speech_generation(request):
             captured_requests.append(request)
@@ -311,13 +325,13 @@ class TestStreamingSpeechWebSocket:
                 assert error["type"] == "error"
                 assert "without --forced-aligner" in error["message"]
 
-    def test_word_timestamps_emit_sidecar_json_frame(self, mocker: MockerFixture):
+    def test_word_timestamps_emit_pipeline_json_frame(self, mocker: MockerFixture):
         captured_requests = []
         speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
         speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
-        speech_service.forced_aligner_config = SimpleNamespace(model="aligner")
+        speech_service.forced_aligner_enabled = True
 
         async def mock_prepare_speech_generation(request):
             captured_requests.append(request)
@@ -328,19 +342,15 @@ class TestStreamingSpeechWebSocket:
         first_chunk = b"\x01" * 1000
         second_chunk = b"\x02" * 1000
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        # The forced-aligner stage rides the same generator: its pooling output
+        # is surfaced via the ``collect`` channel once the audio has streamed.
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
             for chunk in (first_chunk, second_chunk):
                 yield (chunk, 1000) if include_sample_rate else chunk
+            if collect is not None:
+                collect["aligner_res"] = _fake_aligner_res([[0, 200], [200, 900]], ["Hello", "world"])
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
-        # Sentence-level: aligner runs once over the whole sentence audio.
-        mock_align = mocker.AsyncMock(
-            return_value=[
-                WordTimestamp("Hello", 0, 200),
-                WordTimestamp("world", 200, 900),
-            ]
-        )
-        mocker.patch.object(streaming_speech_module, "forced_align", mock_align)
         app, _ = _build_test_app(speech_service)
 
         with TestClient(app) as client:
@@ -406,7 +416,6 @@ class TestStreamingSpeechWebSocket:
                 }
 
         assert captured_requests[0].word_timestamps is True
-        assert mock_align.await_count == 1
 
     def test_word_timestamps_emit_word_dicts(self, mocker: MockerFixture):
         # The streaming layer forwards the aligner's (already monotonic,
@@ -415,21 +424,16 @@ class TestStreamingSpeechWebSocket:
         speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
-        speech_service.forced_aligner_config = SimpleNamespace(model="aligner")
+        speech_service.forced_aligner_enabled = True
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
 
-        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False, collect=None):
             chunk = b"\x01" * 1000
             yield (chunk, 1000) if include_sample_rate else chunk
+            if collect is not None:
+                collect["aligner_res"] = _fake_aligner_res([[0, 1000], [1000, 1200]], ["Hello", "world"])
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
-        mock_align = mocker.AsyncMock(
-            return_value=[
-                WordTimestamp("Hello", 0, 1000),
-                WordTimestamp("world", 1000, 1200),
-            ]
-        )
-        mocker.patch.object(streaming_speech_module, "forced_align", mock_align)
         app, _ = _build_test_app(speech_service)
 
         with TestClient(app) as client:
@@ -572,7 +576,7 @@ class TestStreamingSpeechWebSocket:
         speech_service._generate_pcm_chunks = mocker.AsyncMock()
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
-        speech_service.forced_aligner_config = None
+        speech_service.forced_aligner_enabled = False
         app, _ = _build_test_app(speech_service)
 
         with TestClient(app) as client:
@@ -602,7 +606,7 @@ class TestStreamingSpeechWebSocket:
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-stream-err", object(), {}))
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
-        speech_service.forced_aligner_config = None
+        speech_service.forced_aligner_enabled = False
 
         async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             yield b"\x01\x02"
@@ -696,7 +700,7 @@ class TestStreamingSpeechWebSocket:
         speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-abort", object(), {}))
         speech_service.engine_client = mocker.MagicMock()
         speech_service.engine_client.abort = mocker.AsyncMock()
-        speech_service.forced_aligner_config = None
+        speech_service.forced_aligner_enabled = False
 
         async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False):
             yield b"\x01\x02"

tests/model_executor/stage_input_processors/test_forced_aligner_stage.py

@@ -0,0 +1,117 @@
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.forced_aligner import (
+    ALIGNER_WORDS_KEY,
+    _extract_text,
+    _extract_waveform,
+    code2wav2aligner,
+    decode_pooling_output,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeCompletion:
+    def __init__(self, multimodal_output):
+        self.multimodal_output = multimodal_output
+
+
+class _FakeOutput:
+    def __init__(self, multimodal_output, finished=True):
+        self.finished = finished
+        self.outputs = [_FakeCompletion(multimodal_output)]
+
+
+def _audio_mm(n=2400, sr=24000):
+    return {"audio": torch.zeros(n, dtype=torch.float32), "sr": torch.tensor(sr, dtype=torch.int32)}
+
+
+def test_extract_waveform_from_audio_key():
+    wav, sr = _extract_waveform(_audio_mm(n=100, sr=16000))
+    assert isinstance(wav, np.ndarray) and wav.dtype == np.float32
+    assert wav.shape == (100,)
+    assert sr == 16000
+
+
+def test_extract_waveform_falls_back_to_model_outputs_and_default_sr():
+    wav, sr = _extract_waveform({"model_outputs": [torch.ones(50)]})
+    assert wav.shape == (50,)
+    assert sr == 24000  # default when sr missing
+
+
+def test_extract_text_unwraps_additional_information_list():
+    prompt = {"additional_information": {"text": ["Hello world."]}}
+    assert _extract_text(prompt) == "Hello world."
+
+
+def test_code2wav2aligner_builds_text_prompt_and_audio():
+    prompt = {"additional_information": {"text": ["U.S.A test"], "language": ["English"]}}
+    out = code2wav2aligner([_FakeOutput(_audio_mm())], prompt)
+
+    assert len(out) == 1
+    item = out[0]
+    # text prompt: tokenization + audio features deferred to the stage preprocessor
+    assert "prompt_token_ids" not in item
+    # prompt built from official segmentation: "U.S.A" -> "USA"
+    assert "USA<timestamp><timestamp>test" in item["prompt"]
+    # audio carried as (np.float32, sr) tuple
+    audio, sr = item["multi_modal_data"]["audio"]
+    assert isinstance(audio, np.ndarray) and sr == 24000
+    assert item["additional_information"]["aligner_words"] == [["USA", "test"]]
+    assert item["additional_information"]["language"] == ["English"]
+
+
+def test_code2wav2aligner_skips_unfinished_and_empty():
+    # unfinished -> skipped
+    assert code2wav2aligner([_FakeOutput(_audio_mm(), finished=False)], {}) == []
+    # finished but empty waveform -> skipped
+    empty = _FakeOutput({"audio": torch.zeros(0), "sr": torch.tensor(24000)})
+    assert code2wav2aligner([empty], {"additional_information": {"text": ["hi"]}}) == []
+
+
+def test_decode_pooling_output_produces_word_timestamps_ms():
+    # 2 words -> 4 <timestamp> markers (id=99). classify_num=5, seg=200ms.
+    hf_config = SimpleNamespace(
+        timestamp_token_id=99,
+        timestamp_segment_time=200.0,
+        thinker_config=SimpleNamespace(classify_num=5),
+    )
+    # prompt: marker positions at indices 2,3,5,6
+    prompt_token_ids = [1, 1, 99, 99, 1, 99, 99]
+    request = SimpleNamespace(
+        prompt_token_ids=prompt_token_ids,
+        additional_information={ALIGNER_WORDS_KEY: [["hello", "world"]]},
+    )
+    logits = torch.zeros((7, 5), dtype=torch.float32)
+    logits[2, 0] = 1.0  # word0 start -> bin 0
+    logits[3, 2] = 1.0  # word0 end   -> bin 2
+    logits[5, 2] = 1.0  # word1 start -> bin 2
+    logits[6, 4] = 1.0  # word1 end   -> bin 4
+
+    ms = decode_pooling_output(logits, request, hf_config)
+
+    # decoder returns a bare [n_words, 2] int32 tensor (not a dict)
+    assert ms.tolist() == [[0, 400], [400, 800]]  # bins x 200ms
+
+
+def test_decode_pooling_output_handles_hf_config_without_thinker():
+    # hf_config without thinker_config falls back to itself for classify_num.
+    hf_config = SimpleNamespace(
+        timestamp_token_id=99,
+        timestamp_segment_time=100.0,
+        classify_num=4,
+    )
+    request = SimpleNamespace(
+        prompt_token_ids=[99, 99],
+        additional_information={ALIGNER_WORDS_KEY: [["hi"]]},
+    )
+    logits = torch.zeros((2, 4), dtype=torch.float32)
+    logits[0, 1] = 1.0
+    logits[1, 3] = 1.0
+
+    ms = decode_pooling_output(logits, request, hf_config)
+    assert ms.tolist() == [[100, 300]]

tests/utils/test_forced_aligner.py

@@ -105,9 +105,7 @@ forced_aligner:
 """,
         encoding="utf-8",
     )
-    args = type("Args", (), {"forced_aligner": None, "forced_aligner_config": str(cfg)})()
-
-    out = forced_aligner.build_forced_aligner_config(args)
+    out = forced_aligner.build_forced_aligner_config(None, str(cfg))
 
     assert out == forced_aligner.ForcedAlignerConfig(
         model="Qwen/Qwen3-ForcedAligner-0.6B",
@@ -127,16 +125,7 @@ def test_build_config_cli_model_overrides_yaml(tmp_path):
         "forced_aligner:\n  model: old\n  gpu_memory_utilization: 0.2\n  dtype: float16\n",
         encoding="utf-8",
     )
-    args = type(
-        "Args",
-        (),
-        {
-            "forced_aligner": "new",
-            "forced_aligner_config": str(cfg),
-        },
-    )()
-
-    out = forced_aligner.build_forced_aligner_config(args)
+    out = forced_aligner.build_forced_aligner_config("new", str(cfg))
 
     assert out is not None
     # --forced-aligner overrides the YAML model; gpu_memory_utilization/dtype
@@ -148,16 +137,7 @@ def test_build_config_cli_model_overrides_yaml(tmp_path):
 
 
 def test_build_config_from_cli_model_uses_default_yaml():
-    args = type(
-        "Args",
-        (),
-        {
-            "forced_aligner": "local-aligner",
-            "forced_aligner_config": None,
-        },
-    )()
-
-    out = forced_aligner.build_forced_aligner_config(args)
+    out = forced_aligner.build_forced_aligner_config("local-aligner", None)
 
     assert out is not None
     assert out.model == "local-aligner"

tests/utils/test_qwen3_force_align_processor.py

@@ -46,20 +46,3 @@ def test_fix_timestamp_repairs_dip_onto_monotonic_sequence():
 
 def test_fix_timestamp_passes_through_monotonic():
     assert processor.fix_timestamp([0, 100, 100, 250]) == [0, 100, 100, 250]
-
-
-def test_resolve_timestamp_token_id_defaults_to_marker_token():
-    # Regression: the default must resolve the same <timestamp> marker that
-    # build_prompt inserts, not None (which would degrade every request).
-    seen = {}
-
-    class FakeTokenizer:
-        def convert_tokens_to_ids(self, token):
-            seen["token"] = token
-            return 151705
-
-    tid = processor.resolve_timestamp_token_id(FakeTokenizer())
-
-    assert tid == 151705
-    assert seen["token"] == processor.TIMESTAMP_TOKEN
-    assert processor.TIMESTAMP_TOKEN in processor.build_prompt(["hello", "world"])