[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>

examples/offline_inference/text_to_speech/README.md

@@ -372,18 +372,10 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
 Streaming requires `async_chunk: true` in the stage config.
 
 ### Word Timestamps
-Generate a WAV offline and a JSON sidecar with word-level timestamps from
-`Qwen/Qwen3-ForcedAligner-0.6B`:
-```bash
-python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
-    --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
-    --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B \
-    --text "Hello world." \
-    --output-dir /tmp/qwen3_tts_timestamps
-```
-The script writes `qwen3_tts_word_timestamps.wav` and
-`qwen3_tts_word_timestamps.json`. On machines without a local CUDA toolkit,
-set `VLLM_USE_FLASHINFER_SAMPLER=0` to avoid FlashInfer sampler JIT.
+Word-level timestamps are currently a serving-path feature: launch
+`vllm-omni serve` with `--forced-aligner` and request `word_timestamps`
+(see `examples/online_serving/text_to_speech/README.md`). An offline
+example is not provided in this release.
 
 ### Batched decoding
 The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates):

examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py

@@ -1,172 +0,0 @@
-"""Offline Qwen3-TTS example with word-level timestamps.
-
-This demo runs Qwen3-TTS offline, aligns the synthesized audio in memory with
-the same shared forced aligner used by the streaming server path, then writes
-the WAV and JSON sidecar:
-
-    python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
-        --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
-        --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
-
-On machines without a local CUDA toolkit, set ``VLLM_USE_FLASHINFER_SAMPLER=0``
-to avoid FlashInfer sampler JIT during warmup.
-"""
-
-from __future__ import annotations
-
-import asyncio
-import json
-import os
-from pathlib import Path
-from typing import Any
-
-os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
-# The in-process offline aligner consumes token_classify outputs. Keeping V1
-# multiprocessing off avoids msgspec serialization issues for token-wise logits.
-os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
-
-import numpy as np
-import soundfile as sf
-import torch
-from end2end import _estimate_prompt_len
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-from vllm_omni import Omni
-from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
-from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
-
-
-def _default_stage_config() -> str:
-    repo_root = Path(__file__).resolve().parents[4]
-    return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
-
-
-def _build_custom_voice_input(args: Any) -> dict[str, Any]:
-    additional_information = {
-        "task_type": ["CustomVoice"],
-        "text": [args.text],
-        "language": [args.language],
-        "speaker": [args.speaker],
-        "instruct": [args.instructions],
-        "max_new_tokens": [args.max_new_tokens],
-    }
-    return {
-        "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, args.model),
-        "additional_information": additional_information,
-    }
-
-
-def _audio_tensor_and_sample_rate(mm: dict[str, Any]) -> tuple[torch.Tensor, int]:
-    audio_data = mm["audio"]
-    sr_raw = mm["sr"]
-    sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
-    sample_rate = int(sr_val.item()) if hasattr(sr_val, "item") else int(sr_val)
-    audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
-    return audio_tensor.float().detach().cpu().flatten(), sample_rate
-
-
-def _float_audio_to_pcm16_bytes(audio: torch.Tensor) -> bytes:
-    samples = audio.numpy()
-    pcm = (np.clip(samples, -1.0, 1.0) * 32767.0).astype("<i2")
-    return pcm.tobytes()
-
-
-async def _run_alignment(args: Any, audio: torch.Tensor, sample_rate: int) -> list[dict[str, Any]] | None:
-    aligner_config = build_forced_aligner_config(args)
-    if aligner_config is None:
-        raise ValueError("--forced-aligner or --forced-aligner-config is required")
-    timestamps = await align(
-        audio=_float_audio_to_pcm16_bytes(audio),
-        text=args.text,
-        sample_rate=sample_rate,
-        config=aligner_config,
-        language=args.language,
-    )
-    if timestamps is None:
-        return None
-    return [
-        {
-            "word": item.word,
-            "start_ms": item.start_ms,
-            "end_ms": item.end_ms,
-        }
-        for item in timestamps
-    ]
-
-
-def main(args: Any) -> None:
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    omni_kwargs = vars(args).copy()
-    for key in (
-        "forced_aligner",
-        "text",
-        "language",
-        "speaker",
-        "instructions",
-        "max_new_tokens",
-        "output_dir",
-    ):
-        omni_kwargs.pop(key, None)
-    omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
-    omni_kwargs["log_stats"] = args.log_stats
-    omni = Omni(**omni_kwargs)
-
-    prompt = _build_custom_voice_input(args)
-    final_output = None
-    for stage_outputs in omni.generate([prompt]):
-        final_output = stage_outputs.request_output
-    if final_output is None:
-        raise RuntimeError("Qwen3-TTS did not produce an output.")
-
-    mm = final_output.outputs[0].multimodal_output
-    audio, sample_rate = _audio_tensor_and_sample_rate(mm)
-
-    wav_path = output_dir / "qwen3_tts_word_timestamps.wav"
-    # Align the in-memory tensor, not the saved WAV. The WAV is only a demo
-    # artifact so users can listen to the same audio referenced by the sidecar.
-    timestamps = asyncio.run(_run_alignment(args, audio, sample_rate))
-    sf.write(wav_path, audio.numpy(), samplerate=sample_rate, format="WAV")
-
-    sidecar = {
-        "text": args.text,
-        "sample_rate": sample_rate,
-        "audio_path": str(wav_path),
-        "timestamps": timestamps,
-    }
-    json_path = output_dir / "qwen3_tts_word_timestamps.json"
-    json_path.write_text(json.dumps(sidecar, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
-
-    print(f"Saved audio: {wav_path}")
-    print(f"Saved timestamps: {json_path}")
-    print(json.dumps(sidecar["timestamps"], ensure_ascii=False, indent=2))
-
-
-def parse_args() -> Any:
-    parser = FlexibleArgumentParser(description="Offline Qwen3-TTS word timestamp example")
-    parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", help="Qwen3-TTS model path/name")
-    parser.add_argument(
-        "--forced-aligner",
-        default=None,
-        help="Qwen3 forced aligner model path/name",
-    )
-    parser.add_argument(
-        "--forced-aligner-config",
-        default=None,
-        help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
-    )
-    parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
-    parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
-    parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
-    parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
-    parser.add_argument("--instructions", default="", help="Optional speaking style instruction")
-    parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
-    parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
-    parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
-    nullify_stage_engine_defaults(parser)
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    main(parse_args())

examples/online_serving/text_to_speech/README.md

@@ -605,6 +605,13 @@ python qwen3_tts/streaming_speech_client.py \
 The client writes one PCM file per sentence and a matching
 `sentence_XXX_timestamps.json` sidecar.
 
+Non-streaming requests can also ask for timestamps: pass
+`"word_timestamps": true` to `POST /v1/audio/speech` and read the
+`X-Word-Timestamps` response header (JSON list of `{word, start_ms, end_ms}`,
+ASCII-escaped). Past 4 KB the header is replaced by
+`X-Word-Timestamps-Omitted: oversize; bytes=<n>; limit=4096` and the audio
+still returns — use the WebSocket path for long transcripts.
+
 To *see* the alignment instead of reading a JSON sidecar, run the
 word-timestamp Gradio demo (server must be launched with `--forced-aligner`):
 ```bash

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"])

vllm_omni/config/config_factory.py

@@ -409,6 +409,10 @@ class StageConfigFactory:
         if cli_async_chunk is not None:
             deploy_cfg.async_chunk = bool(cli_async_chunk)
 
+        from vllm_omni.utils.forced_aligner import inject_forced_aligner_stage
+
+        pipeline_cfg, deploy_cfg = inject_forced_aligner_stage(pipeline_cfg, deploy_cfg, cli_overrides)
+
         stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides)
 
         # Overlay declarative parallel strategies (opt-in) before CLI overrides.

vllm_omni/config/model.py

@@ -129,6 +129,9 @@ class OmniModelConfig(ModelConfig):
     model_arch: str | None = None
     worker_type: str | None = None
     engine_output_type: str | None = None
+    # Optional dotted path of a per-stage pooling-output decoder applied
+    # worker-side before IPC. Read by the AR scheduler.
+    pooling_output_decoder: str | None = None
     hf_config_name: str | None = None
     custom_process_next_stage_input_func: str | None = None
     stage_connector_config: dict[str, Any] = field(

vllm_omni/config/stage_config.py

@@ -340,6 +340,7 @@ class StageDeployConfig:
     output_connectors: dict[str, str] | None = None
     input_connectors: dict[str, str] | None = None
     default_sampling_params: dict[str, Any] | None = None
+    default_pooling_params: dict[str, Any] | None = None
     subtalker_sampling_params: dict[str, Any] | None = None
 
     # === Generic stage engine fields ===
@@ -509,6 +510,7 @@ _STAGE_RESERVED_KEYS = frozenset(
         "output_connectors",
         "input_connectors",
         "default_sampling_params",
+        "default_pooling_params",
         "engine_extras",
         "engine_args",
         "runtime",
@@ -562,11 +564,14 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
     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")
+    kwargs["default_pooling_params"] = stage_data.get("default_pooling_params")
     kwargs["engine_extras"] = _get_recursively_merged_dict(explicit_engine_extras, flat_args)
     return StageDeployConfig(**kwargs)
 
 
-_DEEP_MERGE_KEYS = frozenset({"default_sampling_params", "subtalker_sampling_params", "engine_extras", "engine_args"})
+_DEEP_MERGE_KEYS = frozenset(
+    {"default_sampling_params", "default_pooling_params", "subtalker_sampling_params", "engine_extras", "engine_args"}
+)
 
 
 def _deep_merge_stage(base: dict, overlay: dict) -> dict:
@@ -886,6 +891,8 @@ def _build_extras(
     sampling.update(ps.sampling_constraints)
     if sampling:
         extras["default_sampling_params"] = sampling
+    if ds is not None and ds.default_pooling_params:
+        extras["default_pooling_params"] = dict(ds.default_pooling_params)
     if ds is not None and ds.output_connectors:
         extras["output_connectors"] = dict(ds.output_connectors)
     if ds is not None and ds.input_connectors:

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -415,6 +415,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             finish_reason = None
             routed_experts = None
 
+            # Decode the pooling output before stop handling so a decoder
+            # failure finishes the request with FinishReason.ERROR (500).
+            try:
+                pooling_output_payload = self._maybe_decode_pooling_output(request, pooler_output)
+            except Exception as exc:
+                logger.exception("[pooling] decoder hook failed for request %s", req_id)
+                pooling_output_payload = None
+                request.status = RequestStatus.FINISHED_ERROR
+                request.stop_reason = f"pooling output decode failed: {exc}"
+                request.resumable = False
+
             # Check for stop and update request status.
             if new_token_ids:
                 num_sampled_tokens = len(new_token_ids)
@@ -428,7 +439,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
             elif request.pooling_params and pooler_output is not None:
                 # Pooling stops as soon as there is output.
-                request.status = RequestStatus.FINISHED_STOPPED
+                if request.status != RequestStatus.FINISHED_ERROR:
+                    request.status = RequestStatus.FINISHED_STOPPED
                 stopped = True
 
             # If criteria returns True, it means we must STOP the request.
@@ -522,7 +534,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     finish_reason=finish_reason,
                     new_logprobs=new_logprobs,
                     new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                    pooling_output=pooler_output,
+                    pooling_output=pooling_output_payload,
                     multimodal_output=mm_output,
                     stop_reason=request.stop_reason,
                     prefill_stats=request.take_prefill_stats(),

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -5,6 +5,7 @@ from collections import defaultdict
 from typing import Any
 
 from vllm.compilation.cuda_graph import CUDAGraphStat
+from vllm.logger import init_logger
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.interface import PauseState
 from vllm.v1.core.sched.output import SchedulerOutput
@@ -25,6 +26,8 @@ from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestDat
 from vllm_omni.core.sched.utils import omni_routed_experts_for_request
 from vllm_omni.outputs import OmniModelRunnerOutput
 
+logger = init_logger(__name__)
+
 
 class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
     def __init__(self, *args, **kwargs):
@@ -442,10 +445,22 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             is_segment_finished = False
             routed_experts = None
 
+            # Decode the pooling output before stop handling so a decoder
+            # failure finishes the request with FinishReason.ERROR (500).
+            try:
+                pooling_output_payload = self._maybe_decode_pooling_output(request, pooler_output)
+            except Exception as exc:
+                logger.exception("[pooling] decoder hook failed for request %s", req_id)
+                pooling_output_payload = None
+                request.status = RequestStatus.FINISHED_ERROR
+                request.stop_reason = f"pooling output decode failed: {exc}"
+                request.resumable = False
+
             # One-shot generation request: finish after its current input unit
             # has been fully processed.
             if (
-                request.status == RequestStatus.FINISHED_STOPPED
+                request.status == RequestStatus.FINISHED_ERROR
+                or request.status == RequestStatus.FINISHED_STOPPED
                 or (self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens)
                 or (
                     self.chunk_transfer_adapter is not None
@@ -453,7 +468,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     and request.num_computed_tokens >= len(request.prompt_token_ids)
                 )
             ):
-                request.status = RequestStatus.FINISHED_STOPPED
+                if request.status != RequestStatus.FINISHED_ERROR:
+                    request.status = RequestStatus.FINISHED_STOPPED
                 # Optional: set a stop_reason for front-end clarity
                 # (does not affect protocol)
                 stopped = True
@@ -508,7 +524,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     finish_reason=finish_reason,
                     new_logprobs=new_logprobs,
                     new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                    pooling_output=pooler_output,
+                    pooling_output=pooling_output_payload,
                     multimodal_output=mm_output,
                     stop_reason=request.stop_reason,
                     prefill_stats=request.take_prefill_stats(),

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -5,10 +5,12 @@ import time
 from collections.abc import Iterable
 from typing import Any
 
+import torch
 from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.distributed.kv_events import KVEventBatch
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
+from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.utils import remove_all
 from vllm.v1.engine import (
@@ -80,6 +82,29 @@ class OmniSchedulerMixin:
             else None
         )
         self._latest_omni_connector_output = None
+        # Optional per-stage pooling-output decoder hook (dotted path in
+        # model_config); applied worker-side before IPC.
+        self._pooling_output_decoder = None
+        _decoder_path = getattr(model_config, "pooling_output_decoder", None)
+        if _decoder_path:
+            self._pooling_output_decoder = resolve_obj_by_qualname(str(_decoder_path))
+
+    def _maybe_decode_pooling_output(self, request: Request, pooler_output: Any) -> Any:
+        """Apply the stage's pooling-output decoder hook to the pooler tensor
+        before IPC, or pass it through unchanged when none is configured.
+        Decoder exceptions propagate; callers fail the request with
+        FinishReason.ERROR rather than emitting an empty success."""
+        if self._pooling_output_decoder is None:
+            return pooler_output
+        if pooler_output is None or getattr(request, "pooling_params", None) is None:
+            return pooler_output
+        if not isinstance(pooler_output, torch.Tensor):
+            return pooler_output
+        return self._pooling_output_decoder(
+            pooler_output,
+            request,
+            self.vllm_config.model_config.hf_config,
+        )
 
     def _free_input_coordinator_request(self, request_id: str) -> None:
         """Prune full-payload coordinator state for a completed request."""

vllm_omni/deploy/qwen3_tts_forced_aligner.yaml

@@ -1,10 +1,12 @@
-# Reference deploy knobs for Qwen/Qwen3-ForcedAligner-0.6B, the shared forced
-# aligner for streaming TTS word timestamps. Pass with --forced-aligner-config;
-# --forced-aligner overrides the model path without editing this file.
+# Default deploy knobs for Qwen/Qwen3-ForcedAligner-0.6B, the TTS word-timestamp
+# aligner. These configure the pooling stage that --forced-aligner appends to
+# the pipeline tail (see utils.forced_aligner.inject_forced_aligner_stage); they are
+# NOT a standalone engine. Pass --forced-aligner-config to override this file;
+# --forced-aligner sets just the model path.
 #
-# Run as LLM(runner="pooling") with pooling_task="token_classify". Prompt tokens,
-# chat template and word segmentation are Qwen-specific and live in
-# qwen3_force_align_processor (not configurable here).
+# The stage runs the model with runner="pooling" + pooling_task="token_classify".
+# Prompt tokens, chat template and word segmentation are Qwen-specific and live
+# in qwen3_force_align_processor (not configurable here).
 
 forced_aligner:
   model: Qwen/Qwen3-ForcedAligner-0.6B

vllm_omni/engine/arg_utils.py

@@ -191,6 +191,8 @@ class OmniEngineArgs(EngineArgs):
     quantization_config: Any | None = None
     force_cutlass_fp8: bool | None = None
     worker_type: str | None = None
+    # Dotted path of a per-stage pooling-output decoder applied worker-side.
+    pooling_output_decoder: str | None = None
     task_type: str | None = None
     worker_cls: str = None  # type: ignore[assignment]  # Upstream default is "auto"; omni resolves
     # in __post_init__ based on worker_type (ar/generation), so None is safe here.
@@ -380,6 +382,7 @@ class OmniEngineArgs(EngineArgs):
             model_stage=self.model_stage,
             model_arch=self.model_arch,
             worker_type=self.worker_type,
+            pooling_output_decoder=self.pooling_output_decoder,
             engine_output_type=self.engine_output_type,
             hf_config_name=self.hf_config_name,
             custom_process_next_stage_input_func=self.custom_process_next_stage_input_func,
@@ -475,6 +478,11 @@ class OrchestratorArgs:
     # === Mode Switches (orchestrator reads, DeployConfig redistributes) ===
     async_chunk: bool | None = None
 
+    # === Forced aligner (orchestrator injects a pooling stage; never a per-stage knob) ===
+    forced_aligner: str | None = None
+    forced_aligner_config: str | None = None
+    forced_aligner_device: str | None = None
+
     # === Observability ===
     log_stats: bool = False
     enable_orch_monitor: bool = False

vllm_omni/engine/orchestrator.py

@@ -1543,11 +1543,19 @@ class Orchestrator:
             return request
 
         processor = self._get_stage_input_processor(next_stage_id)
+        # A pooling stage is driven by PoolingParams; vLLM validates them against
+        # the stage's supported tasks, so advertise the model's pooling tasks (or
+        # the configured task) instead of generate.
+        if isinstance(params, PoolingParams):
+            model_cfg = next_pool.stage_vllm_config.model_config
+            supported_tasks = tuple(getattr(model_cfg, "supported_tasks", ()) or ()) or (params.task,)
+        else:
+            supported_tasks = ("generate",)
         request = processor.process_inputs(
             request_id=req_id,
             prompt=next_input,
             params=params,
-            supported_tasks=("generate",),
+            supported_tasks=supported_tasks,
             arrival_time=_time.time(),
             resumable=resumable,
         )

vllm_omni/engine/stage_init_utils.py

@@ -20,6 +20,7 @@ from dataclasses import dataclass, fields, replace
 from typing import Any, Literal, cast
 
 from vllm.logger import init_logger
+from vllm.pooling_params import PoolingParams
 from vllm.renderers import BaseRenderer
 from vllm.sampling_params import SamplingParams
 from vllm.tokenizers import cached_tokenizer_from_config
@@ -395,8 +396,19 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata:
     final_output_type: str | None = stage_config.final_output_type
 
     default_sp = _to_dict(_get_attr_or_item(stage_config, "default_sampling_params", {}))
-    SPClass = SamplingParams if stage_type == "llm" else OmniDiffusionSamplingParams
-    default_sampling_params: OmniSamplingParams = SPClass(**default_sp)
+    # A pooling stage carries its task via default_pooling_params, set where the
+    # stage is declared.
+    default_pp = _to_dict(_get_attr_or_item(stage_config, "default_pooling_params", {}))
+    # A pooling stage is an LLM stage run with runner="pooling" (vLLM's
+    # is_pooling_model signal); pick params by that signal, not execution_type.
+    is_pooling = str(engine_args.get("runner", "")).lower() == "pooling"
+    default_params: OmniSamplingParams | PoolingParams
+    if stage_type == "diffusion":
+        default_params = OmniDiffusionSamplingParams(**default_sp)
+    elif is_pooling:
+        default_params = PoolingParams(**default_pp)
+    else:  # generative llm: ar / generation
+        default_params = SamplingParams(**default_sp)
 
     custom_process_input_func: Callable | None = None
     _cpif_path = _get_attr_or_item(stage_config, "custom_process_input_func")
@@ -428,7 +440,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata:
             engine_input_source=engine_input_source,
             final_output=final_output,
             final_output_type=final_output_type,
-            default_sampling_params=default_sampling_params,
+            default_sampling_params=default_params,
             custom_process_input_func=custom_process_input_func,
             model_stage=model_stage,
             runtime_cfg=runtime_cfg,
@@ -448,7 +460,7 @@ def extract_legacy_stage_metadata(stage_config: Any) -> StageMetadata:
         engine_input_source=engine_input_source,
         final_output=final_output,
         final_output_type=final_output_type,
-        default_sampling_params=default_sampling_params,
+        default_sampling_params=default_params,
         custom_process_input_func=custom_process_input_func,
         model_stage=model_stage,
         runtime_cfg=runtime_cfg,

vllm_omni/entrypoints/cli/serve.py

@@ -237,9 +237,8 @@ class OmniServeCommand(CLISubcommand):
             "VoiceDesign, or Base, while diffusion models may use it to select "
             "task-specific weights. If omitted, the model default is used.",
         )
-        # Forced aligner / word timestamps. --forced-aligner is the opt-in
-        # toggle; heavier knobs (gpu_memory_utilization, dtype, max_model_len)
-        # live in the deploy YAML passed via --forced-aligner-config.
+        # Forced aligner / word timestamps. Either flag opts in (--forced-aligner-config
+        # alone works when its YAML sets the model); heavier knobs live in that YAML.
         omni_config_group.add_argument(
             "--forced-aligner",
             type=str,
@@ -260,6 +259,13 @@ class OmniServeCommand(CLISubcommand):
                 "flag, when set, overrides the YAML model field."
             ),
         )
+        omni_config_group.add_argument(
+            "--forced-aligner-device",
+            type=str,
+            default=None,
+            help="Device(s) for the forced-aligner stage (e.g. '2'). Defaults to "
+            "sharing an existing stage's GPU when unset.",
+        )
         omni_config_group.add_argument(
             "--deploy-config",
             type=str,

vllm_omni/entrypoints/openai/api_server.py

@@ -1127,7 +1127,11 @@ async def omni_init_app_state(
         state.openai_serving_models,
         request_logger=request_logger,
         model_name=model_name,
-        forced_aligner_config=build_forced_aligner_config(args),
+        forced_aligner_enabled=build_forced_aligner_config(
+            getattr(args, "forced_aligner", None),
+            getattr(args, "forced_aligner_config", None),
+        )
+        is not None,
     )
 
     # Warm up speech pipeline (CUDA Graph capture, torch.compile) so the first

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -161,10 +161,12 @@ class OpenAICreateSpeechRequest(BaseModel):
     word_timestamps: bool = Field(
         default=False,
         description=(
-            "When true, the server runs a shared forced aligner alongside the streamed "
-            "audio and emits per-chunk word timestamps. Requires the server to be "
-            "launched with --forced-aligner pointing at an aligner model. No effect "
-            "when streaming is off."
+            "When true, non-streaming responses carry per-word timestamps in the "
+            "X-Word-Timestamps header (JSON list of {word, start_ms, end_ms}, "
+            "ASCII-escaped; replaced by X-Word-Timestamps-Omitted past 4 KB). "
+            "Requires the server to be launched with --forced-aligner (400 otherwise). "
+            "Not supported with stream=true; for streaming use the WebSocket "
+            "/v1/audio/speech/stream path."
         ),
     )
 

vllm_omni/entrypoints/openai/serving_speech.py

@@ -423,7 +423,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
     def __init__(self, *args, **kwargs):
         self.model_name = kwargs.pop("model_name", None)
-        self.forced_aligner_config: Any | None = kwargs.pop("forced_aligner_config", None)
+        # True when the server was launched with --forced-aligner (a pooling
+        # aligner stage is appended to the pipeline). Gates word_timestamps.
+        self.forced_aligner_enabled: bool = bool(kwargs.pop("forced_aligner_enabled", False))
         super().__init__(*args, **kwargs)
         self._init_speaker_storage()
 
@@ -2186,6 +2188,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request_start_s: float | None = None,
         include_sample_rate: bool = False,
         usage_acc: SpeechOutputTokenCounter | None = None,
+        collect: dict | None = None,
     ):
         """Generate audio chunks for streaming response.
 
@@ -2218,6 +2221,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     usage_acc.observe(res)
                 audio_output, audio_key = self._extract_audio_output(res)
                 if audio_key is None:
+                    # Stash the aligner's timestamps output for streaming callers.
+                    if collect is not None and self._is_timestamps_output(res):
+                        collect["aligner_res"] = res
                     continue
 
                 sr_raw = audio_output.get("sr")
@@ -2406,6 +2412,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             data = json.dumps(payload, separators=(",", ":"))
             yield f"event: speech.audio.error\ndata: {data}\n\n"
 
+    @staticmethod
+    def _is_timestamps_output(res) -> bool:
+        """True when ``res`` is the forced-aligner stage's terminal timestamps output."""
+        from vllm_omni.model_executor.stage_input_processors.forced_aligner import TIMESTAMPS_MODALITY
+
+        return getattr(res, "final_output_type", None) == TIMESTAMPS_MODALITY
+
     @staticmethod
     def _extract_audio_output(res) -> tuple[dict | None, str | None]:
         """Return (audio_output dict, audio key) or (None, None).
@@ -3332,11 +3345,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     stage0_params.extra_args = {}
                 stage0_params.extra_args.setdefault("tts_local_seed", int(default_seed))
 
+        # When word_timestamps is requested, also ask for the aligner stage's
+        # output so the orchestrator drives the request through the forced-aligner
+        # stage (final_stage_id extends to it). Harmless if no aligner stage exists.
+        output_modalities = ["audio"]
+        if getattr(request, "word_timestamps", False):
+            from vllm_omni.model_executor.stage_input_processors.forced_aligner import TIMESTAMPS_MODALITY
+
+            output_modalities.append(TIMESTAMPS_MODALITY)
+
         generator = self.engine_client.generate(
             prompt=prompt,
             request_id=request_id,
             sampling_params_list=sampling_params_list,
-            output_modalities=["audio"],
+            output_modalities=output_modalities,
         )
         self._track_ref_audio_artifact_warmup(
             request_id,
@@ -3345,17 +3367,22 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         )
         return request_id, generator, tts_params
 
-    async def _generate_pcm_chunks(self, generator, request_id: str, *, include_sample_rate: bool = False):
+    async def _generate_pcm_chunks(
+        self, generator, request_id: str, *, include_sample_rate: bool = False, collect: dict | None = None
+    ):
         """Yield raw PCM byte chunks from the engine generator.
 
         Delegates to ``_generate_audio_chunks`` with ``response_format="pcm"``.
         Used by the WebSocket streaming handler and ``_iter_pcm_audio_bytes``.
+        ``collect`` (when given) receives the forced-aligner stage's pooling
+        output under ``"aligner_res"`` for downstream word-timestamp extraction.
         """
         async for chunk in self._generate_audio_chunks(
             generator,
             request_id,
             response_format="pcm",
             include_sample_rate=include_sample_rate,
+            collect=collect,
         ):
             yield chunk
 
@@ -3375,6 +3402,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         request_id: str | None = None,
         usage_out: list[SpeechTokenUsage] | None = None,
         has_inline_ref_audio: bool | None = None,
+        collect: dict | None = None,
     ) -> tuple[bytes | str, str]:
         # ``usage_out`` is an opt-in output channel: when a list is passed, the
         # computed SpeechTokenUsage is appended to it. The return stays a
@@ -3398,9 +3426,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             # Non-streaming is FINAL_ONLY, so the stage-0 output carries the full
             # token sequence; the counter records its length for output_tokens.
             usage_acc = SpeechOutputTokenCounter()
+            audio_res: OmniRequestOutput | None = None
+            aligner_res: OmniRequestOutput | None = None
             async for res in generator:
                 final_output = res
                 usage_acc.observe(res)
+                # The generator yields both the audio output (Code2Wav) and, with
+                # a forced-aligner stage, a timestamps output. Keep the audio res
+                # for the WAV and the aligner res for word timestamps.
+                if self._is_timestamps_output(res):
+                    aligner_res = res
+                else:
+                    _, audio_key = self._extract_audio_output(res)
+                    if audio_key is not None:
+                        audio_res = res
                 if not is_moss:
                     continue
                 try:
@@ -3422,10 +3461,26 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             if final_output is None:
                 raise ValueError("No output generated from the model.")
 
-            audio_output, audio_key = self._extract_audio_output(final_output)
+            # Extract audio from the audio-bearing res (not necessarily the last
+            # yielded one, which may be the aligner's timestamps output).
+            audio_source = audio_res if audio_res is not None else final_output
+            audio_output, audio_key = self._extract_audio_output(audio_source)
             if audio_key is None:
                 raise ValueError("TTS model did not produce audio output.")
 
+            # Surface forced-aligner word timestamps to the caller (set as a
+            # response header) when requested and an aligner stage produced them.
+            if collect is not None and getattr(request, "word_timestamps", False):
+                from vllm_omni.utils.forced_aligner import extract_word_timestamps
+
+                ts = (
+                    extract_word_timestamps(aligner_res, request.input, getattr(request, "language", None))
+                    if aligner_res is not None
+                    else None
+                )
+                if ts is not None:
+                    collect["word_timestamps"] = ts
+
             audio_tensor = audio_output[audio_key]
             sr_raw = audio_output.get("sr", 24000)
             sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
@@ -3695,6 +3750,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     "/v1/audio/speech/stream path. Use session.config with "
                     "stream_audio=true and response_format='pcm'."
                 )
+            if request.word_timestamps and not self.forced_aligner_enabled:
+                # Fail loud instead of silently returning 200 with no timestamps header.
+                return self.create_error_response(
+                    "word_timestamps=true requires the server to be launched with --forced-aligner."
+                )
 
             if request.is_raw_audio_stream():
                 response_format, error = self._validate_speech_streaming_request(
@@ -3739,7 +3799,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     media_type="text/event-stream",
                 )
 
-            audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id)
+            collect: dict = {}
+            audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id, collect=collect)
             total_ms = (time.perf_counter() - request_start_s) * 1000.0
             logger.info(
                 "[SpeechE2E] request_id=%s stream=false status=ok total_ms=%.2f response_bytes=%d",
@@ -3747,7 +3808,22 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 total_ms,
                 len(audio_bytes) if isinstance(audio_bytes, (bytes, bytearray)) else len(str(audio_bytes)),
             )
-            return Response(content=audio_bytes, media_type=media_type)
+            headers = {}
+            if collect.get("word_timestamps") is not None:
+                # Default ensure_ascii keeps the header latin-1 encodable (non-ASCII words \uXXXX-escaped).
+                ts_json = json.dumps(collect["word_timestamps"])
+                # Cap at 4 KB: oversized headers turn into opaque 502s at common reverse-proxy defaults.
+                if len(ts_json) <= 4096:
+                    headers["X-Word-Timestamps"] = ts_json
+                else:
+                    # Marker header so clients can tell an oversized alignment from no alignment.
+                    headers["X-Word-Timestamps-Omitted"] = f"oversize; bytes={len(ts_json)}; limit=4096"
+                    logger.warning(
+                        "X-Word-Timestamps header omitted: %d bytes exceeds the 4 KB budget "
+                        "(use the WebSocket streaming path for long transcripts)",
+                        len(ts_json),
+                    )
+            return Response(content=audio_bytes, media_type=media_type, headers=headers)
 
         except asyncio.CancelledError:
             total_ms = (time.perf_counter() - request_start_s) * 1000.0

vllm_omni/entrypoints/openai/serving_speech_stream.py

@@ -61,8 +61,7 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
     StreamingSpeechSessionConfig,
 )
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
-from vllm_omni.utils.forced_aligner import ForcedAlignerLoadError
-from vllm_omni.utils.forced_aligner import align as forced_align
+from vllm_omni.utils.forced_aligner import extract_word_timestamps
 
 logger = init_logger(__name__)
 
@@ -273,7 +272,7 @@ class OmniStreamingSpeechHandler:
 
         # Reject unmet word-timestamps preconditions early with a clear reason.
         if config.word_timestamps:
-            if self._speech_service.forced_aligner_config is None:
+            if not self._speech_service.forced_aligner_enabled:
                 await self._send_error(
                     websocket,
                     "word_timestamps=true but the server was launched without "
@@ -285,7 +284,8 @@ class OmniStreamingSpeechHandler:
                 await self._send_error(
                     websocket,
                     "word_timestamps=true requires stream_audio=true and "
-                    "response_format='pcm' (the aligner consumes raw PCM).",
+                    "response_format='pcm' (timestamps ride the per-sentence "
+                    "PCM audio.chunk stream).",
                 )
                 return
 
@@ -395,19 +395,20 @@ class OmniStreamingSpeechHandler:
     ) -> int:
         """Stream PCM as JSON ``audio.chunk`` frames, aligned per sentence.
 
-        Forward each PCM chunk live (``timestamps: null``) while buffering the
-        sentence audio, then run the forced aligner once over the whole
-        sentence and emit a final empty-audio ``audio.chunk`` with the word
-        timestamps. On aligner failure timestamps is ``null``; for silence it
-        is ``[]`` (audio always flows regardless).
+        Forward each PCM chunk live (``timestamps: null``). The forced-aligner
+        pipeline stage (appended when the server is launched with
+        ``--forced-aligner``) consumes the synthesized audio internally and its
+        pooling output rides the same generator, so once the audio finishes we
+        pull the word timestamps straight off that aligner output and emit a
+        final empty-audio ``audio.chunk`` carrying them. Timestamps is ``null``
+        when the aligner produced none; audio always flows regardless.
         """
-        aligner_config = self._speech_service.forced_aligner_config
-        assert aligner_config is not None  # gated by the precondition check
-
-        sentence_audio = bytearray()
+        audio_bytes_seen = 0
         total_bytes = 0
         sample_rate = _PCM_SAMPLE_RATE
         chunk_id = 0
+        # Receives the aligner stage's pooling output from the generator.
+        collect: dict = {}
 
         async def send_chunk(
             chunk: bytes,
@@ -437,62 +438,30 @@ class OmniStreamingSpeechHandler:
                 generator,
                 request_id,
                 include_sample_rate=True,
+                collect=collect,
             )
         ) as stream:
             async for chunk, chunk_sample_rate in stream:
                 sample_rate = chunk_sample_rate
-                chunk_start_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0))
-                sentence_audio.extend(chunk)
-                chunk_end_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0))
+                chunk_start_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0))
+                audio_bytes_seen += len(chunk)
+                chunk_end_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0))
                 total_bytes += len(chunk)
                 # Audio first, timestamps after the whole sentence is aligned.
                 await send_chunk(chunk, chunk_sample_rate, None, chunk_start_ms, chunk_end_ms)
 
-        # Single alignment pass over the full sentence, then emit timestamps.
-        # A load/config failure is permanent, so surface the reason once; audio
-        # has already streamed, so the trailing frame still carries null.
-        try:
-            timestamps_payload = await self._align_sentence(
-                audio=bytes(sentence_audio),
-                text=sentence_text,
-                sample_rate=sample_rate,
-                config=aligner_config,
-                language=language,
-            )
-        except ForcedAlignerLoadError as exc:
-            await self._send_error(websocket, f"forced aligner unavailable: {exc}")
-            timestamps_payload = None
-        sentence_end_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0))
+        # Pull word timestamps off the aligner stage's pooling output (it rode
+        # the same generator); extract_word_timestamps re-segments the sentence
+        # text for the word strings when the aligner output doesn't carry them.
+        aligner_res = collect.get("aligner_res")
+        timestamps_payload = (
+            extract_word_timestamps(aligner_res, sentence_text, language) if aligner_res is not None else None
+        )
+        sentence_end_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0))
         await send_chunk(b"", sample_rate, timestamps_payload, 0, sentence_end_ms)
 
         return total_bytes
 
-    @staticmethod
-    async def _align_sentence(
-        *,
-        audio: bytes,
-        text: str,
-        sample_rate: int,
-        config,
-        language: str | None = None,
-    ) -> list[dict] | None:
-        """Convert a sentence alignment into JSON word-timestamp dicts.
-
-        Returns ``None`` on aligner failure, ``[]`` when it ran but produced no
-        tokens. Monotonic, non-overlapping bounds are guaranteed by the decoder.
-        ``language`` is forwarded to word segmentation.
-        """
-        aligned = await forced_align(
-            audio=audio,
-            text=text,
-            sample_rate=sample_rate,
-            config=config,
-            language=language,
-        )
-        if aligned is None:
-            return None
-        return [{"word": ts.word, "start_ms": ts.start_ms, "end_ms": ts.end_ms} for ts in aligned]
-
     @staticmethod
     async def _send_error(websocket: WebSocket, message: str) -> None:
         """Send an error message to the client."""

vllm_omni/model_executor/stage_input_processors/forced_aligner.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""stage input processor for a forced-aligner stage.
+
+Bridges a Code2Wav (audio) stage to a pooling forced-aligner stage:
+
+    Talker (LLM_AR) -> Code2Wav (LLM_GENERATION) -> ForcedAligner (LLM_AR + runner=pooling)
+
+``code2wav2aligner`` is a synchronous ``custom_process_input_func`` (the same
+shape as ``qwen3_tts.talker2code2wav``). It runs in the orchestrator process,
+reads the synthesized waveform from the upstream Code2Wav output and the
+ground-truth text from the original request, then builds the aligner prompt
+(``<timestamp>`` markers per word) + audio multimodal input.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.model_executor.stage_input_processors.tts_utils import (
+    extract_language_from_prompt,
+)
+from vllm_omni.utils.qwen3_force_align_processor import build_prompt, segment_words
+
+logger = init_logger(__name__)
+
+# Shared keys / names for the forced-aligner stage payloads, so the producer
+# (this processor), the worker-side decode, and the serving consumer don't
+# duplicate magic strings.
+ALIGNER_WORDS_KEY = "aligner_words"
+ALIGNER_DURATION_MS_KEY = "aligner_audio_duration_ms"
+TIMESTAMPS_MODALITY = "timestamps"
+# Dotted path of the pooling-output decoder hook (read by the AR scheduler for
+# a stage whose engine_args set ``pooling_output_decoder``).
+POOLING_OUTPUT_DECODER_PATH = "vllm_omni.model_executor.stage_input_processors.forced_aligner.decode_pooling_output"
+
+
+def decode_pooling_output(logits: Any, request: Any, hf_config: Any) -> Any:
+    """Worker-side decoder for the forced-aligner pooling stage.
+
+    Converts the aligner's ``[n_token, classify_num]`` logits into a bare
+    ``[n_words, 2]`` int32 tensor of per-word ``[start_ms, end_ms]`` (word
+    strings travel separately via ``additional_information``).
+    """
+    from vllm_omni.engine.serialization import deserialize_additional_information
+    from vllm_omni.utils.forced_aligner import _decode_timestamps
+
+    ts_id = int(getattr(hf_config, "timestamp_token_id"))
+    thinker_cfg = getattr(hf_config, "thinker_config", hf_config)
+    classify_num = int(getattr(thinker_cfg, "classify_num"))
+    seg_ms = float(getattr(hf_config, "timestamp_segment_time"))
+
+    prompt_token_ids = list(getattr(request, "prompt_token_ids", []) or [])
+    positions = [i for i, t in enumerate(prompt_token_ids) if t == ts_id]
+
+    info = deserialize_additional_information(getattr(request, "additional_information", None)) or {}
+    words_field = info.get(ALIGNER_WORDS_KEY)
+    words = words_field[0] if isinstance(words_field, list) and words_field else []
+
+    # Real waveform duration (stashed by code2wav2aligner) drives the decode clamp; fall back to the
+    # theoretical ceiling (clamp no-op) if absent.
+    duration_field = info.get(ALIGNER_DURATION_MS_KEY)
+    duration_ms = (
+        float(duration_field[0])
+        if isinstance(duration_field, list) and duration_field
+        else float(classify_num) * seg_ms
+    )
+
+    timestamps = _decode_timestamps(
+        logits=logits,
+        words=list(words),
+        timestamp_positions=positions,
+        classify_num=classify_num,
+        timestamp_segment_time_ms=seg_ms,
+        audio_duration_ms=duration_ms,
+    )
+    # No failure sentinel: an empty [0, 2] tensor decodes downstream to timestamps: null.
+    ms = torch.tensor(
+        [[t.start_ms, t.end_ms] for t in timestamps],
+        dtype=torch.int32,
+    ).reshape(-1, 2)
+    return ms
+
+
+def _extract_waveform(multimodal_output: dict[str, Any]) -> tuple[np.ndarray, int]:
+    """Pull a mono float32 waveform + sample rate from a Code2Wav output."""
+    wav = multimodal_output.get("audio")
+    if wav is None:
+        wav = multimodal_output.get("model_outputs")
+    if isinstance(wav, list):
+        wav = wav[0] if wav else None
+    if isinstance(wav, torch.Tensor):
+        wav_np = wav.detach().cpu().to(torch.float32).flatten().numpy()
+    elif isinstance(wav, np.ndarray):
+        wav_np = wav.astype(np.float32).flatten()
+    else:
+        wav_np = np.zeros(0, dtype=np.float32)
+
+    sr_raw = multimodal_output.get("sr")
+    if isinstance(sr_raw, list):
+        sr_raw = sr_raw[-1] if sr_raw else None
+    if isinstance(sr_raw, torch.Tensor):
+        sr = int(sr_raw.reshape(-1)[-1].item()) if sr_raw.numel() > 0 else 24000
+    elif sr_raw is None:
+        sr = 24000
+    else:
+        sr = int(sr_raw)
+    return wav_np, sr
+
+
+def _extract_text(prompt: Any, index: int = 0) -> str:
+    """Read the ground-truth text the aligner needs from the original prompt."""
+    if not isinstance(prompt, dict):
+        return ""
+    text = (prompt.get("additional_information") or {}).get("text")
+    if isinstance(text, list):
+        return str(text[index]) if index < len(text) else ""
+    return str(text) if text else ""
+
+
+def code2wav2aligner(
+    source_outputs: list[Any],
+    prompt: Any = None,
+    _requires_multimodal_data: bool = False,
+    _streaming_context: Any | None = None,
+) -> list[Any]:
+    """Build forced-aligner stage inputs from finished Code2Wav audio.
+
+    Emits a **text prompt** (``prompt`` string + ``multi_modal_data`` audio).
+    Tokenization and audio feature extraction are deferred to the aligner
+    stage's input preprocessor (run by the orchestrator), so this function
+    needs neither the aligner tokenizer nor its mm processor and stays
+    unit-testable on CPU.
+    """
+    from vllm_omni.inputs.data import OmniTextPrompt
+
+    aligner_inputs: list[OmniTextPrompt] = []
+    for i, src in enumerate(source_outputs):
+        if not getattr(src, "finished", False):
+            # Sentence-final alignment: only run once the audio is complete.
+            continue
+        mm = src.outputs[0].multimodal_output or {}
+        wav_np, sample_rate = _extract_waveform(mm)
+        if wav_np.size == 0:
+            logger.warning("code2wav2aligner: empty waveform for request; skipping alignment.")
+            continue
+
+        text = _extract_text(prompt, index=i)
+        language_list = extract_language_from_prompt(prompt, index=i)
+        language = language_list[0] if language_list else None
+        words = segment_words(text, language)
+        prompt_str = build_prompt(words)
+
+        additional_information: dict[str, Any] = {
+            "text": [text],
+            ALIGNER_WORDS_KEY: [words],
+            ALIGNER_DURATION_MS_KEY: [float(wav_np.size) / float(sample_rate) * 1000.0],
+        }
+        if language_list is not None:
+            additional_information["language"] = language_list
+
+        aligner_inputs.append(
+            OmniTextPrompt(
+                prompt=prompt_str,
+                multi_modal_data={"audio": (wav_np, sample_rate)},
+                additional_information=additional_information,
+            )
+        )
+    return aligner_inputs

vllm_omni/utils/forced_aligner.py

@@ -1,30 +1,29 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Shared forced aligner for streaming TTS word timestamps.
+"""Forced-aligner config + word-timestamp decoding for TTS.
 
-Hosts one lazy-loaded, in-process pooling ``vllm.LLM`` shared by the TTS
-frontend. ``llm.encode`` is sync/blocking, so :func:`align` runs it in
-``asyncio.to_thread``, serialized on a lock since the offline LLM is not
-thread-safe. See :func:`align` for the return-value contract.
+The aligner runs as a pooling stage appended to the pipeline by
+``--forced-aligner`` (see :func:`inject_forced_aligner_stage`). Qwen-specific
+word segmentation / prompt building / marker repair live in
+:mod:`vllm_omni.utils.qwen3_force_align_processor`.
 """
 
 from __future__ import annotations
 
-import asyncio
 import logging
-import os
-import threading
-from dataclasses import dataclass, field
+from dataclasses import dataclass, replace
 from pathlib import Path
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import yaml
 
 from vllm_omni.utils import qwen3_force_align_processor as _processor
 
-logger = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    from vllm_omni.config.stage_config import DeployConfig, PipelineConfig
 
+logger = logging.getLogger(__name__)
 
 _DEFAULT_CONFIG_PATH = Path(__file__).resolve().parents[1] / "deploy" / "qwen3_tts_forced_aligner.yaml"
 
@@ -42,7 +41,8 @@ class WordTimestamp:
 
 @dataclass(frozen=True, slots=True)
 class ForcedAlignerConfig:
-    """Plain-data config for constructing the aligner ``vllm.LLM`` (lazy-loaded)."""
+    """Plain-data config for the forced-aligner pipeline stage; fields are
+    consumed by :func:`inject_forced_aligner_stage`."""
 
     model: str
     runner: str | None = None
@@ -52,7 +52,6 @@ class ForcedAlignerConfig:
     dtype: str | None = None
     max_model_len: int | None = None
     trust_remote_code: bool | None = None
-    extra_llm_kwargs: dict[str, Any] = field(default_factory=dict)
 
 
 def _load_forced_aligner_yaml(path: str | Path) -> dict[str, Any]:
@@ -61,8 +60,8 @@ def _load_forced_aligner_yaml(path: str | Path) -> dict[str, Any]:
     return dict(raw.get("forced_aligner", raw))
 
 
-def build_forced_aligner_config(args: Any) -> ForcedAlignerConfig | None:
-    """Build a config from CLI args, or ``None`` when the feature is off.
+def build_forced_aligner_config(model: str | None, config_path: str | None = None) -> ForcedAlignerConfig | None:
+    """Build a config from the CLI values, or ``None`` when the feature is off.
 
     Precedence (lowest to highest): the packaged default YAML
     (:data:`_DEFAULT_CONFIG_PATH`, Qwen deploy defaults) -> a user YAML passed
@@ -70,254 +69,94 @@ def build_forced_aligner_config(args: Any) -> ForcedAlignerConfig | None:
     feature is off (returns ``None``) unless a model resolves from this chain.
     Per-field overrides such as ``gpu_memory_utilization`` live in the YAML.
     """
-    config_path = getattr(args, "forced_aligner_config", None)
     config_data: dict[str, Any] = {}
-    model = getattr(args, "forced_aligner", None)
     if config_path or model:
         config_data.update(_load_forced_aligner_yaml(_DEFAULT_CONFIG_PATH))
     if config_path:
         config_data.update(_load_forced_aligner_yaml(config_path))
     if model:
         config_data["model"] = str(model)
-    model = config_data.get("model")
-    if not model:
+    if not config_data.get("model"):
         return None
     allowed = set(ForcedAlignerConfig.__dataclass_fields__)
     return ForcedAlignerConfig(**{k: v for k, v in config_data.items() if k in allowed})
 
 
-class ForcedAlignerLoadError(RuntimeError):
-    """The aligner model/config could not be loaded (permanent until restart).
-
-    Distinguished from a per-request alignment failure so the streaming layer
-    can report a clear reason once rather than degrading every request to
-    ``timestamps: null``.
+def inject_forced_aligner_stage(
+    pipeline: PipelineConfig,
+    deploy: DeployConfig,
+    cli_overrides: dict[str, Any],
+) -> tuple[PipelineConfig, DeployConfig]:
+    """Append a forced-aligner pooling stage to the pipeline tail when
+    ``--forced-aligner`` or ``--forced-aligner-config`` resolves an aligner
+    model; no-op otherwise. The stage runs the aligner model with
+    ``runner="pooling"``, consuming the previous stage's audio and emitting a
+    terminal word-timestamps side-output.
     """
+    if not (cli_overrides.get("forced_aligner") or cli_overrides.get("forced_aligner_config")):
+        return pipeline, deploy
 
+    from vllm_omni.config.stage_config import (
+        StageDeployConfig,
+        StageExecutionType,
+        StagePipelineConfig,
+    )
+    from vllm_omni.model_executor.stage_input_processors.forced_aligner import (
+        POOLING_OUTPUT_DECODER_PATH,
+        TIMESTAMPS_MODALITY,
+    )
 
-# --- Singleton state ---
-# A single LLM serves the whole API server. ``_lock`` guards the lazy
-# constructor; ``_encode_lock`` serializes ``encode`` because the offline
-# ``vllm.LLM`` is a sync batch interface that is not safe to call from multiple
-# threads at once (concurrent WebSocket sessions otherwise race the engine).
-_lock = threading.Lock()
-_encode_lock = asyncio.Lock()
-_llm: Any = None
-_classify_num: int | None = None
-_timestamp_token_id: int | None = None
-_timestamp_segment_time_ms: float | None = None
-_loaded_config: ForcedAlignerConfig | None = None
+    fa = build_forced_aligner_config(
+        cli_overrides.get("forced_aligner"),
+        cli_overrides.get("forced_aligner_config"),
+    )
+    if fa is None:
+        return pipeline, deploy
 
+    _FA = "vllm_omni.model_executor.stage_input_processors.forced_aligner"
+    new_id = len(pipeline.stages)
+    aligner_ps = StagePipelineConfig(
+        stage_id=new_id,
+        model_stage="forced_aligner",
+        # Pooling stage = LLM_AR stage run with runner="pooling" (vLLM derives
+        # is_pooling_model from runner_type).
+        execution_type=StageExecutionType.LLM_AR,
+        input_sources=(new_id - 1,),
+        final_output=True,
+        final_output_type=TIMESTAMPS_MODALITY,
+        owns_tokenizer=True,
+        requires_multimodal_data=True,
+        model_arch=fa.architecture,
+        custom_process_input_func=f"{_FA}.code2wav2aligner",
+    )
+    extended = replace(pipeline, stages=pipeline.stages + (aligner_ps,))
 
-async def align(
-    *,
-    audio: bytes,
-    text: str,
-    sample_rate: int,
-    config: ForcedAlignerConfig,
-    language: str | None = None,
-) -> list[WordTimestamp] | None:
-    """Run one forced-alignment pass.
-
-    Args:
-        audio: Signed-int16 little-endian mono PCM bytes.
-        text: Ground-truth text whose words to align.
-        sample_rate: Sample rate of ``audio`` in Hz.
-        config: Aligner config (same instance for every call across the
-            server's lifetime; reload requires a server restart).
-        language: Optional language hint for word segmentation. ``None`` /
-            ``"auto"`` use the space + Chinese-mixed path; ``"japanese"`` /
-            ``"korean"`` (or their codes) need ``qwen_asr`` installed for
-            faithful tokenisation, else they degrade to whitespace splitting.
-
-    Returns:
-        List of :class:`WordTimestamp` on success (possibly empty for
-        silence / no aligned tokens), ``None`` on a per-request alignment
-        failure (decode error).
-
-    Raises:
-        ForcedAlignerLoadError: the model/config could not be loaded. This is
-            permanent until restart, so callers surface it once instead of
-            degrading every request to ``None``.
-    """
-    async with _encode_lock:
-        try:
-            await asyncio.to_thread(_ensure_loaded, config)
-        except Exception as exc:  # noqa: BLE001
-            logger.exception("Forced aligner failed to load")
-            raise ForcedAlignerLoadError(str(exc)) from exc
-        try:
-            return await asyncio.to_thread(_align_sync, audio, text, sample_rate, config, language)
-        except Exception:  # noqa: BLE001
-            logger.exception("Forced aligner failed for text=%r", text)
-            return None
-
-
-def _align_sync(
-    audio: bytes,
-    text: str,
-    sample_rate: int,
-    config: ForcedAlignerConfig,
-    language: str | None = None,
-) -> list[WordTimestamp]:
-    # Precondition: align() has already run _ensure_loaded(config) under
-    # _encode_lock and owns the load-failure -> ForcedAlignerLoadError mapping.
-    # Don't reload here, so a load error can't be swallowed as a per-request None.
-    audio_arr = _pcm_bytes_to_float32(audio)
-    if audio_arr.size == 0:
-        return []
-    audio_duration_ms = (audio_arr.size / sample_rate) * 1000.0
-
-    # Segment once and reuse for both the prompt and the decode: the word
-    # units MUST match between the two or the markers drift out of sync.
-    words = _processor.segment_words(text, language)
-    prompt = _processor.build_prompt(words)
-    request = {
-        "prompt": prompt,
-        "multi_modal_data": {"audio": (audio_arr, sample_rate)},
+    engine_extras: dict[str, Any] = {
+        "model": fa.model,
+        "runner": fa.runner or "pooling",
+        # Worker-side decoder hook: pooler logits -> word-timestamps tensor.
+        "pooling_output_decoder": POOLING_OUTPUT_DECODER_PATH,
     }
-
-    from vllm.pooling_params import PoolingParams
-
-    outputs = _llm.encode(  # type: ignore[union-attr]
-        [request],
-        pooling_params=PoolingParams(),
-        pooling_task=config.pooling_task or "token_classify",
-        use_tqdm=False,
+    if fa.architecture:
+        engine_extras["hf_overrides"] = {"architectures": [fa.architecture]}
+    if fa.trust_remote_code is not None:
+        engine_extras["trust_remote_code"] = fa.trust_remote_code
+    if fa.dtype is not None:
+        engine_extras["dtype"] = fa.dtype
+    aligner_ds = StageDeployConfig(
+        stage_id=new_id,
+        devices=cli_overrides.get("forced_aligner_device"),
+        gpu_memory_utilization=fa.gpu_memory_utilization,
+        max_model_len=fa.max_model_len,
+        engine_extras=engine_extras,
+        default_pooling_params={"task": fa.pooling_task or "token_classify"},
+        # Pooling is one-shot -> sync AR scheduler (the async one assumes
+        # token-by-token generation).
+        async_scheduling=False,
     )
-    if not outputs:
-        return []
-
-    result = outputs[0]
-    logits = result.outputs.data  # [n_token, classify_num]
-    prompt_token_ids = list(result.prompt_token_ids)
-    timestamp_positions = [i for i, tid in enumerate(prompt_token_ids) if tid == _timestamp_token_id]
-    if not timestamp_positions:
-        logger.warning(
-            "No <|timestamp|> tokens found in prompt for text=%r; aligner returned %d rows.",
-            text,
-            logits.shape[0] if hasattr(logits, "shape") else len(logits),
-        )
-        return []
-
-    return _decode_timestamps(
-        logits=logits,
-        words=words,
-        timestamp_positions=timestamp_positions,
-        classify_num=_classify_num,
-        timestamp_segment_time_ms=_timestamp_segment_time_ms,
-        audio_duration_ms=audio_duration_ms,
-    )
-
-
-def _ensure_loaded(config: ForcedAlignerConfig) -> None:
-    """Lazy-load the singleton ``vllm.LLM`` under lock; idempotent."""
-    global _llm, _classify_num, _timestamp_token_id, _timestamp_segment_time_ms, _loaded_config
-
-    if _llm is not None:
-        if _loaded_config is not None and _loaded_config.model != config.model:
-            # Multiple configs from different requests — refuse rather
-            # than swap models silently. A server restart is required.
-            raise RuntimeError(
-                f"Forced aligner already loaded with config={_loaded_config!r}; "
-                f"cannot serve a request that asks for {config!r}. "
-                "Restart the server to change the aligner config."
-            )
-        return
-
-    with _lock:
-        if _llm is not None:
-            return  # raced; another caller did the load
-
-        # Lazy import: vllm pulls torch + CUDA, which we want to avoid at
-        # module import time.
-        from vllm import LLM
-
-        logger.info(
-            "Loading forced aligner %s (gpu_memory_utilization=%s)",
-            config.model,
-            config.gpu_memory_utilization if config.gpu_memory_utilization is not None else "default",
-        )
-        llm_kwargs: dict[str, Any] = {"model": config.model}
-        if config.runner is not None:
-            llm_kwargs["runner"] = config.runner
-        if config.architecture is not None:
-            llm_kwargs["hf_overrides"] = {"architectures": [config.architecture]}
-        if config.gpu_memory_utilization is not None:
-            llm_kwargs["gpu_memory_utilization"] = config.gpu_memory_utilization
-        if config.trust_remote_code is not None:
-            llm_kwargs["trust_remote_code"] = config.trust_remote_code
-        if config.dtype is not None:
-            llm_kwargs["dtype"] = config.dtype
-        if config.max_model_len is not None:
-            llm_kwargs["max_model_len"] = config.max_model_len
-        llm_kwargs.update(config.extra_llm_kwargs)
-
-        # Force the aligner engine in-process: in multiprocessing mode the
-        # token_classify ``pooling_output`` fails to round-trip over the
-        # engine-core IPC on current vLLM. Hard-set (not setdefault) so a
-        # globally exported ``=1`` can't push it into a subprocess; save/restore
-        # to avoid leaking the override to anything built later in this process.
-        prev_mp = os.environ.get("VLLM_ENABLE_V1_MULTIPROCESSING")
-        os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
-        try:
-            llm = LLM(**llm_kwargs)
-        finally:
-            if prev_mp is None:
-                os.environ.pop("VLLM_ENABLE_V1_MULTIPROCESSING", None)
-            else:
-                os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = prev_mp
-
-        thinker_config = getattr(llm.llm_engine.model_config.hf_config, "thinker_config", None)
-        if thinker_config is None or not hasattr(thinker_config, "classify_num"):
-            raise RuntimeError(
-                "Loaded aligner has no thinker_config.classify_num; "
-                "expected a Qwen3ASRForcedAlignerForTokenClassification checkpoint."
-            )
-        timestamp_segment_time_ms = getattr(llm.llm_engine.model_config.hf_config, "timestamp_segment_time", None)
-        if timestamp_segment_time_ms is None:
-            raise RuntimeError(
-                "Loaded aligner has no timestamp_segment_time; expected a Qwen3ASR forced aligner checkpoint."
-            )
-
-        tokenizer = llm.get_tokenizer()
-        timestamp_token_id = _processor.resolve_timestamp_token_id(tokenizer)
-
-        # Publish in this order so a concurrent reader either sees
-        # _llm == None (will block on the lock) or sees a fully
-        # initialized aligner.
-        _classify_num = int(thinker_config.classify_num)
-        _timestamp_token_id = timestamp_token_id
-        _timestamp_segment_time_ms = float(timestamp_segment_time_ms)
-        _loaded_config = config
-        _llm = llm
-
-        logger.info(
-            "Forced aligner ready: timestamp_token_id=%d, classify_num=%d, timestamp_segment_time_ms=%.1f",
-            timestamp_token_id,
-            _classify_num,
-            _timestamp_segment_time_ms,
-        )
-
-
-# --- pure helpers (testable without a GPU / vllm) ---
-#
-# Word segmentation, prompt building, timestamp repair and marker-token
-# lookup are Qwen-specific and live in
-# :mod:`vllm_omni.utils.qwen3_force_align_processor`. This module owns only
-# the generic vLLM orchestration.
-
-
-def _pcm_bytes_to_float32(audio: bytes) -> np.ndarray:
-    """Decode signed-int16 mono PCM bytes into a [-1, 1] float32 array."""
-    if not audio:
-        return np.zeros(0, dtype=np.float32)
-    if len(audio) % 2 != 0:
-        # Drop a trailing odd byte rather than raise; keeps streaming
-        # robust against off-by-one chunk boundaries.
-        audio = audio[:-1]
-    pcm = np.frombuffer(audio, dtype=np.int16)
-    return (pcm.astype(np.float32) / 32768.0).copy()
+    deploy = replace(deploy, stages=list(deploy.stages) + [aligner_ds])
+    logger.info("[stage_init] Injected forced-aligner stage %d (model=%s)", new_id, fa.model)
+    return extended, deploy
 
 
 def _decode_timestamps(
@@ -382,18 +221,40 @@ def _decode_timestamps(
     return out
 
 
-# Test hooks ---------------------------------------------------------------
-# Tests need a way to reset module state without restarting Python. Not
-# part of the public API; do not call in production code.
+def extract_word_timestamps(res: Any, text: str, language: str | None = None) -> list[dict] | None:
+    """Read a forced-aligner stage's engine output into per-word timestamps.
 
+    Returns ``[{word, start_ms, end_ms}, ...]`` from the stage's pooling output
+    (``res.outputs.data`` is an int32 ``[n_words, 2]`` tensor). Words come from
+    re-segmenting ``text``; any mismatch or non-aligner input returns ``None``.
+    """
+    if not text:
+        return None
+    try:
+        outs = getattr(res, "outputs", None)
+        if outs is None:
+            return None
+        # The aligner stage's res.outputs is a single PoolingOutput (not a list).
+        # Its .data is a [n_words, 2] int32 tensor of [start_ms, end_ms] per word.
+        out0 = outs[0] if isinstance(outs, (list, tuple)) else outs
+        ms = getattr(out0, "data", None)
+        if ms is None or not hasattr(ms, "tolist"):
+            return None
+        ms_list = ms.tolist()
+        if not ms_list or not isinstance(ms_list[0], (list, tuple)):
+            return None
 
-def _reset_for_tests() -> None:
-    """Drop the cached aligner state so the next call reloads."""
-    global _llm, _classify_num, _timestamp_token_id, _timestamp_segment_time_ms, _loaded_config
-    with _lock:
-        _llm = None
-        _classify_num = None
-        _timestamp_token_id = None
-        _timestamp_segment_time_ms = None
-        _loaded_config = None
-    _processor._reset_for_tests()
+        words = _processor.segment_words(text, language)
+        if len(words) != len(ms_list):
+            # Segmentation drifted from the engine side; pairing positionally would mislabel every word.
+            logger.warning(
+                "word/timestamp count mismatch (%d words vs %d intervals); returning no timestamps",
+                len(words),
+                len(ms_list),
+            )
+            return None
+
+        return [{"word": w, "start_ms": int(p[0]), "end_ms": int(p[1])} for w, p in zip(words, ms_list)]
+    except Exception:
+        logger.debug("failed to extract word timestamps from aligner output", exc_info=True)
+        return None

vllm_omni/utils/qwen3_force_align_processor.py

@@ -12,15 +12,15 @@
 """Qwen3 forced-aligner text/timestamp processor.
 
 This is the model-specific half of upstream's ``Qwen3ForceAlignProcessor``:
-it turns text into the aligner's word units and prompt, repairs the predicted
-timestamp bins, and resolves the marker token id. The generic "run a vLLM
-pooling model" half lives in :mod:`vllm_omni.utils.forced_aligner`, which
-calls into this module.
+it turns text into the aligner's word units and prompt, and repairs the
+predicted timestamp bins. It feeds the forced-aligner pipeline stage (input
+glue in :mod:`vllm_omni.model_executor.stage_input_processors.forced_aligner`,
+timestamp decoding in :mod:`vllm_omni.utils.forced_aligner`).
 
 Keeping the Qwen-specific pieces here marks the seam for the model-agnostic
 aligner the issue asks for: a different aligner family would supply its own
 processor exposing the same small surface — :func:`segment_words`,
-:func:`build_prompt`, :func:`fix_timestamp`, :func:`resolve_timestamp_token_id`.
+:func:`build_prompt`, :func:`fix_timestamp`.
 
 Word segmentation prefers qwen_asr's official ``Qwen3ForceAlignProcessor``
 when installed (full multilingual fidelity, incl. Japanese/Korean) and
@@ -260,26 +260,3 @@ def fix_timestamp(values: list[int]) -> list[int]:
         i = j
 
     return [int(r) for r in result]
-
-
-def resolve_timestamp_token_id(tokenizer: Any, timestamp_token: str = TIMESTAMP_TOKEN) -> int:
-    """Look up the integer id of the timestamp special token."""
-    convert = getattr(tokenizer, "convert_tokens_to_ids", None)
-    if not callable(convert):
-        raise RuntimeError("Aligner tokenizer has no convert_tokens_to_ids method.")
-    tid = convert(timestamp_token)
-    if isinstance(tid, list):
-        tid = tid[0] if tid else None
-    if tid is None or (isinstance(tid, int) and tid < 0):
-        raise RuntimeError(
-            f"Aligner tokenizer does not recognise {timestamp_token!r} (got id={tid}). "
-            "Check the model card; the marker token may use a different name."
-        )
-    return int(tid)
-
-
-def _reset_for_tests() -> None:
-    """Drop the cached official-processor probe so the next call re-checks."""
-    global _official_processor, _official_processor_unavailable
-    _official_processor = None
-    _official_processor_unavailable = False

vllm_omni/worker/gpu_ar_model_runner.py

@@ -1360,7 +1360,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 hidden_states = model_output
                 aux_hidden_states = None
 
-            hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output)
+            # A pooling stage skips mm extraction + the omni prefix cache (both assume an AR
+            # [n_tok, hidden] output) and exits through the shared _pool() branch below.
+            if self.is_pooling_model:
+                multimodal_outputs = None
+            else:
+                hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output)
             hidden_states_cpu = None
 
             # Async-write pipeline (replaces the per-step blocking
@@ -1369,7 +1374,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             # the actual CPU scatter into ``hidden_states_cache`` /
             # ``mm_outputs_cache`` happens in ``drain_ready_async_writes``
             # at the top of subsequent execute_model() calls.
-            if self.omni_prefix_cache is not None and get_pp_group().is_last_rank:
+            if not self.is_pooling_model and self.omni_prefix_cache is not None and get_pp_group().is_last_rank:
                 hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None
                 # Some models (e.g. qwen3-tts-talker) opt out of full-hidden-state
                 # prefix caching but the downstream pooler payload path still