Merge branch 'main' into dev/rebase-v0.18.0

.buildkite/test-nightly.yml

@@ -73,7 +73,7 @@ steps:
   - label: ":full_moon: Diffusion Model Test with H100"
     timeout_in_minutes: 60
     depends_on: upload-nightly-pipeline
-    # if: build.env("NIGHTLY") == "1"
+    if: build.env("NIGHTLY") == "1"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
       - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
@@ -164,6 +164,8 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
+                    value: "120"
                   - name: HF_TOKEN
                     valueFrom:
                       secretKeyRef:
@@ -183,7 +185,7 @@ steps:
   - label: ":email: Nightly Collection & Email"
     key: nightly-perf-distribution
     depends_on:
-    - nightly-performance
+      - nightly-performance
     if: build.env("NIGHTLY") == "1"
     commands:
       - pip install openpyxl

docs/models/supported_models.md

@@ -52,7 +52,7 @@ th {
 |`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
 |`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
 |`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
-|`VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/tts-model` |
+|`VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` |
 
 
 ## List of Supported Models for NPU

docs/serving/speech_api.md

@@ -4,7 +4,7 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
 
 - **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
 - **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
-- **Voxtral TTS** (`mistralai/tts-model`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
+- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
 
 Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
 
@@ -31,7 +31,7 @@ vllm-omni serve fishaudio/s2-pro \
     --gpu-memory-utilization 0.9
 
 # Voxtral TTS
-vllm serve mistralai/tts-model \
+vllm serve mistralai/Voxtral-4B-TTS-2603 \
     --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
     --omni \
     --port 8091 \
@@ -360,7 +360,7 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne
 
 | Model | Description |
 |-------|-------------|
-| `mistralai/tts-model` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |
+| `mistralai/Voxtral-4B-TTS-2603` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |
 
 ## Error Responses
 

examples/offline_inference/voxtral_tts/README.md

@@ -12,28 +12,28 @@ When `mistral_common` has `SpeechRequest` support, prompt token IDs are built vi
 python3 examples/offline_inference/voxtral_tts/end2end.py \
     --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
     --write-audio --voice cheerful_female \
-    --model mistralai/tts-model \
+    --model mistralai/Voxtral-4B-TTS-2603 \
     --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
 
 # 32 replicate prompts with cheerful_female voice preset
 python3 examples/offline_inference/voxtral_tts/end2end.py \
     --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
     --num-prompts 32 --write-audio --voice cheerful_female \
-    --model mistralai/tts-model \
+    --model mistralai/Voxtral-4B-TTS-2603 \
     --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
 
 # Streaming with neutral_female voice preset
 python3 examples/offline_inference/voxtral_tts/end2end.py \
     --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
     --streaming --write-audio --voice neutral_female \
-    --model mistralai/tts-model \
+    --model mistralai/Voxtral-4B-TTS-2603 \
     --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
 
 # 32 prompts, 8 concurrent requests per wave, streaming with neutral_female voice
 python3 examples/offline_inference/voxtral_tts/end2end.py \
     --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
     --num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
-    --model mistralai/tts-model \
+    --model mistralai/Voxtral-4B-TTS-2603 \
     --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
 
 # Short debug prompt with reference audio
@@ -41,7 +41,7 @@ python3 examples/offline_inference/voxtral_tts/end2end.py \
 python3 examples/offline_inference/voxtral_tts/end2end.py \
     --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
     --write-audio \
-    --model mistralai/tts-model \
+    --model mistralai/Voxtral-4B-TTS-2603 \
     --text "This is a test message." \
     --audio-path path/to/reference_audio.wav
 ```
@@ -50,7 +50,7 @@ python3 examples/offline_inference/voxtral_tts/end2end.py \
 
 | Argument | Description |
 |---|---|
-| `--model PATH` | HuggingFace repo ID or local directory path (default: `mistralai/tts-model`) |
+| `--model PATH` | HuggingFace repo ID or local directory path (default: `mistralai/Voxtral-4B-TTS-2603`) |
 | `--text TEXT` | Text to synthesize (default: `"This is a test message."`) |
 | `--audio-path PATH` | Path to reference audio file for voice cloning |
 | `--output-dir DIR` | Directory to write output WAV files (default: `output_audio`) |

examples/offline_inference/voxtral_tts/end2end.py

@@ -229,7 +229,7 @@ def parse_args() -> Namespace:
     parser.add_argument(
         "--model",
         type=str,
-        default="mistralai/tts-model",
+        default="mistralai/Voxtral-4B-TTS-2603",
         help="Model name or path.",
     )
     parser.add_argument(

examples/online_serving/voxtral_tts/gradio_demo.py

@@ -1,3 +1,14 @@
+"""
+- Make sure to install the following for this example to function correctly:
+- `pip install -e .`
+- `pip install gradio==5.50 mistral_common=1.10.0`
+
+Example use case:
+
+python examples/online_serving/voxtral_tts/gradio_demo.py --host slurm-199-077 --port 8000
+
+"""
+
 import argparse
 import io
 import json
@@ -125,11 +136,12 @@ def _load_from_share(
 
 def main(
     model: str,
+    host: str,
+    port: str,
     output_dir: str | None = None,
-    debug: bool = False,
 ) -> None:
-    base_url = "http://localhost:8091/v1"
-    print(f"Using speech API at: {base_url}/audio/speech")
+    base_url = f"http://{host}:{port}/v1"
+    logger.info(f"Using speech API at: {base_url}/audio/speech")
 
     outputs_dir: Path | None = None
     if output_dir is not None:
@@ -236,7 +248,9 @@ def main(
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Voxtral TTS Gradio Demo")
-    parser.add_argument("--model", type=str, required=True, help="Path to the model checkpoint")
+    parser.add_argument("--model", type=str, default="mistralai/Voxtral-4B-TTS-2603", help="Name of model repo on HF")
+    parser.add_argument("--host", type=str, default="localhost", help="Name of host")
+    parser.add_argument("--port", type=str, default="8091", help="port number")
     parser.add_argument(
         "--output-dir",
         type=str,
@@ -244,12 +258,12 @@ if __name__ == "__main__":
         help="Directory to save generated audio and share links. "
         "If not provided, save/share functionality is disabled.",
     )
-    parser.add_argument("--debug", action="store_true", default=False, help="Enable debug mode")
 
     args = parser.parse_args()
 
     main(
         model=args.model,
+        host=args.host,
+        port=args.port,
         output_dir=args.output_dir,
-        debug=args.debug,
     )

tests/e2e/online_serving/test_voxtral_tts.py

@@ -20,7 +20,7 @@ import pytest
 from tests.conftest import OmniServerParams
 from tests.utils import hardware_test
 
-MODEL = "mistralai/tts-model"
+MODEL = "mistralai/Voxtral-4B-TTS-2603"
 
 STAGE_CONFIG = str(
     Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "voxtral_tts.yaml"

tests/perf/tests/test.json

@@ -12,11 +12,13 @@
                 "endpoint": "/v1/chat/completions",
                 "num_prompts": [
                     10,
-                    40
+                    40,
+                    100
                 ],
                 "max_concurrency": [
                     1,
-                    4
+                    4,
+                    10
                 ],
                 "random_input_len": 100,
                 "random_output_len": 100,
@@ -34,11 +36,13 @@
                 "endpoint": "/v1/chat/completions",
                 "num_prompts": [
                     10,
-                    40
+                    40,
+                    100
                 ],
                 "max_concurrency": [
                     1,
-                    4
+                    4,
+                    10
                 ],
                 "random_input_len": 100,
                 "random_output_len": 100,
@@ -96,11 +100,13 @@
                 "endpoint": "/v1/chat/completions",
                 "num_prompts": [
                     10,
-                    40
+                    40,
+                    100
                 ],
                 "max_concurrency": [
                     1,
-                    4
+                    4,
+                    10
                 ],
                 "random_input_len": 100,
                 "random_output_len": 100,
@@ -118,11 +124,13 @@
                 "endpoint": "/v1/chat/completions",
                 "num_prompts": [
                     10,
-                    40
+                    40,
+                    100
                 ],
                 "max_concurrency": [
                     1,
-                    4
+                    4,
+                    10
                 ],
                 "random_input_len": 100,
                 "random_output_len": 100,

vllm_omni/benchmarks/patch/patch.py

@@ -114,91 +114,127 @@ async def async_request_openai_chat_omni_completions(
 
     output = MixRequestFuncOutput()
     output.prompt_len = request_func_input.prompt_len
-
-    generated_text = ""
-    generated_audio = None
-    ttft = 0.0
-    st = time.perf_counter()
-    output.start_time = st
-    most_recent_timestamp = st
-    timestamp = st
-    audio_generate_time = 0.0
-    try:
-        async with session.post(url=api_url, json=payload, headers=headers) as response:
-            if response.status == 200:
-                handler = StreamedResponseHandler()
-                async for chunk_bytes in response.content.iter_any():
-                    chunk_bytes = chunk_bytes.strip()
-                    if not chunk_bytes:
-                        continue
-
-                    messages = handler.add_chunk(chunk_bytes)
-                    for message in messages:
-                        if type(message) is bytes:
-                            message = message.decode("utf-8")
-                        # NOTE: SSE comments (often used as pings) start with
-                        # a colon. These are not JSON data payload and should
-                        # be skipped.
-                        if message.startswith(":"):
+    max_retries = 3
+    retry_delay = 0.1
+    for attempt in range(max_retries + 1):
+        # Reset per-attempt state so that retries do not mix partial
+        # outputs or metrics from previous attempts.
+        generated_text = ""
+        generated_audio = None
+        ttft = 0.0
+        st = time.perf_counter()
+        output.start_time = st
+        most_recent_timestamp = st
+        timestamp = st
+        audio_generate_time = 0.0
+        output.itl = []
+        output.generated_text = ""
+        output.ttft = 0.0
+        output.audio_ttfp = 0.0
+        output.audio_duration = 0.0
+        output.audio_frames = 0
+        output.audio_rtf = 0.0
+        output.text_latency = 0.0
+        output.output_tokens = 0
+        output.error = ""
+        output.success = False
+        try:
+            async with session.post(url=api_url, json=payload, headers=headers) as response:
+                if response.status == 200:
+                    handler = StreamedResponseHandler()
+                    async for chunk_bytes in response.content.iter_any():
+                        chunk_bytes = chunk_bytes.strip()
+                        if not chunk_bytes:
                             continue
 
-                        chunk = message.removeprefix("data: ")
-                        if chunk != "[DONE]":
-                            timestamp = time.perf_counter()
-                            data = json.loads(chunk)
-                            if choices := data.get("choices"):
-                                modality = data.get("modality")
-                                content = choices[0]["delta"].get("content")
-                                if modality == "text":
-                                    # First token
-                                    if ttft == 0.0:
-                                        ttft = timestamp - st
-                                        output.ttft = ttft
-                                    else:
-                                        output.itl.append(timestamp - most_recent_timestamp)
-                                    generated_text += content or ""
-                                    most_recent_timestamp = timestamp
-                                    output.text_latency = timestamp - st
-                                elif modality == "audio":
-                                    if output.audio_ttfp == 0.0:
-                                        output.audio_ttfp = timestamp - st
-                                    audio_generate_time = timestamp - st
-                                    if content != "":
-                                        audio_bytes = base64.b64decode(content)
-                                        seg = AudioSegment.from_file(io.BytesIO(audio_bytes))
-                                        if seg is not None:
-                                            if generated_audio is None:
-                                                generated_audio = seg
-                                            else:
-                                                generated_audio = generated_audio + seg
+                        messages = handler.add_chunk(chunk_bytes)
+                        for message in messages:
+                            if type(message) is bytes:
+                                message = message.decode("utf-8")
+                            # NOTE: SSE comments (often used as pings) start with
+                            # a colon. These are not JSON data payload and should
+                            # be skipped.
+                            if message.startswith(":"):
+                                continue
 
-                            if metrics := data.get("metrics"):
-                                output.output_tokens = metrics.get("num_tokens_out", 0)
+                            chunk = message.removeprefix("data: ")
+                            if chunk != "[DONE]":
+                                timestamp = time.perf_counter()
+                                data = json.loads(chunk)
+                                if choices := data.get("choices"):
+                                    modality = data.get("modality")
+                                    content = choices[0]["delta"].get("content")
+                                    if modality == "text":
+                                        # First token
+                                        if ttft == 0.0:
+                                            ttft = timestamp - st
+                                            output.ttft = ttft
+                                        else:
+                                            output.itl.append(timestamp - most_recent_timestamp)
+                                        generated_text += content or ""
+                                        most_recent_timestamp = timestamp
+                                        output.text_latency = timestamp - st
+                                    elif modality == "audio":
+                                        if output.audio_ttfp == 0.0:
+                                            output.audio_ttfp = timestamp - st
+                                        audio_generate_time = timestamp - st
+                                        if content != "":
+                                            audio_bytes = base64.b64decode(content)
+                                            seg = AudioSegment.from_file(io.BytesIO(audio_bytes))
+                                            if seg is not None:
+                                                if generated_audio is None:
+                                                    generated_audio = seg
+                                                else:
+                                                    generated_audio = generated_audio + seg
 
-                output.latency = timestamp - st
-                output.generated_text = generated_text
-                if generated_audio is not None:
-                    output.audio_duration = len(generated_audio) / 1000.0
-                    frame_width = generated_audio.frame_width
-                    if frame_width > 0:
-                        output.audio_frames = len(generated_audio.raw_data) // frame_width
-                    else:
-                        output.audio_frames = 0
-                        logger.warning("Audio frame width is zero")
-                    audio_duration = output.audio_duration
-                    if audio_duration > 0:
-                        output.audio_rtf = audio_generate_time / output.audio_duration
-                    else:
-                        output.audio_rtf = 0
-                        logger.warning("Audio duration is zero")
-                output.success = True
-            else:
-                output.error = response.reason or ""
-                output.success = False
-    except Exception:
-        output.success = False
-        output.error = traceback.format_exc()
-        logger.error(f"ERROR: send request failed, reason is: {output.error}")
+                                if metrics := data.get("metrics"):
+                                    output.output_tokens = metrics.get("num_tokens_out", 0)
+
+                    output.latency = timestamp - st
+                    output.generated_text = generated_text
+                    if generated_audio is not None:
+                        output.audio_duration = len(generated_audio) / 1000.0
+                        frame_width = generated_audio.frame_width
+                        if frame_width > 0:
+                            output.audio_frames = len(generated_audio.raw_data) // frame_width
+                        else:
+                            output.audio_frames = 0
+                            logger.warning("Audio frame width is zero")
+                        audio_duration = output.audio_duration
+                        if audio_duration > 0:
+                            output.audio_rtf = audio_generate_time / output.audio_duration
+                        else:
+                            output.audio_rtf = 0
+                            logger.warning("Audio duration is zero")
+                    output.success = True
+                else:
+                    output.error = response.reason or ""
+                    output.success = False
+            break
+        except aiohttp.ClientError as e:
+            # transient transport error: may retry
+            output.success = False
+            output.error = traceback.format_exc()
+            if attempt < max_retries:
+                logger.warning(
+                    "ClientError in omni benchmark request (will retry): attempt=%d/%d delay=%.2fs: %s",
+                    attempt + 1,
+                    max_retries + 1,
+                    retry_delay,
+                    str(e),
+                )
+                await asyncio.sleep(retry_delay)
+                continue
+            logger.error(
+                "ClientError in omni benchmark request (giving up):\n%s",
+                output.error,
+            )
+            break
+        except Exception:
+            output.success = False
+            output.error = traceback.format_exc()
+            logger.error(f"ERROR: send request failed, reason is: {output.error}")
+            break
 
     if pbar:
         pbar.update(1)

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -52,7 +52,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.get_req_chunk: dict[str, int] = defaultdict(int)
         self.finished_requests: set[str] = set()
         self.request_payload = {}
-        self.code_prompt_token_ids: dict[str, list[list[int]]] = defaultdict(list)
+        self.code_prompt_token_ids: dict[str, list[torch.Tensor]] = defaultdict(list)
         self.request_ids_mapping: dict[str, str] = {}
 
         self.waiting_for_chunk_waiting_requests: deque[Any] = deque()

vllm_omni/entrypoints/openai/serving_chat.py

@@ -299,7 +299,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 if not extracted_prompt:
                     return self.create_error_response("No text prompt found in messages")
 
-                extra_body = getattr(request, "extra_body", None) or {}
+                # [NOTE] When sending request via openai client Python library,
+                #   `extra_body` is flattented and merged into the payload's root.
+                #   These extra fields are accessible via `model_extra` property (from Pydantic base class).
+                #   When sending raw request with curl, no flattening happens. Directly read the `extra_body` dict.
+                extra_body = getattr(request, "extra_body", None)
+                if not extra_body:
+                    extra_body = request.model_extra or {}
                 height = extra_body.get("height")
                 width = extra_body.get("width")
                 if "size" in extra_body:
@@ -2038,7 +2044,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
             # Extract generation parameters from extra_body (preferred)
             # Reference: text_to_image.py and text_to_video.py for supported parameters
-            extra_body = getattr(request, "extra_body", None) or {}
+            # [NOTE] When sending request via openai client Python library,
+            #   `extra_body` is flattented and merged into the payload's root.
+            #   These extra fields are accessible via `model_extra` property (from Pydantic base class).
+            #   When sending raw request with curl, no flattening happens. Directly read the `extra_body` dict.
+            extra_body = getattr(request, "extra_body", None)
+            if not extra_body:
+                extra_body = request.model_extra or {}
 
             # Parse size if provided (supports "1024x1024" format)
             height = extra_body.get("height")

vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py

@@ -17,6 +17,7 @@ from typing import Any
 
 import torch
 import torch.nn as nn
+from torch.nn.utils.parametrize import remove_parametrizations
 from vllm.config import VllmConfig
 from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.logger import init_logger
@@ -58,6 +59,56 @@ class FishSpeechDACDecoder(nn.Module):
         self._hop_length: int = DAC_HOP_LENGTH
         self._logged_codec_stats = False
 
+    def _bake_weight_norm(self, codec: nn.Module) -> None:
+        baked = 0
+        for module in codec.modules():
+            parametrizations = getattr(module, "parametrizations", None)
+            if not parametrizations:
+                continue
+            for name in list(parametrizations.keys()):
+                remove_parametrizations(module, name, leave_parametrized=True)
+                baked += 1
+        if baked > 0:
+            logger.info("Baked %d DAC parametrized weights for inference", baked)
+
+    def _cache_attention_masks(self, codec: nn.Module) -> None:
+        for module in codec.modules():
+            if not hasattr(module, "make_mask") or not hasattr(module, "make_window_limited_mask"):
+                continue
+
+            base_make_mask = module.make_mask
+            base_make_window_mask = module.make_window_limited_mask
+            mask_cache: dict[int, torch.Tensor] = {}
+            window_mask_cache: dict[int, torch.Tensor] = {}
+
+            def make_mask_cached(max_length: int, x_lens: torch.Tensor | None = None, *, _orig=base_make_mask):
+                if x_lens is not None:
+                    return _orig(max_length, x_lens)
+                key = int(max_length)
+                cached = mask_cache.get(key)
+                if cached is None:
+                    cached = _orig(max_length, x_lens)
+                    mask_cache[key] = cached
+                return cached
+
+            def make_window_mask_cached(
+                max_length: int,
+                x_lens: torch.Tensor | None = None,
+                *,
+                _orig=base_make_window_mask,
+            ):
+                if x_lens is not None:
+                    return _orig(max_length, x_lens)
+                key = int(max_length)
+                cached = window_mask_cache.get(key)
+                if cached is None:
+                    cached = _orig(max_length, x_lens)
+                    window_mask_cache[key] = cached
+                return cached
+
+            module.make_mask = make_mask_cached
+            module.make_window_limited_mask = make_window_mask_cached
+
     def _ensure_codec_loaded(self) -> None:
         if self._codec is not None:
             return
@@ -87,6 +138,8 @@ class FishSpeechDACDecoder(nn.Module):
         if "generator" in state_dict:
             state_dict = state_dict["generator"]
         codec.load_state_dict(state_dict, strict=False)
+        self._bake_weight_norm(codec)
+        self._cache_attention_masks(codec)
 
         device = self.vllm_config.device_config.device
         codec = codec.to(device=device, dtype=torch.float32)
@@ -160,7 +213,7 @@ class FishSpeechDACDecoder(nn.Module):
         ids = input_ids.reshape(-1).to(dtype=torch.long)
         request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
 
-        parsed: list[tuple[int, int]] = []
+        parsed_ctx_frames: list[int] = []
         valid_codes_qf: list[torch.Tensor] = []
         valid_indices: list[int] = []
         left_context_size = [0] * len(request_ids_list)
@@ -173,7 +226,7 @@ class FishSpeechDACDecoder(nn.Module):
 
         for i, req_ids in enumerate(request_ids_list):
             if req_ids.numel() < 1:
-                parsed.append((0, 0))
+                parsed_ctx_frames.append(0)
                 continue
             ctx_frames = left_context_size[i]
             flat = req_ids
@@ -185,11 +238,11 @@ class FishSpeechDACDecoder(nn.Module):
                         n,
                         q,
                     )
-                parsed.append((0, 0))
+                parsed_ctx_frames.append(0)
                 continue
             frames = n // q
             codes_qf = flat.reshape(q, frames)
-            parsed.append((ctx_frames, frames))
+            parsed_ctx_frames.append(ctx_frames)
             valid_codes_qf.append(codes_qf)
             valid_indices.append(i)
 
@@ -219,23 +272,33 @@ class FishSpeechDACDecoder(nn.Module):
             except Exception:
                 pass
 
-        # Decode each request individually.
-        wav_tensors: list[torch.Tensor] = []
-        for codes_qf in valid_codes_qf:
-            codes_bqf = codes_qf.unsqueeze(0)  # [1, num_codebooks, num_frames]
-            num_frames = codes_qf.shape[1]
-            feature_lengths = torch.tensor([num_frames], device=codes_bqf.device)
-            with torch.cuda.amp.autocast(enabled=False):
-                wav, audio_lengths = self._codec.decode(codes_bqf, feature_lengths)
-            # wav shape: [1, 1, wav_len]
-            wav_tensors.append(wav.squeeze(0).squeeze(0))  # [wav_len]
+        feature_lengths = torch.tensor(
+            [codes_qf.shape[1] for codes_qf in valid_codes_qf],
+            device=valid_codes_qf[0].device,
+            dtype=torch.long,
+        )
+        max_frames = int(feature_lengths.max().item())
+        batch_size = len(valid_codes_qf)
+
+        codes_bqf = torch.zeros(
+            (batch_size, q, max_frames),
+            device=valid_codes_qf[0].device,
+            dtype=torch.long,
+        )
+        for i, codes_qf in enumerate(valid_codes_qf):
+            frame_count = int(feature_lengths[i].item())
+            codes_bqf[i, :, :frame_count] = codes_qf
+
+        with torch.amp.autocast("cuda", enabled=False):
+            wav_batch, audio_lengths = self._codec.decode(codes_bqf, feature_lengths)
 
         audios: list[torch.Tensor] = [empty] * num_req
         srs = [sr_tensor] * num_req
 
         for j, idx in enumerate(valid_indices):
-            ctx_frames, actual_frames = parsed[idx]
-            wav = wav_tensors[j]
+            ctx_frames = parsed_ctx_frames[idx]
+            audio_len = int(audio_lengths[j].item()) if audio_lengths.numel() > j else int(wav_batch.shape[-1])
+            wav = wav_batch[j, 0, :audio_len]
             # Trim context frames (left overlap for streaming).
             if ctx_frames > 0:
                 cut = ctx_frames * self._hop_length

vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py

@@ -308,6 +308,8 @@ class FishSpeechFastAR(nn.Module):
         self._embed_buf: torch.Tensor | None = None
         self._pos_ids: torch.Tensor | None = None
         self._compiled_model_fwd: object | None = None
+        self._compile_attempted = False
+        self._compile_failed = False
 
     def _ensure_buffers(self, bsz: int, device: torch.device, dtype: torch.dtype) -> None:
         max_seq = self._num_codebooks + 1  # hidden_state + num_codebooks codes
@@ -322,11 +324,61 @@ class FishSpeechFastAR(nn.Module):
         self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device)
 
     def _setup_compile(self) -> None:
-        if self._compiled_model_fwd is not None:
+        if self._compile_attempted:
             return
-        # TODO: Enable torch.compile for performance.  Eager for now to avoid
-        # potential graph-break issues during initial bring-up.
-        self._compiled_model_fwd = self.model.forward
+        self._compile_attempted = True
+        try:
+            self._compiled_model_fwd = torch.compile(
+                self.model.forward,
+                # Keep the helper compiler separate from vLLM's outer
+                # cudagraph-managed Stage-0 execution.
+                mode="default",
+                dynamic=True,
+                fullgraph=False,
+            )
+        except Exception as exc:
+            self._compile_failed = True
+            logger.warning("Failed to enable torch.compile for Fish Speech Fast AR: %s", exc)
+            self._compiled_model_fwd = self.model.forward
+        else:
+            logger.info("Enabled torch.compile for Fish Speech Fast AR forward (mode=default)")
+
+    @torch.inference_mode()
+    def warmup_compile(
+        self,
+        device: torch.device,
+        dtype: torch.dtype,
+        batch_sizes: tuple[int, ...] = (1,),
+    ) -> None:
+        self._setup_compile()
+        if self._compiled_model_fwd is self.model.forward or self._compile_failed:
+            return
+        for batch_size in batch_sizes:
+            hidden = torch.zeros((batch_size, self.slow_ar_config.hidden_size), device=device, dtype=dtype)
+            semantic = torch.full(
+                (batch_size,),
+                self.slow_ar_config.semantic_begin_id,
+                device=device,
+                dtype=torch.long,
+            )
+            self(hidden, semantic, do_sample=False)
+        torch.cuda.synchronize(device)
+
+    @torch.inference_mode()
+    def _run_model(self, step_input: torch.Tensor, step_pos_ids: torch.Tensor, bsz: int) -> torch.Tensor:
+        # Default-on compile only pays off for single-request decode. For
+        # batched decode, eager preserves loaded throughput and avoids the
+        # regression seen with batch>1 compiled execution.
+        model_fwd = self._compiled_model_fwd if bsz == 1 else self.model.forward
+        try:
+            return model_fwd(step_input, step_pos_ids)
+        except Exception as exc:
+            if model_fwd is self.model.forward or self._compile_failed:
+                raise
+            self._compile_failed = True
+            self._compiled_model_fwd = self.model.forward
+            logger.warning("Fish Speech Fast AR torch.compile fallback to eager after runtime failure: %s", exc)
+            return self.model.forward(step_input, step_pos_ids)
 
     @torch.inference_mode()
     def forward(
@@ -369,7 +421,6 @@ class FishSpeechFastAR(nn.Module):
 
         embed_buf = self._embed_buf
         pos_ids = self._pos_ids
-        model_fwd = self._compiled_model_fwd
 
         # Position 0: projected Slow AR hidden state.
         projected = self.fast_project_in(slow_ar_hidden.reshape(bsz, -1))
@@ -390,9 +441,11 @@ class FishSpeechFastAR(nn.Module):
         for step in range(1, num_cb):
             seq_len = step + 1
             step_input = embed_buf[:bsz, :seq_len, :]
-            step_pos_ids = pos_ids[:seq_len] if bsz == 1 else pos_ids[:seq_len].repeat(bsz)
+            # Use a dense 2D position tensor for every batch size; stride-0
+            # views from expand() were fragile under compiled execution.
+            step_pos_ids = pos_ids[:seq_len].unsqueeze(0).repeat(bsz, 1)
 
-            hidden_out = model_fwd(step_input, step_pos_ids)
+            hidden_out = self._run_model(step_input, step_pos_ids, bsz)
             logits = self.fast_output(self.fast_norm(hidden_out[:, -1, :]))
 
             # Residual codebooks (step >= 1) only have 1024 entries.

vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py

@@ -190,6 +190,7 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
         self.has_postprocess = True
         self.mtp_hidden_size = int(self.text_config.hidden_size)
         self.talker_mtp_output_key = "audio_codes"
+        self.gpu_resident_buffer_keys: set[str] = {"last_slow_ar_hidden"}
 
         # Qwen3 transformer backbone.
         self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
@@ -354,28 +355,31 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
 
         if span_len > 1:
             # --- Prefill ---
-            prompt_embeds_cpu = info_dict.get("slow_ar_prompt_embeds")
-            is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2
+            prompt_embeds_buf = info_dict.get("slow_ar_prompt_embeds")
+            is_first_prefill = not isinstance(prompt_embeds_buf, torch.Tensor) or prompt_embeds_buf.ndim != 2
             dev = input_ids.device
 
             if is_first_prefill:
                 prompt_embeds = self._build_prefill_embeds(input_ids, info_dict)
-                prompt_embeds_cpu = prompt_embeds.detach().to("cpu").contiguous()
+                prompt_embeds_buf = prompt_embeds.detach().to("cpu").contiguous()
+                if not prompt_embeds_buf.is_pinned():
+                    prompt_embeds_buf = prompt_embeds_buf.pin_memory()
+                total_prompt_len = int(prompt_embeds_buf.shape[0])
+                next_offset = min(span_len, total_prompt_len)
 
                 info_update: dict[str, Any] = {
-                    "slow_ar_prompt_embeds": prompt_embeds_cpu,
-                    "prefill_offset": 0,
+                    "slow_ar_prompt_embeds": prompt_embeds_buf if next_offset < total_prompt_len else None,
+                    "prefill_offset": next_offset,
                 }
 
-                take = prompt_embeds_cpu[:span_len]
+                take = prompt_embeds_buf[:span_len]
                 if int(take.shape[0]) < span_len:
                     pad_n = span_len - int(take.shape[0])
                     pad_embed = self.embed_input_ids(
                         torch.tensor([self._audio_pad_token_id], device=dev, dtype=torch.long)
                     ).reshape(1, -1)
-                    take = torch.cat([take, pad_embed.detach().cpu().expand(pad_n, -1)], dim=0)
-                prompt_embeds = take.to(device=dev, dtype=torch.bfloat16)
-                info_update["prefill_offset"] = span_len
+                    take = torch.cat([take, pad_embed.expand(pad_n, -1)], dim=0)
+                prompt_embeds = take.to(device=dev, dtype=torch.bfloat16, non_blocking=True)
 
                 zeros = torch.zeros(
                     (prompt_embeds.shape[0], self._num_codebooks),
@@ -391,23 +395,26 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
             else:
                 # Subsequent prefill chunk.
                 offset = int(info_dict.get("prefill_offset", 0) or 0)
-                s = max(0, min(offset, int(prompt_embeds_cpu.shape[0])))
-                e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0])))
-                take = prompt_embeds_cpu[s:e]
+                total_prompt_len = int(prompt_embeds_buf.shape[0])
+                s = max(0, min(offset, total_prompt_len))
+                e = max(0, min(offset + span_len, total_prompt_len))
+                take = prompt_embeds_buf[s:e]
                 if int(take.shape[0]) < span_len:
                     pad_n = span_len - int(take.shape[0])
                     pad_embed = self.embed_input_ids(
                         torch.tensor([self._audio_pad_token_id], device=dev, dtype=torch.long)
                     ).reshape(1, -1)
-                    take = torch.cat([take, pad_embed.detach().cpu().expand(pad_n, -1)], dim=0)
-                prompt_embeds = take.to(device=dev, dtype=torch.bfloat16)
+                    take = torch.cat([take, pad_embed.expand(pad_n, -1)], dim=0)
+                prompt_embeds = take.to(device=dev, dtype=torch.bfloat16, non_blocking=True)
+                next_offset = offset + span_len
 
                 zeros = torch.zeros((prompt_embeds.shape[0], self._num_codebooks), device=dev, dtype=torch.long)
                 return (
                     input_ids.clone().fill_(self._audio_pad_token_id),
                     prompt_embeds,
                     {
-                        "prefill_offset": offset + span_len,
+                        "slow_ar_prompt_embeds": prompt_embeds_buf if next_offset < total_prompt_len else None,
+                        "prefill_offset": next_offset,
                         "audio_codes": zeros,
                     },
                 )
@@ -415,8 +422,8 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
         # --- Decode: span_len == 1 ---
         dev = input_ids.device
 
-        last_hidden_cpu = info_dict.get("last_slow_ar_hidden")
-        if not isinstance(last_hidden_cpu, torch.Tensor):
+        last_hidden = info_dict.get("last_slow_ar_hidden")
+        if not isinstance(last_hidden, torch.Tensor):
             # First decode step after prefill -- just embed the token directly.
             logger.warning(
                 "preprocess decode: last_slow_ar_hidden not found (keys=%s), "
@@ -437,7 +444,7 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
 
         info_update = {
             "mtp_inputs": (
-                last_hidden_cpu.to(device=dev, dtype=torch.bfloat16).reshape(1, -1),
+                last_hidden.to(device=dev, dtype=torch.bfloat16).reshape(1, -1),
                 torch.zeros(1, self.text_config.hidden_size, device=dev, dtype=torch.bfloat16),
             ),
         }
@@ -447,7 +454,7 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
         if hidden_states.numel() == 0:
             logger.debug("postprocess: empty hidden_states")
             return {}
-        last = hidden_states[-1, :].detach().to("cpu").contiguous()
+        last = hidden_states[-1, :].detach().contiguous()
         logger.debug("postprocess: saved last_slow_ar_hidden shape=%s", tuple(last.shape))
         return {"last_slow_ar_hidden": last}
 
@@ -542,20 +549,19 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
         # This ensures the Slow AR sees codes from FastAR(hidden_{t-1}).
         inputs_embeds_out = input_embeds.reshape(bsz, -1).clone()
 
-        for b in range(bsz):
-            token_id = int(input_ids[b, 0].item())
-            is_semantic = self._semantic_begin_id <= token_id <= self._semantic_end_id
-            if is_semantic:
-                codes = audio_codes[b]  # [num_codebooks]
-                codebook_sum = torch.zeros(self.text_config.hidden_size, device=dev, dtype=torch.bfloat16)
-                for i in range(self._num_codebooks):
-                    code_with_offset = codes[i].clamp(min=0) + i * self._codebook_size
-                    emb = self.codebook_embeddings(code_with_offset.unsqueeze(0))
-                    codebook_sum += emb.squeeze(0).to(dtype=torch.bfloat16)
+        semantic_mask = (input_ids[:, 0] >= self._semantic_begin_id) & (input_ids[:, 0] <= self._semantic_end_id)
+        if semantic_mask.any():
+            semantic_codes = audio_codes[semantic_mask].clamp(min=0)
+            offsets = (
+                torch.arange(self._num_codebooks, device=dev, dtype=semantic_codes.dtype) * self._codebook_size
+            ).unsqueeze(0)
+            codebook_sum = self.codebook_embeddings(semantic_codes + offsets).sum(dim=1).to(dtype=torch.bfloat16)
 
-                # Normalize by sqrt(num_codebooks + 1) as in the reference model
-                # (scale_codebook_embeddings=True for fish_qwen3_omni).
-                inputs_embeds_out[b] = (inputs_embeds_out[b] + codebook_sum) / math.sqrt(self._num_codebooks + 1)
+            # Normalize by sqrt(num_codebooks + 1) as in the reference model
+            # (scale_codebook_embeddings=True for fish_qwen3_omni).
+            inputs_embeds_out[semantic_mask] = (inputs_embeds_out[semantic_mask] + codebook_sum) / math.sqrt(
+                self._num_codebooks + 1
+            )
 
         return inputs_embeds_out, audio_codes.to(dtype=torch.long)
 
@@ -666,4 +672,13 @@ class FishSpeechSlowARForConditionalGeneration(nn.Module):
         if truncated:
             logger.info("Truncated %d RoPE cos_sin_cache buffers to bf16 precision", truncated)
 
+        try:
+            self.fast_ar.warmup_compile(
+                device=self.codebook_embeddings.weight.device,
+                dtype=torch.bfloat16,
+                batch_sizes=(1,),
+            )
+        except Exception as exc:
+            logger.warning("Fish Speech Fast AR compile warmup failed: %s", exc)
+
         return loaded_params

vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml

@@ -4,6 +4,7 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
+      max_batch_size: 16
     engine_args:
       max_num_seqs: 4
       model_stage: fish_speech_slow_ar
@@ -15,10 +16,10 @@ stage_args:
       async_scheduling: false
       enable_prefix_caching: false
       engine_output_type: latent
-      gpu_memory_utilization: 0.4
+      gpu_memory_utilization: 0.6
       distributed_executor_backend: "mp"
-      max_num_batched_tokens: 2048
-      max_model_len: 32768
+      max_num_batched_tokens: 3072
+      max_model_len: 16384
       custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.fish_speech.slow_ar_to_dac_decoder_async_chunk
     output_connectors:
       to_stage_1: connector_of_shared_memory
@@ -37,6 +38,7 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
+      max_batch_size: 16
     engine_args:
       max_num_seqs: 1
       model_stage: dac_decoder
@@ -48,10 +50,10 @@ stage_args:
       async_scheduling: false
       enable_prefix_caching: false
       engine_output_type: audio
-      gpu_memory_utilization: 0.2
+      gpu_memory_utilization: 0.1
       distributed_executor_backend: "mp"
       max_num_batched_tokens: 8192
-      max_model_len: 32768
+      max_model_len: 16384
     engine_input_source: [0]
     final_output: true
     final_output_type: audio
@@ -70,7 +72,7 @@ runtime:
   enabled: true
   defaults:
     window_size: -1
-    max_inflight: 1
+    max_inflight: 16
 
   connectors:
     connector_of_shared_memory:
@@ -85,7 +87,7 @@ runtime:
         # 25 frames ≈ 1.16s of audio at 21.5 Hz.
         codec_chunk_frames: 25
         codec_left_context_frames: 25
-        initial_codec_chunk_frames: 0
+        initial_codec_chunk_frames: 4
 
   edges:
     - from: 0

vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml

@@ -1,46 +0,0 @@
-async_chunk: true
-stage_args:
-- stage_id: 0
-  stage_type: llm
-  runtime: {process: true, devices: '0'}
-  engine_args: {max_num_seqs: 32, model_stage: audio_generation, model_arch: VoxtralTTSForConditionalGeneration,
-    worker_type: ar, worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker, scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler,
-    gpu_memory_utilization: 0.8, enforce_eager: false, trust_remote_code: true, async_scheduling: true,
-    engine_output_type: latent, enable_prefix_caching: false, tokenizer_mode: mistral,
-    config_format: mistral, skip_mm_profiling: true, enable_chunked_prefill: false,
-    max_model_len: 4096, custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.voxtral_tts.generator2tokenizer_async_chunk}
-  output_connectors: {to_stage_1: connector_of_shared_memory}
-  is_comprehension: true
-  final_output: false
-  final_output_type: text
-  default_sampling_params: {temperature: 0.0, top_p: 1.0, top_k: -1, max_tokens: 2048,
-    seed: 42, detokenize: true, repetition_penalty: 1.1}
-- stage_id: 1
-  stage_type: llm
-  runtime: {process: true, devices: '0'}
-  engine_args: {max_num_seqs: 32, model_stage: audio_tokenizer, model_arch: VoxtralTTSForConditionalGeneration,
-    worker_type: generation, worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker,
-    scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler,
-    async_scheduling: false, gpu_memory_utilization: 0.1, enforce_eager: true, trust_remote_code: true,
-    enable_prefix_caching: false, skip_mm_profiling: true, engine_output_type: audio,
-    tokenizer_mode: mistral, config_format: mistral, max_num_batched_tokens: 65536,
-    max_model_len: 65536}
-  engine_input_source: [0]
-  is_comprehension: false
-  final_output: true
-  final_output_type: audio
-  input_connectors: {from_stage_0: connector_of_shared_memory}
-  tts_args: {max_instructions_length: 500}
-  default_sampling_params: {temperature: 0.9, top_p: 0.8, top_k: 40, max_tokens: 2048,
-    seed: 42, detokenize: true, repetition_penalty: 1.05}
-runtime:
-  enabled: true
-  defaults: {window_size: -1, max_inflight: 1}
-  connectors:
-    connector_of_shared_memory:
-      name: SharedMemoryConnector
-      extra: {shm_threshold_bytes: 65536, codec_streaming: true, connector_get_sleep_s: 0.01,
-        connector_get_max_wait_first_chunk: 3000, connector_get_max_wait: 300, codec_chunk_frames: 25,
-        codec_chunk_frames_at_begin: 5, codec_left_context_frames: 25}
-  edges:
-  - {from: 0, to: 1, window_size: -1}

vllm_omni/model_executor/stage_input_processors/fish_speech.py

@@ -7,8 +7,6 @@ from vllm.logger import init_logger
 
 logger = init_logger(__name__)
 
-_NUM_CODEBOOKS = 10  # 1 semantic + 9 residual
-
 
 def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None:
     """Extract the last frame of audio codes from the pooling output."""
@@ -76,8 +74,7 @@ def slow_ar_to_dac_decoder_async_chunk(
     if isinstance(pooling_output, dict):
         frame = _extract_last_frame(pooling_output)
         if frame is not None:
-            codec_codes = frame.cpu().tolist()
-            transfer_manager.code_prompt_token_ids[request_id].append(codec_codes)
+            transfer_manager.code_prompt_token_ids[request_id].append(frame.detach().to(device="cpu", dtype=torch.long))
     elif not finished:
         return None
 
@@ -114,7 +111,7 @@ def slow_ar_to_dac_decoder_async_chunk(
         if finished:
             return {
                 "code_predictor_codes": [],
-                "finished": torch.tensor(True, dtype=torch.bool),
+                "finished": True,
             }
         return None
 
@@ -142,10 +139,11 @@ def slow_ar_to_dac_decoder_async_chunk(
         window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
 
     # Pack into codebook-major flat codes.
-    code_predictor_codes = torch.tensor(window_frames).transpose(0, 1).reshape(-1).tolist()
+    stacked_frames = torch.stack(window_frames, dim=0)
+    code_predictor_codes = stacked_frames.transpose(0, 1).reshape(-1).tolist()
 
     return {
         "code_predictor_codes": code_predictor_codes,
         "left_context_size": left_context_size,
-        "finished": torch.tensor(finished, dtype=torch.bool),
+        "finished": finished,
     }