[Feat][OmniVoice]Support Varlen Attn, Request-Batch and Step-Execution

docs/serving/speech_api.md

@@ -752,6 +752,10 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne
 | ------- | ------------- |
 | `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. |
 
+OmniVoice uses packed variable-length attention for batched generator execution. The attention operator accepts FP16 and BF16 inputs, so its query,
+key, and value tensors are evaluated in BF16 even when the stage is configured with `dtype: float32`; the attention output is converted back to the model's
+hidden-state dtype before the output projection. Consequently, a float32 stage configuration does not imply FP32 attention arithmetic.
+
 ### VoxCPM2
 
 | Model | Description |

tests/e2e/online_serving/test_omnivoice_expansion.py

@@ -7,15 +7,18 @@
 accessed through the standard OpenAI-compatible speech API.
 """
 
+import io
 import os
+import wave
+from collections.abc import Sequence
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 import pytest
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import get_asset_path
-from tests.helpers.runtime import OmniServerParams
+from tests.helpers.runtime import OmniResponse, OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.entrypoints.openai.serving_speech import _DEFAULT_VOICE_NAME
 
@@ -34,6 +37,10 @@
 EXTRA_ARGS = [
     "--trust-remote-code",
     "--disable-log-stats",
+    "--max-num-seqs",
+    "8",
+    "--request-batch-max-wait-ms",
+    "50",
 ]
 TEST_PARAMS = [
     OmniServerParams(
@@ -42,7 +49,14 @@
         server_args=EXTRA_ARGS,
     )
 ]
-
+STEP_EXECUTION_ARGS = EXTRA_ARGS + ["--step-execution"]
+STEP_EXECUTION_PARAMS = [
+    OmniServerParams(
+        model=MODEL,
+        stage_config_path=STAGE_CONFIG,
+        server_args=STEP_EXECUTION_ARGS,
+    )
+]
 # Lower this in ``request_config`` via ``min_audio_bytes`` if a run produces legitimately short WAVs.
 _DEFAULT_MIN_AUDIO_BYTES = 5000
 
@@ -58,6 +72,33 @@ def get_prompt(prompt_type="text"):
     return prompts.get(prompt_type, prompts["text"])
 
 
+def _assert_valid_omnivoice_wav_responses(responses: Sequence[OmniResponse]) -> None:
+    """Validate that every concurrent response contains plausible OmniVoice audio."""
+    audio_shapes: list[tuple[int, int, int, int]] = []
+    for response in responses:
+        audio_bytes = response.audio_bytes
+        assert audio_bytes is not None
+        assert audio_bytes.startswith(b"RIFF")
+        with wave.open(io.BytesIO(audio_bytes), "rb") as wav_file:
+            num_channels = wav_file.getnchannels()
+            sample_width = wav_file.getsampwidth()
+            sample_rate = wav_file.getframerate()
+            num_frames = wav_file.getnframes()
+            assert num_channels == 1
+            assert sample_width == 2
+            assert sample_rate == 24000
+            assert num_frames > 0
+            audio_shapes.append((num_channels, sample_width, sample_rate, num_frames))
+
+        duration_s = num_frames / sample_rate
+        assert 0.1 <= duration_s <= 30.0
+
+    # Identical prompts use the same estimated target length. Their samples may
+    # differ numerically across batch layouts, but the decoded audio shape must
+    # remain consistent across concurrent requests.
+    assert len(set(audio_shapes)) == 1
+
+
 @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
 class TestOmniVoiceTTS:
     """E2E tests for OmniVoice TTS model."""
@@ -74,6 +115,57 @@ def test_speech_auto_voice(self, omni_server, online_client) -> None:
         }
         online_client.send_audio_speech_request(request_config)
 
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    def test_speech_auto_voice_batch(self, omni_server, openai_client) -> None:
+        """Test concurrent request-batch TTS generation."""
+        batch_request_config = {
+            "model": omni_server.model,
+            "input": get_prompt("text"),
+            "response_format": "wav",
+            "seed": 42,
+            "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
+        }
+        batch_r = openai_client.send_audio_speech_request(batch_request_config, request_num=2)
+        assert len(batch_r) == 2
+        _assert_valid_omnivoice_wav_responses(batch_r)
+
+
+@pytest.mark.parametrize("omni_server", STEP_EXECUTION_PARAMS, indirect=True)
+class TestOmniVoiceStepExecution:
+    """E2E tests for OmniVoice TTS model."""
+
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    def test_speech_auto_voice_step_execution(self, omni_server, openai_client) -> None:
+        """Test auto voice TTS generation (text only, no reference audio)."""
+        request_config = {
+            "model": omni_server.model,
+            "input": get_prompt("text"),
+            "response_format": "wav",
+            "timeout": 180.0,
+            "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
+            "extra_params": {
+                "num_inference_steps": 32,
+            },
+        }
+        openai_client.send_audio_speech_request(request_config)
+
+    @hardware_test(res={"cuda": "L4"}, num_cards=1)
+    def test_speech_auto_voice_batch_step_execution(self, omni_server, openai_client) -> None:
+        """Test concurrent step-execution TTS generation."""
+        batch_request_config = {
+            "model": omni_server.model,
+            "input": get_prompt("text"),
+            "response_format": "wav",
+            "seed": 42,
+            "min_audio_bytes": _DEFAULT_MIN_AUDIO_BYTES,
+            "extra_params": {
+                "num_inference_steps": 32,
+            },
+        }
+        batch_r = openai_client.send_audio_speech_request(batch_request_config, request_num=2)
+        assert len(batch_r) == 2
+        _assert_valid_omnivoice_wav_responses(batch_r)
+
 
 @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
 class TestOmniVoiceSeed:

tests/e2e/online_serving/test_omnivoice_parity.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cross-mode full-model parity tests for OmniVoice."""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+import requests
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
+
+MODEL = "k2-fsa/OmniVoice"
+STAGE_CONFIG = get_deploy_config_path("omnivoice.yaml")
+PROMPT = "The weather is nice today, perfect for a walk in the park."
+
+payload = {
+    "model": MODEL,
+    "input": PROMPT,
+    "language": "English",
+    "seed": 42,
+    "response_format": "wav",
+    "extra_params": {"num_inference_steps": 32},
+}
+
+
+def _generate_without_graph(server_args: list[str]) -> bytes:
+    with OmniServer(
+        MODEL,
+        server_args,
+        use_omni=True,
+        env_dict={"OMNIVOICE_CUDA_GRAPH": "0"},
+    ) as server:
+        response = requests.post(
+            f"http://{server.host}:{server.port}/v1/audio/speech",
+            json=payload,
+            timeout=600,
+        )
+        response.raise_for_status()
+        assert response.content.startswith(b"RIFF")
+        return response.content
+
+
+def _generate_with_graph(server_args: list[str]) -> bytes:
+    with OmniServer(
+        MODEL,
+        server_args,
+        use_omni=True,
+    ) as server:
+        response = requests.post(
+            f"http://{server.host}:{server.port}/v1/audio/speech",
+            json=payload,
+            timeout=600,
+        )
+        response.raise_for_status()
+        assert response.content.startswith(b"RIFF")
+        return response.content
+
+
+def _common_args() -> list[str]:
+    """Return the shared B=1 server arguments for parity tests."""
+    return [
+        "--trust-remote-code",
+        "--disable-log-stats",
+        "--deploy-config",
+        STAGE_CONFIG,
+        "--max-num-seqs",
+        "1",
+    ]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_request_mode_and_step_execution_b1_parity_without_graph() -> None:
+    """B=1 eager request and step modes must produce identical seeded WAV bytes."""
+    common_args = _common_args()
+
+    request_audio = _generate_without_graph(common_args)
+    step_audio = _generate_without_graph([*common_args, "--step-execution", "--enforce-eager"])
+
+    assert request_audio == step_audio
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_request_mode_and_step_execution_b1_parity_with_graph() -> None:
+    """B=1 Graph request and step modes must produce identical seeded WAV bytes."""
+    common_args = _common_args()
+
+    request_audio = _generate_with_graph(common_args)
+    step_audio = _generate_with_graph([*common_args, "--step-execution", "--enforce-eager"])
+
+    assert request_audio == step_audio

tests/entrypoints/openai_api/test_serving_speech.py

@@ -27,6 +27,9 @@
 from pytest_mock import MockerFixture
 from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
 
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
+from vllm_omni.diffusion.sched.step_scheduler import StepScheduler
 from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
 from vllm_omni.entrypoints.openai import api_server as api_server_module
 from vllm_omni.entrypoints.openai import serving_speech as serving_speech_module
@@ -55,6 +58,7 @@
 from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter
 from vllm_omni.entrypoints.openai.tts_adapters.qwen3_tts import Qwen3TTSCodecLimitError
 from vllm_omni.entrypoints.openai.tts_adapters.voxtral import VoxtralTTSAdapter
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
     FISH_TEXT_ONLY_SYSTEM_PROMPT,
     build_fish_voice_clone_prompt_ids,
@@ -753,13 +757,12 @@ async def test_diffusion_create_speech_with_unknown_voice(self, mocker: MockerFi
 
     @pytest.mark.asyncio
     async def test_create_diffusion_speech_extra_params(self, mocker: MockerFixture):
-        """Test public diffusion speech success and extra_params propagation."""
+        """Test diffusion parameters reach StepScheduler as standard fields."""
         # Mock the engine client
         mock_engine = mocker.MagicMock()
 
         # Mock default sampling params
-        mock_sampling_param = mocker.MagicMock()
-        mock_sampling_param.extra_args = {"existing_arg": "value"}
+        mock_sampling_param = OmniDiffusionSamplingParams(extra_args={"existing_arg": "value"})
         mock_engine.default_sampling_params_list = [mock_sampling_param]
 
         # Mock generate to yield a valid OmniRequestOutput
@@ -775,7 +778,15 @@ async def mock_generate(*args, **kwargs):
             server, "create_audio", return_value=mocker.MagicMock(audio_data=b"dummy", media_type="audio/wav")
         )
 
-        req = OpenAICreateSpeechRequest(input="Hello", extra_params={"new_arg": 123, "existing_arg": "new_value"})
+        req = OpenAICreateSpeechRequest(
+            input="Hello",
+            extra_params={
+                "new_arg": 123,
+                "existing_arg": "new_value",
+                "num_inference_steps": 12,
+                "guidance_scale": 7.0,
+            },
+        )
 
         response = await server.create_speech(req)
 
@@ -792,7 +803,98 @@ async def mock_generate(*args, **kwargs):
 
         # Verify it was deepcopied and updated
         assert passed_params is not mock_engine.default_sampling_params_list
-        assert passed_params[0].extra_args == {"existing_arg": "new_value", "new_arg": 123}
+        assert passed_params[0].extra_args == {
+            "existing_arg": "new_value",
+            "new_arg": 123,
+            "num_inference_steps": 12,
+            "guidance_scale": 7.0,
+        }
+        assert passed_params[0].num_inference_steps == 12
+        assert passed_params[0].guidance_scale == 7.0
+
+        # Regression: StepScheduler.add_request() used to receive
+        # num_inference_steps=None and fail while converting it to int.
+        scheduler = StepScheduler()
+        scheduler.add_request(
+            OmniDiffusionRequest(
+                prompt="Hello",
+                sampling_params=passed_params[0],
+                request_id="speech-test",
+            )
+        )
+        assert scheduler._request_progress["speech-test"].total_steps == 12
+
+    @pytest.mark.asyncio
+    async def test_diffusion_speech_guidance_promotion_controls_request_batch_admission(
+        self,
+        mocker: MockerFixture,
+    ) -> None:
+        """Different request guidance values must not enter one request batch."""
+        mock_engine = mocker.MagicMock()
+        mock_engine.default_sampling_params_list = [OmniDiffusionSamplingParams(num_inference_steps=12)]
+        passed_sampling_params = []
+
+        async def mock_generate(*args, **kwargs):
+            passed_sampling_params.append(kwargs["sampling_params_list"][0])
+            yield create_mock_audio_output_for_test()
+
+        mock_engine.generate = mocker.MagicMock(side_effect=mock_generate)
+        server = OmniOpenAIServingSpeech.for_diffusion(diffusion_engine=mock_engine, model_name="test-model")
+        mocker.patch.object(
+            server,
+            "create_audio",
+            return_value=mocker.MagicMock(audio_data=b"dummy", media_type="audio/wav"),
+        )
+
+        for guidance_scale in (2.0, 7.0):
+            response = await server.create_speech(
+                OpenAICreateSpeechRequest(
+                    input="Hello",
+                    extra_params={"guidance_scale": guidance_scale},
+                )
+            )
+            assert response.status_code == 200
+
+        scheduler = RequestScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+        for index, sampling_params in enumerate(passed_sampling_params):
+            scheduler.add_request(
+                OmniDiffusionRequest(
+                    prompt="Hello",
+                    sampling_params=sampling_params,
+                    request_id=f"speech-{index}",
+                )
+            )
+
+        first = scheduler.schedule()
+
+        assert [request.request_id for request in first.scheduled_new_reqs] == ["speech-0"]
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("extra_params", "expected_message"),
+        [
+            ({"num_inference_steps": "invalid"}, "num_inference_steps must be an integer"),
+            ({"guidance_scale": "invalid"}, "guidance_scale must be a number"),
+        ],
+    )
+    async def test_create_diffusion_speech_rejects_invalid_scheduler_params(
+        self,
+        mocker: MockerFixture,
+        extra_params: dict[str, str],
+        expected_message: str,
+    ) -> None:
+        mock_engine = mocker.MagicMock()
+        mock_engine.default_sampling_params_list = [OmniDiffusionSamplingParams()]
+        server = OmniOpenAIServingSpeech.for_diffusion(diffusion_engine=mock_engine, model_name="test-model")
+
+        response = await server.create_speech(OpenAICreateSpeechRequest(input="Hello", extra_params=extra_params))
+
+        assert response.status_code == 400
+        assert expected_message in response.body.decode()
+        mock_engine.generate.assert_not_called()
 
 
 class TestTTSMethods: