[Bugfix] Report sample rate and channels for streamed speech PCM (#8567)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

docs/serving/speech_api.md

@@ -183,6 +183,16 @@ emit these headers.
 **Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or
 WAV) as they are decoded.
 
+`pcm` is 16-bit signed little-endian, interleaved, at the model's native sample
+rate and channel count (24 kHz mono for most models; MOSS-TTS-Local v1.5 emits
+48 kHz stereo). Raw streams and non-streaming responses state the format in
+response headers:
+
+| Header | Description |
+| --- | --- |
+| `x-audio-sample-rate` | Sample rate of the returned audio in Hz (after any `sample_rate` resampling). |
+| `x-audio-channels` | Channel count of the returned audio. |
+
 Both streaming modes (`stream_format="audio"` and `"sse"`) require
 `response_format="pcm"` or `"wav"`, and `speed` must be `1.0` (or omitted).
 
@@ -193,9 +203,13 @@ Each event has an `event:` line and a JSON `data:` line:
 - `speech.audio.delta` — a base64 audio chunk:
 
     ```json
-    { "type": "speech.audio.delta", "audio": "<base64>", "response_format": "pcm" }
+    { "type": "speech.audio.delta", "audio": "<base64>", "response_format": "pcm",
+      "sample_rate": 24000, "channels": 1 }
     ```
 
+    `sample_rate` and `channels` extend the OpenAI event and describe the
+    decoded chunk.
+
 - `speech.audio.done` — terminal event, carrying token `usage`:
 
     ```json
@@ -340,7 +354,7 @@ Server -> Client:
 
 | Message | Description |
 | --------- | ------------- |
-| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
+| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000, "channels": 1}` | Audio generation starting for the buffered input. With `stream_audio=true` it is sent with the first audio chunk, and for `pcm` it carries that audio's `sample_rate` and `channels` (omitted when no audio was produced). |
 | Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
 | `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
 | `{"type": "session.done", "utterance_index": 0, "total_sentences": N}` | Flushed utterance complete |

tests/benchmarks/patch/test_patch.py

@@ -35,6 +35,7 @@ from vllm_omni.benchmarks.patch.patch import (
     _iter_image_reference_inputs,
     _iter_video_reference_inputs,
     _omni_request_timeout_s,
+    async_request_openai_audio_speech,
     async_request_openai_chat_omni_completions,
     async_request_openai_image_edits_omni,
     async_request_openai_image_generations_omni,
@@ -474,6 +475,48 @@ async def test_streaming_error_chunk_marks_request_failed(mocker: MockerFixture)
     assert output.error == "EngineCore encountered an issue"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("headers", "env", "expected_duration_s"),
+    [
+        # Server-stated format wins over the env fallback.
+        ({"X-Audio-Sample-Rate": "48000", "X-Audio-Channels": "2"}, {}, 1.0),
+        ({"X-Audio-Sample-Rate": "48000", "X-Audio-Channels": "2"}, {"rate": "24000", "channels": "1"}, 1.0),
+        # Older servers send no headers: the env (default 24 kHz mono) applies.
+        ({}, {}, 4.0),
+        ({}, {"rate": "48000", "channels": "2"}, 1.0),
+    ],
+)
+async def test_audio_speech_stream_uses_server_pcm_format(
+    mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, headers, env, expected_duration_s
+):
+    from vllm_omni.metrics import definitions as defs
+
+    for key, var in (("rate", defs.AUDIO_SAMPLE_RATE_ENV), ("channels", defs.AUDIO_CHANNELS_ENV)):
+        if key in env:
+            monkeypatch.setenv(var, env[key])
+        else:
+            monkeypatch.delenv(var, raising=False)
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/audio/speech",
+        prompt_len=10,
+        output_len=20,
+    )
+    # 1 s of 48 kHz stereo s16le = 192000 bytes.
+    mock_response = MockResponse(200, [b"\x00" * 96000, b"\x00" * 96000])
+    mock_response.headers = headers
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)
+
+    output = await async_request_openai_audio_speech(request_input, mock_session)
+
+    assert output.success is True
+    assert output.audio_duration == pytest.approx(expected_duration_s)
+
+
 @pytest.mark.asyncio
 async def test_positive_client_text_timings_take_precedence_over_stage0(mocker: MockerFixture):
     request_input = RequestFuncInput(

tests/entrypoints/openai_api/test_serving_speech.py

@@ -3468,17 +3468,25 @@ class TestStreamingProtocolValidation:
         assert req.is_raw_audio_stream() is False
 
 
-class TestStreamingResponse:
-    """Integration tests for the streaming audio response path."""
+def _build_streaming_speech_app(
+    mocker: MockerFixture,
+    chunk: torch.Tensor | None = None,
+    sample_rate: int | None = None,
+) -> FastAPI:
+    """Speech app whose mock engine yields one intermediate then one final chunk.
 
-    @pytest.fixture
-    def streaming_app(self, mocker: MockerFixture):
-        """Test app whose mock engine yields one intermediate chunk then a final chunk."""
+    ``sample_rate`` is reported as the model's ``sr`` when given; omitted, the
+    serving layer's 24 kHz default applies (as for most TTS models).
+    """
     mock_engine_client = mocker.MagicMock()
     mock_engine_client.errored = False
+    if chunk is None:
+        chunk = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000))
 
     def _make_output(finished: bool) -> OmniRequestOutput:
-            chunk = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000))
+        mm_output: dict = {"audio": chunk}
+        if sample_rate is not None:
+            mm_output["sr"] = torch.tensor(sample_rate)
 
         class MockCompletionOutput:
             def __init__(self, index: int = 0):
@@ -3490,10 +3498,10 @@ class TestStreamingResponse:
                 self.logprobs = None
 
         class MockRequestOutput:
-                def __init__(self, audio_tensor: torch.Tensor):
+            def __init__(self, multimodal_output: dict):
                 self.request_id = "speech-stream-test"
                 self.outputs = [MockCompletionOutput(index=0)]
-                    self.multimodal_output = {"audio": audio_tensor}
+                self.multimodal_output = multimodal_output
                 self.finished = finished
                 self.prompt_token_ids = None
                 self.encoder_prompt_token_ids = None
@@ -3502,7 +3510,7 @@ class TestStreamingResponse:
                 self.kv_transfer_params = None
 
         return OmniRequestOutput.from_stage_output(
-                MockRequestOutput(audio_tensor=chunk),
+            MockRequestOutput(mm_output),
             stage_id=0,
             final_output_type="audio",
             finished=finished,
@@ -3551,6 +3559,15 @@ class TestStreamingResponse:
     app.state.speech_server = speech_server
     return app
 
+
+class TestStreamingResponse:
+    """Integration tests for the streaming audio response path."""
+
+    @pytest.fixture
+    def streaming_app(self, mocker: MockerFixture):
+        """Test app whose mock engine yields one intermediate chunk then a final chunk."""
+        return _build_streaming_speech_app(mocker)
+
     @staticmethod
     def _assert_sse_audio_response(response, response_format: str = "pcm"):
         assert response.status_code == 200
@@ -3603,6 +3620,7 @@ class TestStreamingResponse:
             tts_params=None,
             collect=None,
             target_sample_rate=None,
+            audio_format=None,
         ):
             captured["tts_params"] = tts_params
             captured["target_sample_rate"] = target_sample_rate
@@ -3790,6 +3808,146 @@ class TestStreamingResponse:
         assert "audio/wav" in response.headers["content-type"]
 
 
+class TestSpeechAudioFormatMetadata:
+    """pcm carries no header, so responses must state the model-native format."""
+
+    # 24 kHz mono (most TTS models) and 48 kHz [channels, samples] stereo
+    # (e.g. MOSS-TTS-Local v1.5); the serving layer must not care which.
+    _MONO = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 2400))
+    _STEREO = torch.stack([_MONO.repeat(2) * 0.5, -_MONO.repeat(2) * 0.25])
+
+    _CASES = [
+        pytest.param(None, None, 24000, 1, id="default-24k-mono"),
+        pytest.param(_MONO, 24000, 24000, 1, id="explicit-24k-mono"),
+        pytest.param(_STEREO, 48000, 48000, 2, id="48k-stereo"),
+    ]
+
+    @staticmethod
+    def _expected_pcm(chunk: torch.Tensor | None) -> bytes:
+        """Two engine outputs, each the chunk as interleaved s16le."""
+        # ``None`` is the builder's default chunk.
+        audio = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000)) if chunk is None else chunk
+        samples = audio.numpy().T if audio.ndim == 2 else audio.numpy()
+        # libsndfile-compatible PCM_16 conversion used by ``create_audio``.
+        one = np.clip(np.floor(samples.astype(np.float64) * 32768.0), -32768, 32767).astype("<i2").tobytes()
+        return one * 2
+
+    @pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES)
+    def test_raw_pcm_stream_headers(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
+        app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
+        response = TestClient(app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
+        )
+
+        assert response.status_code == 200
+        assert response.headers["content-type"].startswith("audio/pcm")
+        assert response.headers["x-audio-sample-rate"] == str(rate)
+        assert response.headers["x-audio-channels"] == str(channels)
+        exposed = {h.strip().lower() for h in response.headers["access-control-expose-headers"].split(",")}
+        assert {"x-audio-sample-rate", "x-audio-channels"} <= exposed
+        # The body stays the model-native interleaved s16le stream, byte for byte.
+        assert response.content == self._expected_pcm(chunk)
+
+    # WAV streaming requires the model to report ``sr`` on its first chunk.
+    @pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES[1:])
+    def test_raw_wav_stream_header_matches_metadata(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
+        app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
+        response = TestClient(app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "wav"},
+        )
+
+        assert response.status_code == 200
+        assert response.headers["x-audio-sample-rate"] == str(rate)
+        assert response.headers["x-audio-channels"] == str(channels)
+        num_channels, sample_rate = struct.unpack("<HI", response.content[22:28])
+        assert (sample_rate, num_channels) == (rate, channels)
+
+    @pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES)
+    def test_sse_deltas_carry_format(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
+        app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
+        response = TestClient(app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream": True, "response_format": "pcm"},
+        )
+
+        deltas = [
+            json.loads(line.removeprefix("data: "))
+            for line in response.text.splitlines()
+            if line.startswith("data: ") and '"speech.audio.delta"' in line
+        ]
+        assert deltas
+        assert all((d["sample_rate"], d["channels"]) == (rate, channels) for d in deltas)
+
+    @pytest.mark.parametrize(("chunk", "model_sr", "rate", "channels"), _CASES)
+    def test_non_streaming_pcm_headers(self, mocker: MockerFixture, chunk, model_sr, rate, channels):
+        app = _build_streaming_speech_app(mocker, chunk=chunk, sample_rate=model_sr)
+        response = TestClient(app).post("/v1/audio/speech", json={"input": "Hello", "response_format": "pcm"})
+
+        assert response.status_code == 200
+        assert response.headers["x-audio-sample-rate"] == str(rate)
+        assert response.headers["x-audio-channels"] == str(channels)
+        assert len(response.content) % (2 * channels) == 0
+
+    def test_requested_sample_rate_is_reported(self, mocker: MockerFixture):
+        """The headers describe the bytes sent, i.e. after request.sample_rate resampling."""
+        app = _build_streaming_speech_app(mocker)
+        speech_server = app.state.speech_server
+        mocker.patch.object(speech_server, "_validate_speech_sample_rate", return_value=None)
+        response = TestClient(app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "pcm", "sample_rate": 8000},
+        )
+
+        assert response.status_code == 200
+        assert response.headers["x-audio-sample-rate"] == "8000"
+        assert response.headers["x-audio-channels"] == "1"
+        # Two 1 s chunks at 8 kHz mono s16le, within resampler edge tolerance.
+        assert abs(len(response.content) - 2 * 8000 * 2) <= 2 * 64
+
+    def test_raw_stream_error_before_audio_is_an_http_error(self, mocker: MockerFixture):
+        """Priming the first chunk lets a pre-audio failure return a real status code."""
+        app = _build_streaming_speech_app(mocker)
+        speech_server = app.state.speech_server
+
+        async def failing_generate(*args, **kwargs):
+            raise RuntimeError("boom before audio")
+            yield  # pragma: no cover - generator marker, unreachable
+
+        speech_server.engine_client.generate = mocker.MagicMock(side_effect=failing_generate)
+        response = TestClient(app).post(
+            "/v1/audio/speech",
+            json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"},
+        )
+
+        assert response.status_code >= 400
+        assert "audio/pcm" not in response.headers.get("content-type", "")
+        assert "x-audio-sample-rate" not in response.headers
+
+    def test_primed_stream_closes_inner_generator(self):
+        """Closing the response body before iteration still closes the engine stream."""
+        from vllm_omni.entrypoints.openai.serving_speech import _PrimedAudioStream
+
+        closed: list[bool] = []
+
+        async def chunks():
+            try:
+                yield b"\x01\x00"
+                yield b"\x02\x00"
+            finally:
+                closed.append(True)
+
+        async def run() -> list[bytes]:
+            body = await _PrimedAudioStream.start(chunks())
+            await body.aclose()
+            assert closed == [True]
+            body = await _PrimedAudioStream.start(chunks())
+            return [c async for c in body]
+
+        assert asyncio.run(run()) == [b"\x01\x00", b"\x02\x00"]
+
+
 class TestSpeechBatchAPI:
     """Tests for the /v1/audio/speech/batch endpoint."""
 

tests/entrypoints/openai_api/test_serving_speech_stream.py

@@ -33,6 +33,12 @@ def _fake_aligner_res(pairs, words):
     )
 
 
+def _fill_audio_format(audio_format: dict | None, sample_rate: int = 24000, channels: int = 1) -> None:
+    """Mirror ``_generate_audio_chunks``: report the format before the first chunk."""
+    if audio_format is not None and not audio_format:
+        audio_format.update(sample_rate=sample_rate, channels=channels)
+
+
 def _build_test_app(
     speech_service=None,
     *,
@@ -57,8 +63,10 @@ def _build_test_app(
             tts_params=None,
             collect=None,
             cumulative_audio=False,
+            audio_format=None,
         ):
             for chunk in (b"\x01\x02", b"\x03\x04\x05"):
+                _fill_audio_format(audio_format)
                 yield (chunk, 24000) if include_sample_rate else chunk
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
@@ -282,6 +290,7 @@ class TestStreamingSpeechWebSocket:
             request_arrival_ts=None,
             include_sample_rate=False,
             tts_params=None,
+            audio_format=None,
         ):
             assert request_start_s is not None
             assert request_arrival_ts is not None
@@ -289,6 +298,7 @@ class TestStreamingSpeechWebSocket:
             captured_timing["chunk_arrival"] = request_arrival_ts
             captured_tts_params.append(tts_params)
             for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
+                _fill_audio_format(audio_format)
                 yield (chunk, 24000) if include_sample_rate else chunk
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
@@ -312,6 +322,7 @@ class TestStreamingSpeechWebSocket:
                 assert start["type"] == "audio.start"
                 assert start["format"] == "pcm"
                 assert start["sample_rate"] == 24000
+                assert start["channels"] == 1
 
                 assert ws.receive_bytes() == b"\x01\x02"
                 assert ws.receive_bytes() == b"\x03\x04\x05"
@@ -337,6 +348,91 @@ class TestStreamingSpeechWebSocket:
         assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}]
         assert speech_service._generate_audio_bytes.await_count == 0
 
+    @staticmethod
+    def _stream_service(mocker: MockerFixture, generate_pcm_chunks, *, forced_aligner: bool = False):
+        speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
+        speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
+        speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
+        speech_service._generate_pcm_chunks = generate_pcm_chunks
+        speech_service.engine_client = mocker.MagicMock()
+        speech_service.engine_client.abort = mocker.AsyncMock()
+        speech_service.forced_aligner_enabled = forced_aligner
+        return speech_service
+
+    @staticmethod
+    def _send_pcm_session(ws, **config) -> None:
+        ws.send_json(
+            {"type": "session.config", "voice": "Vivian", "stream_audio": True, "response_format": "pcm", **config}
+        )
+        ws.send_json({"type": "input.text", "text": "Hello world. "})
+        ws.send_json({"type": "input.done"})
+
+    def test_streaming_pcm_start_reports_model_native_format(self, mocker: MockerFixture):
+        """audio.start states the real rate/channels, not a nominal 24 kHz."""
+
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, audio_format=None, **_kwargs):
+            _fill_audio_format(audio_format, sample_rate=48000, channels=2)
+            yield b"\x01\x02\x03\x04"
+
+        app, _ = _build_test_app(self._stream_service(mocker, mock_generate_pcm_chunks))
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                self._send_pcm_session(ws)
+                start = ws.receive_json()
+                assert start["type"] == "audio.start"
+                assert start["sample_rate"] == 48000
+                assert start["channels"] == 2
+                assert ws.receive_bytes() == b"\x01\x02\x03\x04"
+                assert ws.receive_json()["type"] == "audio.done"
+
+    def test_streaming_pcm_error_before_audio_keeps_start_done_pairing(self, mocker: MockerFixture):
+        """No audio means no known format: audio.start omits it rather than guessing."""
+
+        async def mock_generate_pcm_chunks(_generator, _request_id, *, audio_format=None, **_kwargs):
+            raise RuntimeError("early boom")
+            yield  # pragma: no cover - generator marker, unreachable
+
+        app, _ = _build_test_app(self._stream_service(mocker, mock_generate_pcm_chunks))
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                self._send_pcm_session(ws)
+                start = ws.receive_json()
+                assert start["type"] == "audio.start"
+                assert "sample_rate" not in start
+                assert "channels" not in start
+                assert ws.receive_json() == {
+                    "type": "error",
+                    "message": "Generation failed for utterance 0, sentence 0: early boom",
+                }
+                done = ws.receive_json()
+                assert done["type"] == "audio.done"
+                assert done["error"] is True
+
+    def test_word_timestamps_chunk_ms_account_for_channels(self, mocker: MockerFixture):
+        """Stereo s16le is 4 bytes per frame; chunk offsets must not double."""
+
+        async def mock_generate_pcm_chunks(
+            _generator, _request_id, *, include_sample_rate=False, collect=None, audio_format=None, **_kwargs
+        ):
+            _fill_audio_format(audio_format, sample_rate=1000, channels=2)
+            yield (b"\x01" * 4000, 1000) if include_sample_rate else b"\x01" * 4000
+
+        speech_service = self._stream_service(mocker, mock_generate_pcm_chunks, forced_aligner=True)
+        app, _ = _build_test_app(speech_service)
+
+        with TestClient(app) as client:
+            with client.websocket_connect("/v1/audio/speech/stream") as ws:
+                self._send_pcm_session(ws, word_timestamps=True)
+                start = ws.receive_json()
+                assert (start["sample_rate"], start["channels"]) == (1000, 2)
+                chunk = ws.receive_json()
+                assert (chunk["chunk_start_ms"], chunk["chunk_end_ms"]) == (0, 1000)
+                assert (chunk["sample_rate"], chunk["channels"]) == (1000, 2)
+                final = ws.receive_json()
+                assert final["chunk_end_ms"] == 1000
+
     def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture):
         app, _ = _build_test_app(mocker=mocker)
 
@@ -388,11 +484,13 @@ class TestStreamingSpeechWebSocket:
             tts_params=None,
             collect=None,
             cumulative_audio=False,
+            audio_format=None,
         ):
             assert request_start_s is not None
             assert request_arrival_ts is not None
             assert cumulative_audio is True
             for chunk in (first_chunk, second_chunk):
+                _fill_audio_format(audio_format, sample_rate=1000)
                 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"])
@@ -484,8 +582,10 @@ class TestStreamingSpeechWebSocket:
             tts_params=None,
             collect=None,
             cumulative_audio=False,
+            audio_format=None,
         ):
             chunk = b"\x01" * 1000
+            _fill_audio_format(audio_format, sample_rate=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"])
@@ -673,7 +773,9 @@ class TestStreamingSpeechWebSocket:
             request_arrival_ts=None,
             include_sample_rate=False,
             tts_params=None,
+            audio_format=None,
         ):
+            _fill_audio_format(audio_format)
             yield b"\x01\x02"
             raise RuntimeError("stream boom")
 
@@ -777,7 +879,9 @@ class TestStreamingSpeechWebSocket:
             request_arrival_ts=None,
             include_sample_rate=False,
             tts_params=None,
+            audio_format=None,
         ):
+            _fill_audio_format(audio_format)
             yield b"\x01\x02"
 
         speech_service._generate_pcm_chunks = mock_generate_pcm_chunks

vllm_omni/benchmarks/data_modules/seed_tts_eval.py

@@ -42,8 +42,10 @@ Env: ``SEED_TTS_EVAL_DEVICE`` (e.g. ``cuda:0``, ``cpu``); ``SEED_TTS_HF_WHISPER_
 defaults to ``openai/whisper-large-v3`` (override for debugging only). Set
 ``SEED_TTS_WER_SAVE_AUDIO_DIR`` to save the captured 24 kHz mono WAV used by
 WER evaluation for each synthesized utterance.
-Streaming PCM is decoded using ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` /
-``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` (default: 24 kHz mono).
+Streaming PCM is decoded using the server's ``X-Audio-Sample-Rate`` /
+``X-Audio-Channels`` response headers; ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` /
+``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` (default: 24 kHz mono) apply when a server
+does not send them.
 """
 
 from __future__ import annotations

vllm_omni/benchmarks/patch/patch.py

@@ -207,6 +207,32 @@ def _pcm_s16le_to_seed_tts_wer_bytes(
     return (pcm_f32 * 32767).astype(np.int16).tobytes()
 
 
+def _stream_pcm_format_from_headers(
+    headers: Mapping[str, str],
+    sample_rate: int,
+    channels: int,
+) -> tuple[int, int]:
+    """Prefer the server's ``X-Audio-Sample-Rate``/``X-Audio-Channels`` headers.
+
+    Raw PCM has no self-describing header; the server states the model-native
+    format. ``sample_rate``/``channels`` (env-configured) remain the fallback.
+    """
+    resolved = []
+    for name, fallback in (("X-Audio-Sample-Rate", sample_rate), ("X-Audio-Channels", channels)):
+        raw = headers.get(name)
+        value = fallback
+        if raw:
+            try:
+                value = int(raw)
+            except ValueError:
+                logger.warning("Ignoring invalid %s header %r; using %d", name, raw, fallback)
+            else:
+                if value != fallback:
+                    logger.debug("Streamed PCM %s=%d from server header (configured %d)", name, value, fallback)
+        resolved.append(max(value, 1))
+    return resolved[0], resolved[1]
+
+
 get_samples_old = datasets.get_samples
 
 _DEFAULT_DAILY_OMNI_REPO = "liarliar/Daily-Omni"
@@ -2610,6 +2636,7 @@ async def async_request_openai_audio_speech(
     output.prompt_len = request_func_input.prompt_len
 
     # PCM format: 16-bit signed; sample_rate/channels are model-dependent.
+    # The env vars are a fallback for servers that do not send format headers.
     sample_rate, channels = defs.stream_pcm_format_from_env()
     sample_width = defs.DEFAULT_AUDIO_SAMPLE_WIDTH
 
@@ -2623,6 +2650,7 @@ async def async_request_openai_audio_speech(
     try:
         async with session.post(url=api_url, json=payload, headers=headers) as response:
             if response.status == 200:
+                sample_rate, channels = _stream_pcm_format_from_headers(response.headers, sample_rate, channels)
                 async for chunk in response.content.iter_any():
                     if not chunk:
                         continue

vllm_omni/entrypoints/openai/serving_speech.py

@@ -79,6 +79,10 @@ _SPEECH_USAGE_OUTPUT_TOKENS_HEADER = "X-VLLM-OMNI-OUTPUT-TOKENS"
 _SPEECH_USAGE_TOTAL_TOKENS_HEADER = "X-VLLM-OMNI-TOTAL-TOKENS"
 _SPEECH_USAGE_INPUT_TEXT_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-TEXT-TOKENS"
 _SPEECH_USAGE_INPUT_AUDIO_TOKENS_HEADER = "X-VLLM-OMNI-INPUT-AUDIO-TOKENS"
+# Raw ``pcm`` bodies carry no format header and the rate/channel count is
+# model-native, so responses state the s16le layout explicitly.
+SPEECH_AUDIO_SAMPLE_RATE_HEADER = "X-Audio-Sample-Rate"
+SPEECH_AUDIO_CHANNELS_HEADER = "X-Audio-Channels"
 
 
 def _stage_speech_metadata(stage: Any) -> tuple[str | None, str | None, str | None]:
@@ -176,6 +180,64 @@ def _infer_audio_num_channels(audio: np.ndarray) -> int:
     return 1
 
 
+def _audio_format_headers(audio_format: dict[str, int] | None) -> dict[str, str]:
+    """Response headers describing the emitted audio's rate and channel count."""
+    if not audio_format:
+        return {}
+    return {
+        SPEECH_AUDIO_SAMPLE_RATE_HEADER: str(int(audio_format["sample_rate"])),
+        SPEECH_AUDIO_CHANNELS_HEADER: str(int(audio_format["channels"])),
+        # Without this, CORS hides both headers from cross-origin browser clients.
+        "Access-Control-Expose-Headers": f"{SPEECH_AUDIO_SAMPLE_RATE_HEADER}, {SPEECH_AUDIO_CHANNELS_HEADER}",
+    }
+
+
+class _PrimedAudioStream:
+    """Async iterator over an audio stream whose first chunk was already pulled.
+
+    Raw audio streams pull the first chunk before the HTTP response starts so
+    its headers can state the real sample rate and channel count.
+    """
+
+    def __init__(self, chunks: Any, first: Any, has_first: bool) -> None:
+        self._chunks = chunks
+        self._first = first
+        self._has_first = has_first
+
+    @classmethod
+    async def start(cls, chunks: Any) -> "_PrimedAudioStream":
+        try:
+            first = await anext(chunks)
+        except StopAsyncIteration:
+            return cls(chunks, None, False)
+        except BaseException:
+            with anyio.CancelScope(shield=True):
+                await chunks.aclose()
+            raise
+        return cls(chunks, first, True)
+
+    def __aiter__(self) -> "_PrimedAudioStream":
+        return self
+
+    async def __anext__(self) -> Any:
+        if self._has_first:
+            first, self._first, self._has_first = self._first, None, False
+            return first
+        return await anext(self._chunks)
+
+    async def aclose(self) -> None:
+        self._first, self._has_first = None, False
+        await self._chunks.aclose()
+
+
+def _encoded_audio_format(audio_response: Any, sample_rate: int, audio: np.ndarray) -> dict[str, int]:
+    """Rate and channel count of an encoded chunk, preferring encoder metadata."""
+    metadata = getattr(audio_response, "audio_metadata", None)
+    if metadata is not None:
+        return {"sample_rate": int(metadata.sample_rate_hz), "channels": int(metadata.channels)}
+    return {"sample_rate": int(sample_rate), "channels": _infer_audio_num_channels(np.asarray(audio))}
+
+
 def _sanitize_filename(filename: str) -> str:
     """Sanitize filename to prevent path traversal attacks.
 
@@ -1568,6 +1630,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         collect: dict | None = None,
         target_sample_rate: int | None = None,
         cumulative_audio: bool = False,
+        audio_format: dict[str, int] | None = None,
     ):
         """Generate audio chunks for streaming response.
 
@@ -1581,6 +1644,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             generator: Async generator from the engine
             request_id: Request identifier for logging
             response_format: Audio format (pcm or wav)
+            audio_format: Optional out-parameter. Before the first audio
+                chunk is yielded it receives that chunk's ``sample_rate`` and
+                ``channels``, so callers can describe a raw PCM stream.
 
         Yields:
             Raw audio bytes for each chunk (with WAV header for first chunk if wav format)
@@ -1722,8 +1788,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         speed=1.0,
                         base64_encode=False,
                     )
-                    audio_bytes = cast(bytes, self.create_audio(audio_obj).audio_data)
+                    audio_response = self.create_audio(audio_obj)
+                    audio_bytes = cast(bytes, audio_response.audio_data)
                     record_audio_chunk(audio_bytes, chunk_np, res)
+                    if audio_format is not None and not audio_format and audio_bytes:
+                        audio_format.update(_encoded_audio_format(audio_response, output_sample_rate, chunk_np))
                     if wav_header is not None:
                         yield wav_header
                     if include_sample_rate:
@@ -1750,8 +1819,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         speed=1.0,
                         base64_encode=False,
                     )
-                    audio_bytes = cast(bytes, self.create_audio(audio_obj).audio_data)
+                    audio_response = self.create_audio(audio_obj)
+                    audio_bytes = cast(bytes, audio_response.audio_data)
                     record_audio_chunk(audio_bytes, final_chunk, last_audio_result)
+                    if audio_format is not None and not audio_format and audio_bytes:
+                        audio_format.update(_encoded_audio_format(audio_response, output_sample_rate, final_chunk))
                     if wav_header is not None:
                         yield wav_header
                     if include_sample_rate:
@@ -1880,6 +1952,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         """
         usage_acc = SpeechOutputTokenCounter()
         emitted_audio = False
+        audio_format: dict[str, int] = {}
         try:
             async with aclosing(
                 self._generate_audio_chunks(
@@ -1892,6 +1965,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     usage_acc=usage_acc,
                     tts_params=tts_params,
                     target_sample_rate=request.sample_rate if request is not None else None,
+                    audio_format=audio_format,
                 )
             ) as chunks:
                 async for chunk in chunks:
@@ -1900,6 +1974,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         "audio": base64.b64encode(chunk).decode("ascii"),
                         "response_format": response_format,
                     }
+                    if audio_format:
+                        # Additive to the OpenAI schema: pcm deltas are
+                        # model-native s16le, not always 24 kHz mono.
+                        payload["sample_rate"] = audio_format["sample_rate"]
+                        payload["channels"] = audio_format["channels"]
                     data = json.dumps(payload, separators=(",", ":"))
                     emitted_audio = True
                     yield f"event: speech.audio.delta\ndata: {data}\n\n"
@@ -2134,6 +2213,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         collect: dict | None = None,
         target_sample_rate: int | None = None,
         cumulative_audio: bool = False,
+        audio_format: dict[str, int] | None = None,
     ):
         """Yield raw PCM byte chunks from the engine generator.
 
@@ -2154,6 +2234,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 collect=collect,
                 target_sample_rate=target_sample_rate,
                 cumulative_audio=cumulative_audio,
+                audio_format=audio_format,
             )
         ) as chunks:
             async for chunk in chunks:
@@ -2355,6 +2436,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                 base64_encode=base64_encode,
             )
             audio_response: AudioResponse = self.create_audio(audio_obj)
+            if collect is not None:
+                collect["audio_format"] = _encoded_audio_format(
+                    audio_response, request.sample_rate or sample_rate, audio_tensor
+                )
             self._mark_ref_audio_artifact_ready_for_request(request_id)
             artifact_ready = True
             if usage_out is not None:
@@ -2634,7 +2719,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     request_id=request_id,
                     arrival_time=request_arrival_ts,
                 )
-                return _SpeechStreamingResponse(
+                audio_format: dict[str, int] = {}
+                # Pull the first chunk before sending headers: the body has no
+                # self-describing format for pcm and the rate/channel count is
+                # model-native (not always 24 kHz mono), so the headers state it.
+                body = await _PrimedAudioStream.start(
                     self._generate_audio_chunks(
                         generator,
                         request_id,
@@ -2644,8 +2733,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                         request_arrival_ts=request_arrival_ts,
                         tts_params=raw_tts_params,
                         target_sample_rate=request.sample_rate,
-                    ),
+                        audio_format=audio_format,
+                    )
+                )
+                return _SpeechStreamingResponse(
+                    body,
                     media_type=media_type,
+                    headers=_audio_format_headers(audio_format),
                 )
 
             if request.is_sse_stream():
@@ -2735,6 +2829,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
                     )
             if collect.get("audio_truncated") is not None:
                 headers["X-Audio-Truncated"] = "true" if collect["audio_truncated"] else "false"
+            headers.update(_audio_format_headers(collect.get("audio_format")))
             return Response(content=audio_bytes, media_type=media_type, headers=headers)
 
         except asyncio.CancelledError:

vllm_omni/entrypoints/openai/serving_speech_stream.py

@@ -30,6 +30,9 @@ Protocol:
     Server -> Client (default, word_timestamps=false):
         {"type": "audio.start", "utterance_index": 0, "sentence_index": 0,
          "sentence_text": "...", "format": "wav"}
+        # With stream_audio=true, audio.start is sent with the first audio
+        # chunk; for format "pcm" it also carries that audio's "sample_rate"
+        # and "channels" (s16le, interleaved, model-native).
         <binary frame: audio bytes>
         ...
         {"type": "audio.done", "utterance_index": 0, "sentence_index": 0}
@@ -56,6 +59,7 @@ import asyncio
 import base64
 import json
 import time
+from collections.abc import Awaitable, Callable
 from contextlib import aclosing
 
 from fastapi import WebSocket, WebSocketDisconnect
@@ -75,7 +79,7 @@ logger = init_logger(__name__)
 _DEFAULT_IDLE_TIMEOUT = 30.0  # seconds
 _DEFAULT_CONFIG_TIMEOUT = 10.0  # seconds
 _PCM_SAMPLE_RATE = 24000
-_BYTES_PER_SAMPLE = 2  # 16-bit mono PCM
+_BYTES_PER_SAMPLE = 2  # 16-bit PCM, per channel
 _MAX_CONFIG_MESSAGE_SIZE = 4 * 1024 * 1024  # allow large ref_audio payloads
 _MAX_INPUT_TEXT_MESSAGE_SIZE = 128 * 1024
 
@@ -336,14 +340,29 @@ class OmniStreamingSpeechHandler:
             "sentence_text": sentence_text,
             "format": response_format,
         }
-        if config.stream_audio and response_format == "pcm":
-            # Nominal stream rate; each audio.chunk carries the authoritative
-            # per-chunk sample_rate.
-            start_payload["sample_rate"] = _PCM_SAMPLE_RATE
         if config.word_timestamps:
             start_payload["word_timestamps"] = True
+        # Streamed pcm is model-native s16le. Its rate and channel count are
+        # known once the first chunk is encoded, so for pcm audio.start is
+        # deferred to that point and states them; it is sent bare if no audio
+        # arrives. Non-streamed responses still send it before generation.
+        describe_pcm = config.stream_audio and response_format == "pcm"
+        audio_format: dict[str, int] = {}
+        start_sent = False
+
+        async def send_start() -> None:
+            nonlocal start_sent
+            if start_sent:
+                return
+            start_sent = True
+            if describe_pcm and audio_format:
+                start_payload["sample_rate"] = audio_format["sample_rate"]
+                start_payload["channels"] = audio_format["channels"]
             await websocket.send_json(start_payload)
 
+        if not describe_pcm:
+            await send_start()
+
         total_bytes = 0
         generation_failed = False
         request_id = None
@@ -365,6 +384,8 @@ class OmniStreamingSpeechHandler:
                         request_start_s=request_start_s,
                         request_arrival_ts=request_arrival_ts,
                         tts_params=tts_params,
+                        audio_format=audio_format,
+                        send_start=send_start,
                     )
                 else:
                     async with aclosing(
@@ -374,9 +395,11 @@ class OmniStreamingSpeechHandler:
                             request_start_s=request_start_s,
                             request_arrival_ts=request_arrival_ts,
                             tts_params=tts_params,
+                            audio_format=audio_format,
                         )
                     ) as stream:
                         async for chunk in stream:
+                            await send_start()
                             total_bytes += len(chunk)
                             await websocket.send_bytes(chunk)
             else:
@@ -401,6 +424,10 @@ class OmniStreamingSpeechHandler:
                 sentence_index,
                 e,
             )
+            try:
+                await send_start()
+            except Exception:
+                logger.debug("Failed to send audio.start for sentence %d", sentence_index, exc_info=True)
             await self._send_error(
                 websocket,
                 f"Generation failed for utterance {utterance_index}, sentence {sentence_index}: {e}",
@@ -408,6 +435,9 @@ class OmniStreamingSpeechHandler:
             )
         finally:
             try:
+                if not start_sent:
+                    # Keep audio.start/audio.done paired when no audio was produced.
+                    await send_start()
                 await websocket.send_json(
                     {
                         "type": "audio.done",
@@ -433,6 +463,8 @@ class OmniStreamingSpeechHandler:
         request_arrival_ts: float,
         language: str | None = None,
         tts_params: dict | None = None,
+        audio_format: dict[str, int] | None = None,
+        send_start: Callable[[], Awaitable[None]] | None = None,
     ) -> int:
         """Stream PCM as JSON ``audio.chunk`` frames, aligned per sentence.
 
@@ -450,6 +482,10 @@ class OmniStreamingSpeechHandler:
         chunk_id = 0
         # Receives the aligner stage's pooling output from the generator.
         collect: dict = {}
+        audio_format = audio_format if audio_format is not None else {}
+
+        def frame_bytes() -> int:
+            return _BYTES_PER_SAMPLE * int(audio_format.get("channels", 1))
 
         async def send_chunk(
             chunk: bytes,
@@ -459,6 +495,8 @@ class OmniStreamingSpeechHandler:
             chunk_end_ms: int,
         ) -> None:
             nonlocal chunk_id
+            if send_start is not None:
+                await send_start()
             await websocket.send_json(
                 {
                     "type": "audio.chunk",
@@ -468,6 +506,7 @@ class OmniStreamingSpeechHandler:
                     "chunk_start_ms": chunk_start_ms,
                     "chunk_end_ms": chunk_end_ms,
                     "sample_rate": chunk_sample_rate,
+                    "channels": int(audio_format.get("channels", 1)),
                     "audio_b64": base64.b64encode(chunk).decode("ascii"),
                     "timestamps": timestamps_payload,
                 }
@@ -484,13 +523,14 @@ class OmniStreamingSpeechHandler:
                 tts_params=tts_params,
                 collect=collect,
                 cumulative_audio=True,
+                audio_format=audio_format,
             )
         ) as stream:
             async for chunk, chunk_sample_rate in stream:
                 sample_rate = chunk_sample_rate
-                chunk_start_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0))
+                chunk_start_ms = int(round((audio_bytes_seen / frame_bytes() / sample_rate) * 1000.0))
                 audio_bytes_seen += len(chunk)
-                chunk_end_ms = int(round((audio_bytes_seen / _BYTES_PER_SAMPLE / sample_rate) * 1000.0))
+                chunk_end_ms = int(round((audio_bytes_seen / frame_bytes() / 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)
@@ -502,7 +542,7 @@ class OmniStreamingSpeechHandler:
         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))
+        sentence_end_ms = int(round((audio_bytes_seen / frame_bytes() / sample_rate) * 1000.0))
         await send_chunk(b"", sample_rate, timestamps_payload, 0, sentence_end_ms)
 
         return total_bytes