diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 6d5674bf5..ad71e061f 100644 --- a/.buildkite/test-nightly.yml +++ b/.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 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index ece7dffff..a3b1dbe76 100644 --- a/docs/models/supported_models.md +++ b/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 diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 17787e682..7969ca441 100644 --- a/docs/serving/speech_api.md +++ b/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 --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 diff --git a/examples/offline_inference/voxtral_tts/README.md b/examples/offline_inference/voxtral_tts/README.md index f552fd165..d754b0b38 100644 --- a/examples/offline_inference/voxtral_tts/README.md +++ b/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`) | diff --git a/examples/offline_inference/voxtral_tts/end2end.py b/examples/offline_inference/voxtral_tts/end2end.py index d3052e5ca..f2f2a551b 100644 --- a/examples/offline_inference/voxtral_tts/end2end.py +++ b/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( diff --git a/examples/online_serving/voxtral_tts/gradio_demo.py b/examples/online_serving/voxtral_tts/gradio_demo.py index 293f1c391..465167ba1 100644 --- a/examples/online_serving/voxtral_tts/gradio_demo.py +++ b/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, ) diff --git a/tests/e2e/online_serving/test_voxtral_tts.py b/tests/e2e/online_serving/test_voxtral_tts.py index d4de073f8..f795288f3 100644 --- a/tests/e2e/online_serving/test_voxtral_tts.py +++ b/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" diff --git a/tests/perf/tests/test.json b/tests/perf/tests/test.json index dc76d1dc4..98b2ef25c 100644 --- a/tests/perf/tests/test.json +++ b/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, diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 84bcfa93a..4e13dface 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/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) diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 453792f4c..e190ab6b7 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/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() diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index dc8c5a8ea..49fd0c03a 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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") diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py index a889f66fd..3a8042eb2 100644 --- a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py +++ b/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 diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py index f73914c28..8bbb643eb 100644 --- a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py +++ b/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. diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py index 9c41d87f9..6145815aa 100644 --- a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml b/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml index c00528d10..9b37972f1 100644 --- a/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml +++ b/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 diff --git a/vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml b/vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml deleted file mode 100644 index 287d82895..000000000 --- a/vllm_omni/model_executor/stage_configs/voxtral_tts_1773881731.yaml +++ /dev/null @@ -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} diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py index 5be70bc18..d857c9123 100644 --- a/vllm_omni/model_executor/stage_input_processors/fish_speech.py +++ b/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, }