From 496870295050ba404365e82bbe00cb464618f8d9 Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Tue, 10 Feb 2026 14:37:14 +0800 Subject: [PATCH 01/13] [Chore] Update copyright year. (#1256) Signed-off-by: lishunyang --- mkdocs.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mkdocs.yml b/mkdocs.yml index 15a07450c..d3fc52073 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -8,7 +8,7 @@ repo_url: https://github.com/vllm-project/vllm-omni edit_uri: edit/main/docs/ # Copyright -copyright: Copyright © 2025 vLLM-Omni Team +copyright: Copyright © 2026 vLLM-Omni Team # Theme theme: From f73fa1fbe79484cda72e72d54dce00814e48c5fd Mon Sep 17 00:00:00 2001 From: JohnJan Date: Tue, 10 Feb 2026 14:54:35 +0800 Subject: [PATCH 02/13] [feature]: support Flux.1-dev CFG-Parallel (#1269) --- .../diffusion/parallelism_acceleration.md | 2 +- docs/user_guide/diffusion_acceleration.md | 2 +- .../diffusion/models/flux/pipeline_flux.py | 97 +++++++++++-------- 3 files changed, 57 insertions(+), 44 deletions(-) diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index d536f4bc7..e2dae5c84 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -30,7 +30,7 @@ The following table shows which models are currently supported by parallelism me | **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ | | **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | -| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ | +| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | !!! note "TP Limitations for Diffusion Models" We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP. diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 859f8c0a2..ba122a1cb 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -50,7 +50,7 @@ The following table shows which models are currently supported by each accelerat | **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ | | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | -| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ❌ | +| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ### VideoGen diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py index b90aaa8ca..b67c8eb87 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py @@ -22,6 +22,7 @@ from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokeniz from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -129,9 +130,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class FluxPipeline( - nn.Module, -): +class FluxPipeline(nn.Module, CFGParallelMixin): def __init__( self, *, @@ -500,21 +499,23 @@ class FluxPipeline( def diffuse( self, - prompt_embeds, - pooled_prompt_embeds, - negative_prompt_embeds, - negative_pooled_prompt_embeds, - latents, - latent_image_ids, - text_ids, - negative_text_ids, - timesteps, - do_true_cfg, - guidance, - true_cfg_scale, - ): + prompt_embeds: torch.Tensor, + pooled_prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor, + negative_pooled_prompt_embeds: torch.Tensor, + latents: torch.Tensor, + latent_image_ids: torch.Tensor, + text_ids: torch.Tensor, + negative_text_ids: torch.Tensor, + timesteps: torch.Tensor, + do_true_cfg: bool, + guidance: torch.Tensor, + true_cfg_scale: float, + cfg_normalize: bool = False, + ) -> torch.Tensor: """Diffusion loop with optional image conditioning.""" self.scheduler.set_begin_index(0) + self.transformer.do_true_cfg = do_true_cfg for i, t in enumerate(timesteps): if self.interrupt: continue @@ -523,36 +524,46 @@ class FluxPipeline( # broadcast to batch dimension and place on same device/dtype as latents timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype) - self.transformer.do_true_cfg = do_true_cfg # used in teacache hook - # Forward pass for positive prompt (or unconditional if no CFG) - noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] + positive_kwargs = { + "hidden_states": latents, + "timestep": timestep / 1000, + "guidance": guidance, + "pooled_projections": pooled_prompt_embeds, + "encoder_hidden_states": prompt_embeds, + "txt_ids": text_ids, + "img_ids": latent_image_ids, + "joint_attention_kwargs": self.joint_attention_kwargs, + "return_dict": False, + } # Forward pass for negative prompt (CFG) if do_true_cfg: - neg_noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=negative_pooled_prompt_embeds, - encoder_hidden_states=negative_prompt_embeds, - txt_ids=negative_text_ids, - img_ids=latent_image_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) + negative_kwargs = { + "hidden_states": latents, + "timestep": timestep / 1000, + "guidance": guidance, + "pooled_projections": negative_pooled_prompt_embeds, + "encoder_hidden_states": negative_prompt_embeds, + "txt_ids": negative_text_ids, + "img_ids": latent_image_ids, + "joint_attention_kwargs": self.joint_attention_kwargs, + "return_dict": False, + } + else: + negative_kwargs = None + + # Predict noise with automatic CFG parallel handling + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg, + true_cfg_scale, + positive_kwargs, + negative_kwargs, + cfg_normalize, + ) + # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + return latents def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool): @@ -613,6 +624,7 @@ class FluxPipeline( req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale ) generator = req.sampling_params.generator or generator + true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale num_images_per_prompt = ( req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 @@ -723,6 +735,7 @@ class FluxPipeline( do_true_cfg, guidance, true_cfg_scale, + cfg_normalize=False, ) self._current_timestep = None From 9327cccde615374031c44adcf9833951b0a0254d Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Tue, 10 Feb 2026 15:25:08 +0800 Subject: [PATCH 03/13] [Bugfix] Fix 'NoneType' AttributeError in stable-diffusion model detect (#1254) Signed-off-by: Yan Ma --- vllm_omni/diffusion/utils/hf_utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/vllm_omni/diffusion/utils/hf_utils.py b/vllm_omni/diffusion/utils/hf_utils.py index cfc1807a1..7935df1e4 100644 --- a/vllm_omni/diffusion/utils/hf_utils.py +++ b/vllm_omni/diffusion/utils/hf_utils.py @@ -18,13 +18,13 @@ def _looks_like_bagel(model_name: str) -> bool: """Best-effort detection for Bagel (non-diffusers) diffusion models.""" try: cfg = get_hf_file_to_dict("config.json", model_name) + model_type = cfg.get("model_type") + if model_type == "bagel": + return True + architectures = cfg.get("architectures") or [] + return "BagelForConditionalGeneration" in architectures except Exception: return False - model_type = cfg.get("model_type") - if model_type == "bagel": - return True - architectures = cfg.get("architectures") or [] - return "BagelForConditionalGeneration" in architectures @lru_cache From acfe99e4d9e75caa1c0dae6704aabbbb22f5e397 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Tue, 10 Feb 2026 02:52:44 -0500 Subject: [PATCH 04/13] [Doc] Update Qwen3-TTS docs for consistency with Omni examples (#1226) Signed-off-by: linyueqian Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- .../examples/online_serving/qwen3_tts.md | 206 +++++++++++++----- examples/online_serving/qwen3_tts/README.md | 98 ++++++--- .../qwen3_tts/openai_speech_client.py | 10 +- .../online_serving/qwen3_tts/run_server.sh | 2 +- 4 files changed, 230 insertions(+), 86 deletions(-) diff --git a/docs/user_guide/examples/online_serving/qwen3_tts.md b/docs/user_guide/examples/online_serving/qwen3_tts.md index 48bd58b00..f899e362e 100644 --- a/docs/user_guide/examples/online_serving/qwen3_tts.md +++ b/docs/user_guide/examples/online_serving/qwen3_tts.md @@ -1,9 +1,11 @@ -# Qwen3-TTS Online Serving +# Qwen3-TTS Source . -This directory contains examples for running Qwen3-TTS models with vLLM-Omni's online serving API. +## 🛠️ Installation + +Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md) ## Supported Models @@ -12,34 +14,77 @@ This directory contains examples for running Qwen3-TTS models with vLLM-Omni's o | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control | | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description | | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio | +| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant | +| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning | -## Quick Start +## Run examples (Qwen3-TTS) -### 1. Start the Server +### Launch the Server ```bash -# CustomVoice model (default) -./run_server.sh +# CustomVoice model (predefined speakers) +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager -# Or specify task type -./run_server.sh CustomVoice -./run_server.sh VoiceDesign -./run_server.sh Base +# VoiceDesign model +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \ + --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager + +# Base model (voice cloning) +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \ + --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager ``` -### 2. Run the Client +If you have custom stage configs file, launch the server with command below +```bash +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + --stage-configs-path /path/to/stage_configs_file \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager +``` + +Alternatively, use the convenience script: +```bash +./run_server.sh # Default: CustomVoice model +./run_server.sh CustomVoice # CustomVoice model +./run_server.sh VoiceDesign # VoiceDesign model +./run_server.sh Base # Base (voice clone) model +``` + +### Send TTS Request + +Get into the example folder +```bash +cd examples/online_serving/qwen3_tts +``` + +#### Send request via python ```bash # CustomVoice: Use predefined speaker python openai_speech_client.py \ --text "你好，我是通义千问" \ - --voice Vivian \ + --voice vivian \ --language Chinese # CustomVoice with style instruction python openai_speech_client.py \ --text "今天天气真好" \ - --voice Ryan \ + --voice ryan \ --instructions "用开心的语气说" # VoiceDesign: Describe the voice style @@ -58,29 +103,86 @@ python openai_speech_client.py \ --ref-text "Original transcript of the reference audio" ``` -### 3. Using curl +The Python client supports the following command-line arguments: + +- `--api-base`: API base URL (default: `http://localhost:8091`) +- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice`) +- `--task-type` (or `-t`): TTS task type. Options: `CustomVoice`, `VoiceDesign`, `Base` +- `--text`: Text to synthesize (required) +- `--voice`: Speaker/voice name (default: `vivian`). Options: `vivian`, `ryan`, `aiden`, etc. +- `--language`: Language. Options: `Auto`, `Chinese`, `English`, `Japanese`, `Korean`, `German`, `French`, `Russian`, `Portuguese`, `Spanish`, `Italian` +- `--instructions`: Voice style/emotion instructions +- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task) +- `--ref-text`: Reference audio transcript for voice cloning (Base task) +- `--response-format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`, `aac`, `opus` +- `--output` (or `-o`): Output audio file path (default: `tts_output.wav`) + +#### Send request via curl ```bash # Simple TTS request -curl -X POST http://localhost:8000/v1/audio/speech \ +curl -X POST http://localhost:8091/v1/audio/speech \ -H "Content-Type: application/json" \ -d '{ "input": "Hello, how are you?", - "voice": "Vivian", + "voice": "vivian", "language": "English" }' --output output.wav # With style instruction -curl -X POST http://localhost:8000/v1/audio/speech \ +curl -X POST http://localhost:8091/v1/audio/speech \ -H "Content-Type: application/json" \ -d '{ "input": "I am so excited!", - "voice": "Vivian", + "voice": "vivian", "instructions": "Speak with great enthusiasm" }' --output excited.wav # List available voices in CustomVoice models -curl http://localhost:8000/v1/audio/voices +curl http://localhost:8091/v1/audio/voices +``` + +### Using OpenAI SDK + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="none") + +response = client.audio.speech.create( + model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", + voice="vivian", + input="Hello, how are you?", +) + +response.stream_to_file("output.wav") +``` + +### Using Python httpx + +```python +import httpx + +response = httpx.post( + "http://localhost:8091/v1/audio/speech", + json={ + "input": "Hello, how are you?", + "voice": "vivian", + "language": "English", + }, + timeout=300.0, +) + +with open("output.wav", "wb") as f: + f.write(response.content) +``` + +### FAQ + +If you encounter error about backend of librosa, try to install ffmpeg with command below. +``` +sudo apt update +sudo apt install ffmpeg ``` ## API Reference @@ -89,16 +191,31 @@ curl http://localhost:8000/v1/audio/voices ``` POST /v1/audio/speech +Content-Type: application/json ``` This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters. +### Voices Endpoint + +``` +GET /v1/audio/voices +``` + +Lists available voices for the loaded model: + +```json +{ + "voices": ["aiden", "dylan", "eric", "one_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"] +} +``` + ### Request Body ```json { "input": "Text to synthesize", - "voice": "Vivian", + "voice": "vivian", "response_format": "wav", "task_type": "CustomVoice", "language": "Auto", @@ -114,56 +231,38 @@ This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/ ### Response -Returns audio data in the requested format (default: WAV). +Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`). ## Parameters -### Standard OpenAI Parameters +### OpenAI Standard Parameters | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `input` | string | required | Text to synthesize | -| `voice` | string | "Vivian" | Speaker/voice name | +| `input` | string | **required** | Text to synthesize | +| `model` | string | server's model | Model to use (optional, should match server if specified) | +| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) | | `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus | | `speed` | float | 1.0 | Playback speed (0.25-4.0) | -| `model` | string | optional | Model name (optional when serving single model) | -### Qwen3-TTS Parameters +### vLLM-Omni Extension Parameters | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base | -| `language` | string | "Auto" | Language: Auto, Chinese, English, Japanese, Korean | +| `language` | string | "Auto" | Language (see supported languages below) | | `instructions` | string | "" | Voice style/emotion instructions | | `max_new_tokens` | int | 2048 | Maximum tokens to generate | +**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian + ### Voice Clone Parameters (Base task) | Parameter | Type | Required | Description | |-----------|------|----------|-------------| -| `ref_audio` | string | Yes* | Reference audio (file path, URL, or base64) | +| `ref_audio` | string | **Yes** | Reference audio (URL or base64 data URL) | | `ref_text` | string | No | Transcript of reference audio (for ICL mode) | -| `x_vector_only_mode` | bool | false | Use speaker embedding only (no ICL) | - -## Python Usage - -```python -import httpx - -# Simple request -response = httpx.post( - "http://localhost:8000/v1/audio/speech", - json={ - "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", - "input": "Hello world", - "voice": "Vivian", - }, - timeout=300.0, -) - -with open("output.wav", "wb") as f: - f.write(response.content) -``` +| `x_vector_only_mode` | bool | No | Use speaker embedding only (no ICL) | ## Limitations @@ -172,10 +271,11 @@ with open("output.wav", "wb") as f: ## Troubleshooting -1. **Connection refused**: Make sure the server is running on the correct port -2. **Out of memory**: Reduce `--gpu-memory-utilization` in run_server.sh -3. **Unsupported speaker**: Check supported speakers via model documentation -4. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning +1. **TTS model did not produce audio output**: Ensure you're using the correct model variant for your task type (CustomVoice task → CustomVoice model, etc.) +2. **Connection refused**: Make sure the server is running on the correct port +3. **Out of memory**: Use smaller model variant (`Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`) or reduce `--gpu-memory-utilization` +4. **Unsupported speaker**: Use `/v1/audio/voices` to list available voices for the loaded model +5. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning ## Example materials diff --git a/examples/online_serving/qwen3_tts/README.md b/examples/online_serving/qwen3_tts/README.md index 1c9bd4820..e1ad3edc2 100644 --- a/examples/online_serving/qwen3_tts/README.md +++ b/examples/online_serving/qwen3_tts/README.md @@ -9,6 +9,8 @@ This directory contains examples for running Qwen3-TTS models with vLLM-Omni's o | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control | | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description | | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio | +| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant | +| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning | ## Quick Start @@ -24,19 +26,27 @@ This directory contains examples for running Qwen3-TTS models with vLLM-Omni's o ./run_server.sh Base ``` +Or launch directly with vllm serve: + +```bash +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ + --omni --port 8091 --trust-remote-code --enforce-eager +``` + ### 2. Run the Client ```bash # CustomVoice: Use predefined speaker python openai_speech_client.py \ --text "你好，我是通义千问" \ - --voice Vivian \ + --voice vivian \ --language Chinese # CustomVoice with style instruction python openai_speech_client.py \ --text "今天天气真好" \ - --voice Ryan \ + --voice ryan \ --instructions "用开心的语气说" # VoiceDesign: Describe the voice style @@ -59,25 +69,25 @@ python openai_speech_client.py \ ```bash # Simple TTS request -curl -X POST http://localhost:8000/v1/audio/speech \ +curl -X POST http://localhost:8091/v1/audio/speech \ -H "Content-Type: application/json" \ -d '{ "input": "Hello, how are you?", - "voice": "Vivian", + "voice": "vivian", "language": "English" }' --output output.wav # With style instruction -curl -X POST http://localhost:8000/v1/audio/speech \ +curl -X POST http://localhost:8091/v1/audio/speech \ -H "Content-Type: application/json" \ -d '{ "input": "I am so excited!", - "voice": "Vivian", + "voice": "vivian", "instructions": "Speak with great enthusiasm" }' --output excited.wav # List available voices in CustomVoice models -curl http://localhost:8000/v1/audio/voices +curl http://localhost:8091/v1/audio/voices ``` ## API Reference @@ -90,12 +100,26 @@ POST /v1/audio/speech This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters. +### Voices Endpoint + +``` +GET /v1/audio/voices +``` + +Lists available voices for the loaded model: + +```json +{ + "voices": ["aiden", "dylan", "eric", "one_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"] +} +``` + ### Request Body ```json { "input": "Text to synthesize", - "voice": "Vivian", + "voice": "vivian", "response_format": "wav", "task_type": "CustomVoice", "language": "Auto", @@ -115,45 +139,64 @@ Returns audio data in the requested format (default: WAV). ## Parameters -### Standard OpenAI Parameters +### OpenAI Standard Parameters | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `input` | string | required | Text to synthesize | -| `voice` | string | "Vivian" | Speaker/voice name | +| `input` | string | **required** | Text to synthesize | +| `model` | string | server's model | Model to use (optional, should match server if specified) | +| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) | | `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus | | `speed` | float | 1.0 | Playback speed (0.25-4.0) | -| `model` | string | optional | Model name (optional when serving single model) | -### Qwen3-TTS Parameters +### vLLM-Omni Extension Parameters | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base | -| `language` | string | "Auto" | Language: Auto, Chinese, English, Japanese, Korean | +| `language` | string | "Auto" | Language (see supported languages below) | | `instructions` | string | "" | Voice style/emotion instructions | | `max_new_tokens` | int | 2048 | Maximum tokens to generate | +**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian + ### Voice Clone Parameters (Base task) -| Parameter | Type | Required | Description | +| Parameter | Type | Default | Description | |-----------|------|----------|-------------| -| `ref_audio` | string | Yes* | Reference audio (file path, URL, or base64) | -| `ref_text` | string | No | Transcript of reference audio (for ICL mode) | -| `x_vector_only_mode` | bool | false | Use speaker embedding only (no ICL) | +| `ref_audio` | string | null | Reference audio (URL or base64 data URL) | +| `ref_text` | string | null | Transcript of reference audio | +| `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) | ## Python Usage +### Using OpenAI SDK + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="none") + +response = client.audio.speech.create( + model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", + voice="vivian", + input="Hello, how are you?", +) + +response.stream_to_file("output.wav") +``` + +### Using httpx + ```python import httpx -# Simple request response = httpx.post( - "http://localhost:8000/v1/audio/speech", + "http://localhost:8091/v1/audio/speech", json={ - "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", - "input": "Hello world", - "voice": "Vivian", + "input": "Hello, how are you?", + "voice": "vivian", + "language": "English", }, timeout=300.0, ) @@ -169,7 +212,8 @@ with open("output.wav", "wb") as f: ## Troubleshooting -1. **Connection refused**: Make sure the server is running on the correct port -2. **Out of memory**: Reduce `--gpu-memory-utilization` in run_server.sh -3. **Unsupported speaker**: Check supported speakers via model documentation -4. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning +1. **"TTS model did not produce audio output"**: Ensure you're using the correct model variant for your task type (CustomVoice task → CustomVoice model, etc.) +2. **Connection refused**: Make sure the server is running on the correct port +3. **Out of memory**: Use smaller model variant (`Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`) or reduce `--gpu-memory-utilization` +4. **Unsupported speaker**: Use `/v1/audio/voices` to list available voices for the loaded model +5. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning diff --git a/examples/online_serving/qwen3_tts/openai_speech_client.py b/examples/online_serving/qwen3_tts/openai_speech_client.py index ea45a7482..e008eb139 100644 --- a/examples/online_serving/qwen3_tts/openai_speech_client.py +++ b/examples/online_serving/qwen3_tts/openai_speech_client.py @@ -5,10 +5,10 @@ to generate audio from text using Qwen3-TTS models. Examples: # CustomVoice task (predefined speaker) - python openai_speech_client.py --text "Hello, how are you?" --voice Vivian + python openai_speech_client.py --text "Hello, how are you?" --voice vivian # CustomVoice with emotion instruction - python openai_speech_client.py --text "I'm so happy!" --voice Vivian \ + python openai_speech_client.py --text "I'm so happy!" --voice vivian \ --instructions "Speak with excitement" # VoiceDesign task (voice from description) @@ -30,7 +30,7 @@ import os import httpx # Default server configuration -DEFAULT_API_BASE = "http://localhost:8000" +DEFAULT_API_BASE = "http://localhost:8091" DEFAULT_API_KEY = "EMPTY" @@ -178,8 +178,8 @@ def parse_args(): parser.add_argument( "--voice", type=str, - default="Vivian", - help="Speaker/voice name (default: Vivian). Options: Vivian, Ryan, etc.", + default="vivian", + help="Speaker/voice name (default: vivian). Options: vivian, ryan, aiden, etc.", ) parser.add_argument( "--language", diff --git a/examples/online_serving/qwen3_tts/run_server.sh b/examples/online_serving/qwen3_tts/run_server.sh index 203ed76e8..fd020f18b 100755 --- a/examples/online_serving/qwen3_tts/run_server.sh +++ b/examples/online_serving/qwen3_tts/run_server.sh @@ -33,7 +33,7 @@ echo "Starting Qwen3-TTS server with model: $MODEL" vllm-omni serve "$MODEL" \ --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ --host 0.0.0.0 \ - --port 8000 \ + --port 8091 \ --gpu-memory-utilization 0.9 \ --trust-remote-code \ --enforce-eager \ From 109bb97b2c99a7f68c76a06ea69c0b8d9a733661 Mon Sep 17 00:00:00 2001 From: zhou zhuoxin Date: Tue, 10 Feb 2026 16:05:12 +0800 Subject: [PATCH 05/13] [Fix]Ensure HuggingFace downloads complete before initialization. (#1213) Signed-off-by: zhou zhuoxin Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- vllm_omni/entrypoints/omni.py | 22 +++++++++++++-- .../model_loader/weight_utils.py | 27 ++++++++++++++----- 2 files changed, 41 insertions(+), 8 deletions(-) diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 9e49f951e..510813c6a 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -10,6 +10,7 @@ from collections.abc import Callable, Generator, Sequence from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Literal, overload +import huggingface_hub from omegaconf import OmegaConf from tqdm.auto import tqdm from vllm import SamplingParams @@ -40,6 +41,9 @@ from vllm_omni.entrypoints.utils import ( ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams from vllm_omni.metrics import OrchestratorAggregator, StageRequestStats +from vllm_omni.model_executor.model_loader.weight_utils import ( + download_weights_from_hf_specific, +) from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -66,14 +70,28 @@ def _dummy_snapshot_download(model_id): def omni_snapshot_download(model_id) -> str: + # If it's already a local path, just return it + if os.path.exists(model_id): + return model_id # TODO: this is just a workaround for quickly use modelscope, we should support # modelscope in weight loading feature instead of using `snapshot_download` if os.environ.get("VLLM_USE_MODELSCOPE", False): from modelscope.hub.snapshot_download import snapshot_download return snapshot_download(model_id) - else: - return _dummy_snapshot_download(model_id) + # For other cases (Hugging Face), perform a real download to ensure all + # necessary files (including *.pt for audio/diffusion) are available locally + # before stage workers are spawned. This prevents initialization timeouts. + try: + return download_weights_from_hf_specific( + model_name_or_path=model_id, + cache_dir=None, + allow_patterns=["*"], + require_all=True, + ) + except huggingface_hub.errors.RepositoryNotFoundError: + logger.warning(f"Repository not found for '{model_id}'.") + return model_id class OmniBase: diff --git a/vllm_omni/model_executor/model_loader/weight_utils.py b/vllm_omni/model_executor/model_loader/weight_utils.py index 7432ad9a2..c5225ec02 100644 --- a/vllm_omni/model_executor/model_loader/weight_utils.py +++ b/vllm_omni/model_executor/model_loader/weight_utils.py @@ -20,6 +20,7 @@ def download_weights_from_hf_specific( allow_patterns: list[str], revision: str | None = None, ignore_patterns: str | list[str] | None = None, + require_all: bool = False, ) -> str: """Download model weights from Hugging Face Hub. Users can specify the allow_patterns to download only the necessary weights. @@ -35,6 +36,9 @@ def download_weights_from_hf_specific( ignore_patterns (Optional[Union[str, list[str]]]): The patterns to filter out the weight files. Files matched by any of the patterns will be ignored. + require_all (bool): If True, will iterate through and download files + matching all patterns in allow_patterns. If False, will stop after + the first pattern that matches any files. Returns: str: The path to the downloaded model weights. @@ -48,20 +52,31 @@ def download_weights_from_hf_specific( # downloading the same model weights at the same time. with get_lock(model_name_or_path, cache_dir): start_time = time.perf_counter() - for allow_pattern in allow_patterns: + if require_all: hf_folder = snapshot_download( model_name_or_path, - allow_patterns=allow_pattern, + allow_patterns=allow_patterns, ignore_patterns=ignore_patterns, cache_dir=cache_dir, revision=revision, local_files_only=local_only, **download_kwargs, ) - # If we have downloaded weights for this allow_pattern, - # we don't need to check the rest. - if any(Path(hf_folder).glob(allow_pattern)): - break + else: + for allow_pattern in allow_patterns: + hf_folder = snapshot_download( + model_name_or_path, + allow_patterns=allow_pattern, + ignore_patterns=ignore_patterns, + cache_dir=cache_dir, + revision=revision, + local_files_only=local_only, + **download_kwargs, + ) + # If we have downloaded weights for this allow_pattern, + # we don't need to check the rest, unless require_all is set. + if any(Path(hf_folder).glob(allow_pattern)): + break time_taken = time.perf_counter() - start_time if time_taken > 0.5: logger.info( From ee624d0686d0d333d865cd9a35e30fcd6b6bbaa0 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Tue, 10 Feb 2026 16:08:20 +0800 Subject: [PATCH 06/13] [BugFix] Fixed the issue where ignore_eos was not working. (#1286) Signed-off-by: amy-why-3459 --- .../test_serving_chat_sampling_params.py | 22 +++++++------------ vllm_omni/entrypoints/openai/serving_chat.py | 4 ++++ 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py index 58d4d253f..a015b9b5f 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py +++ b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py @@ -86,9 +86,13 @@ def mock_request(): # OpenAI standard sampling fields request.temperature = None request.top_p = None + request.top_k = None request.max_tokens = None + request.min_tokens = None request.seed = None + request.ignore_eos = None request.stop = None + request.stop_token_ids = None request.frequency_penalty = None request.presence_penalty = None return request @@ -106,9 +110,13 @@ def test_openai_sampling_fields_contains_expected_fields(): expected_fields = { "temperature", "top_p", + "top_k", "max_tokens", + "min_tokens", "seed", + "ignore_eos", "stop", + "stop_token_ids", "frequency_penalty", "presence_penalty", } @@ -236,20 +244,6 @@ def test_multiple_params_override_together(serving_chat, mock_request): assert comprehension_params.repetition_penalty == 1.05 -def test_yaml_custom_params_not_overridden_by_request(serving_chat, mock_request): - """Test that YAML custom params (top_k, repetition_penalty) are not affected.""" - # Even if request has these attributes, they should not override YAML - # because they're not in _OPENAI_SAMPLING_FIELDS - mock_request.top_k = 100 # Not in allowlist - mock_request.repetition_penalty = 2.0 # Not in allowlist - - result = serving_chat._build_sampling_params_list_from_request(mock_request) - - comprehension_params = result[0] - assert comprehension_params.top_k == 1 # YAML default preserved - assert comprehension_params.repetition_penalty == 1.05 # YAML default preserved - - # ============================================================================= # Tests for _apply_request_overrides # ============================================================================= diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index cfee79157..c025c5f06 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -489,9 +489,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): _OPENAI_SAMPLING_FIELDS: set[str] = { "temperature", "top_p", + "top_k", "max_tokens", + "min_tokens", "seed", + "ignore_eos", "stop", + "stop_token_ids", "frequency_penalty", "presence_penalty", } From f55ce9c79338bf2446ae1b0193eff8b45b187829 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Tue, 10 Feb 2026 03:33:39 -0500 Subject: [PATCH 07/13] [Test] Add e2e tests for Qwen3-TTS speech endpoint (#1206) Signed-off-by: linyueqian Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> --- .buildkite/pipeline.yml | 20 ++ docker/Dockerfile.ci | 2 +- docs/.nav.yml | 1 + docs/serving/speech_api.md | 267 +++++++++++++++++++++ pyproject.toml | 2 +- tests/e2e/online_serving/test_qwen3_tts.py | 226 +++++++++++++++++ vllm_omni/metrics/stats.py | 45 ++-- 7 files changed, 543 insertions(+), 20 deletions(-) create mode 100644 docs/serving/speech_api.md create mode 100644 tests/e2e/online_serving/test_qwen3_tts.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 01897db96..75cc818a6 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -228,6 +228,26 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Qwen3-TTS E2E Test" + timeout_in_minutes: 10 + depends_on: image-build + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py + agents: + queue: "gpu_4_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + # - label: "Omni Model Test with H100" # timeout_in_minutes: 30 # depends_on: image-build diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index c5d847345..c927e139a 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -8,7 +8,7 @@ COPY . . # Install system dependencies RUN apt-get update && \ - apt-get install -y ffmpeg && \ + apt-get install -y ffmpeg sox libsox-fmt-all && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* diff --git a/docs/.nav.yml b/docs/.nav.yml index 16821c595..e0eb6d674 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -8,6 +8,7 @@ nav: - OpenAI-Compatible API: - Image Generation: serving/image_generation_api.md - Image Edit: serving/image_edit_api.md + - Text to Speech: serving/speech_api.md - Examples: - examples/README.md - Offline Inference: diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md new file mode 100644 index 000000000..86b7c73f8 --- /dev/null +++ b/docs/serving/speech_api.md @@ -0,0 +1,267 @@ +# Speech API + +vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation using Qwen3-TTS models. + +Each server instance runs a single model (specified at startup via `vllm serve --omni`). + +## Quick Start + +### Start the Server + +```bash +# CustomVoice model (predefined speakers) +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager +``` + +### Generate Speech + +**Using curl:** + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "vivian", + "language": "English" + }' --output output.wav +``` + +**Using Python:** + +```python +import httpx + +response = httpx.post( + "http://localhost:8091/v1/audio/speech", + json={ + "input": "Hello, how are you?", + "voice": "vivian", + "language": "English", + }, + timeout=300.0, +) + +with open("output.wav", "wb") as f: + f.write(response.content) +``` + +**Using OpenAI SDK:** + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="none") + +response = client.audio.speech.create( + model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", + voice="vivian", + input="Hello, how are you?", +) + +response.stream_to_file("output.wav") +``` + +## API Reference + +### Endpoint + +``` +POST /v1/audio/speech +Content-Type: application/json +``` + +### Request Parameters + +#### OpenAI Standard Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `input` | string | **required** | The text to synthesize into speech | +| `model` | string | server's model | Model to use (optional, should match server if specified) | +| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) | +| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus | +| `speed` | float | 1.0 | Playback speed (0.25-4.0) | + +#### vLLM-Omni Extension Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `task_type` | string | "CustomVoice" | TTS task type: CustomVoice, VoiceDesign, or Base | +| `language` | string | "Auto" | Language (see supported languages below) | +| `instructions` | string | "" | Voice style/emotion instructions | +| `max_new_tokens` | integer | 2048 | Maximum tokens to generate | + +**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian + +#### Voice Clone Parameters (Base task) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `ref_audio` | string | null | Reference audio (URL or base64 data URL) | +| `ref_text` | string | null | Transcript of reference audio | +| `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) | + +### Response Format + +Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`). + +### Voices Endpoint + +``` +GET /v1/audio/voices +``` + +Lists available voices for the loaded model. + +```json +{ + "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"] +} +``` + +## Examples + +### CustomVoice with Style Instruction + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "I am so excited!", + "voice": "vivian", + "instructions": "Speak with great enthusiasm" + }' --output excited.wav +``` + +### VoiceDesign (Natural Language Voice Description) + +```bash +# Start server with VoiceDesign model first +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \ + --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager +``` + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello world", + "task_type": "VoiceDesign", + "instructions": "A warm, friendly female voice with a gentle tone" + }' --output designed.wav +``` + +### Base (Voice Cloning) + +```bash +# Start server with Base model first +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \ + --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager +``` + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is a cloned voice", + "task_type": "Base", + "ref_audio": "https://example.com/reference.wav", + "ref_text": "Original transcript of the reference audio" + }' --output cloned.wav +``` + +## Supported Models + +| Model | Task Type | Description | +|-------|-----------|-------------| +| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control | +| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description | +| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio | +| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant | +| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning | + +## Error Responses + +### 400 Bad Request + +Invalid parameters: + +```json +{ + "error": { + "message": "Input text cannot be empty", + "type": "BadRequestError", + "param": null, + "code": 400 + } +} +``` + +### 404 Not Found + +Model not found: + +```json +{ + "error": { + "message": "The model `xxx` does not exist.", + "type": "NotFoundError", + "param": "model", + "code": 404 + } +} +``` + +## Troubleshooting + +### "TTS model did not produce audio output" + +Ensure you're using the correct model variant for your task type: +- CustomVoice task → CustomVoice model +- VoiceDesign task → VoiceDesign model +- Base task → Base model + +### Server Not Running + +```bash +# Check if server is responding +curl http://localhost:8091/v1/audio/voices +``` + +### Out of Memory + +If you encounter OOM errors: +1. Use smaller model variant: `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` +2. Reduce `--gpu-memory-utilization` + +### Unsupported Speaker + +Use `/v1/audio/voices` to list available voices for the loaded model. + +## Development + +Enable debug logging: + +```bash +vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager \ + --uvicorn-log-level debug +``` diff --git a/pyproject.toml b/pyproject.toml index 0748b7eab..6af193953 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -192,6 +192,6 @@ extend-ignore-identifiers-re = [ ".*nothink.*", ".*NOTHINK.*", ".*nin.*", - "Ono_Anna", + ".*[Oo]no_[Aa]nna.*", ".*cann.*", ] diff --git a/tests/e2e/online_serving/test_qwen3_tts.py b/tests/e2e/online_serving/test_qwen3_tts.py new file mode 100644 index 000000000..9a1ff8b97 --- /dev/null +++ b/tests/e2e/online_serving/test_qwen3_tts.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E Online tests for Qwen3-TTS model with text input and audio output. + +These tests verify the /v1/audio/speech endpoint works correctly with +actual model inference, not mocks. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +from pathlib import Path + +import httpx +import pytest + +from tests.conftest import OmniServer +from tests.utils import hardware_test + +MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" + + +def get_stage_config(): + """Get the stage config path for Qwen3-TTS.""" + return str( + Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "qwen3_tts.yaml" + ) + + +@pytest.fixture(scope="module") +def omni_server(): + """Start vLLM-Omni server with CustomVoice model.""" + stage_config_path = get_stage_config() + + print(f"Starting OmniServer with model: {MODEL}") + + with OmniServer( + MODEL, + [ + "--stage-configs-path", + stage_config_path, + "--stage-init-timeout", + "120", + "--trust-remote-code", + "--enforce-eager", + ], + ) as server: + print("OmniServer started successfully") + yield server + print("OmniServer stopping...") + + print("OmniServer stopped") + + +def make_speech_request( + host: str, + port: int, + text: str, + voice: str = "vivian", + language: str = "English", + task_type: str | None = None, + instructions: str | None = None, + timeout: float = 120.0, +) -> httpx.Response: + """Make a request to the /v1/audio/speech endpoint.""" + url = f"http://{host}:{port}/v1/audio/speech" + payload = { + "input": text, + "voice": voice, + "language": language, + } + if task_type: + payload["task_type"] = task_type + if instructions: + payload["instructions"] = instructions + + with httpx.Client(timeout=timeout) as client: + return client.post(url, json=payload) + + +def verify_wav_audio(content: bytes) -> bool: + """Verify that content is valid WAV audio data.""" + # WAV files start with "RIFF" header + if len(content) < 44: # Minimum WAV header size + return False + return content[:4] == b"RIFF" and content[8:12] == b"WAVE" + + +# Minimum expected audio size for a short sentence (~1 second of 24kHz 16-bit mono WAV) +MIN_AUDIO_BYTES = 10000 + + +class TestQwen3TTSCustomVoice: + """E2E tests for Qwen3-TTS CustomVoice model.""" + + @pytest.mark.core_model + @pytest.mark.omni + @hardware_test(res={"cuda": "L4"}, num_cards=4) + def test_speech_english_basic(self, omni_server) -> None: + """Test basic English TTS generation.""" + response = make_speech_request( + host=omni_server.host, + port=omni_server.port, + text="Hello, how are you?", + voice="vivian", + language="English", + ) + + assert response.status_code == 200, f"Request failed: {response.text}" + assert response.headers.get("content-type") == "audio/wav" + assert verify_wav_audio(response.content), "Response is not valid WAV audio" + assert len(response.content) > MIN_AUDIO_BYTES, ( + f"Audio content too small ({len(response.content)} bytes), expected at least {MIN_AUDIO_BYTES} bytes" + ) + + @pytest.mark.core_model + @pytest.mark.omni + @hardware_test(res={"cuda": "L4"}, num_cards=4) + def test_speech_chinese_basic(self, omni_server) -> None: + """Test basic Chinese TTS generation.""" + response = make_speech_request( + host=omni_server.host, + port=omni_server.port, + text="你好，我是通义千问", + voice="vivian", + language="Chinese", + ) + + assert response.status_code == 200, f"Request failed: {response.text}" + assert response.headers.get("content-type") == "audio/wav" + assert verify_wav_audio(response.content), "Response is not valid WAV audio" + assert len(response.content) > MIN_AUDIO_BYTES, ( + f"Audio content too small ({len(response.content)} bytes), expected at least {MIN_AUDIO_BYTES} bytes" + ) + + @pytest.mark.core_model + @pytest.mark.omni + @hardware_test(res={"cuda": "L4"}, num_cards=4) + def test_speech_different_voices(self, omni_server) -> None: + """Test TTS with different voice options.""" + voices = ["vivian", "ryan"] + for voice in voices: + response = make_speech_request( + host=omni_server.host, + port=omni_server.port, + text="Testing voice selection.", + voice=voice, + language="English", + ) + + assert response.status_code == 200, f"Request failed for voice {voice}: {response.text}" + assert verify_wav_audio(response.content), f"Invalid WAV for voice {voice}" + + @pytest.mark.core_model + @pytest.mark.omni + @hardware_test(res={"cuda": "L4"}, num_cards=4) + def test_speech_binary_response_not_utf8_error(self, omni_server) -> None: + """ + Regression test: Verify binary audio is returned, not UTF-8 error. + + This test ensures the multimodal_output property correctly retrieves + audio from completion outputs, preventing the "TTS model did not + produce audio output" error. + """ + response = make_speech_request( + host=omni_server.host, + port=omni_server.port, + text="This should return binary audio, not a JSON error.", + voice="vivian", + language="English", + ) + + # Should NOT be a JSON error response + assert response.status_code == 200, f"Request failed: {response.text}" + + # Verify it's binary audio, not JSON + try: + # If this succeeds and starts with {"error", it's a bug + text = response.content.decode("utf-8") + assert not text.startswith('{"error"'), f"Got error response instead of audio: {text}" + except UnicodeDecodeError: + # This is expected - binary audio can't be decoded as UTF-8 + pass + + assert verify_wav_audio(response.content), "Response is not valid WAV audio" + + +class TestQwen3TTSAPIEndpoints: + """Test API endpoint functionality.""" + + @pytest.mark.core_model + @pytest.mark.omni + @hardware_test(res={"cuda": "L4"}, num_cards=4) + def test_list_voices_endpoint(self, omni_server) -> None: + """Test the /v1/audio/voices endpoint returns available voices.""" + url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/voices" + + with httpx.Client(timeout=30.0) as client: + response = client.get(url) + + assert response.status_code == 200 + data = response.json() + assert "voices" in data + assert isinstance(data["voices"], list) + assert len(data["voices"]) > 0 + # Check some expected voices are present + voices_lower = [v.lower() for v in data["voices"]] + assert "vivian" in voices_lower or "ryan" in voices_lower + + @pytest.mark.core_model + @pytest.mark.omni + @hardware_test(res={"cuda": "L4"}, num_cards=4) + def test_models_endpoint(self, omni_server) -> None: + """Test the /v1/models endpoint returns loaded model.""" + url = f"http://{omni_server.host}:{omni_server.port}/v1/models" + + with httpx.Client(timeout=30.0) as client: + response = client.get(url) + + assert response.status_code == 200 + data = response.json() + assert "data" in data + assert len(data["data"]) > 0 diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 037db2ea8..d5c45a0a9 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -221,24 +221,33 @@ class OrchestratorAggregator: stage_id: int, request_id: str, ) -> None: - if ( - output_to_yield.final_output_type == "audio" - and finished - and (multimodal_output := output_to_yield.multimodal_output.get("audio")) is not None - ): - nframes = int(multimodal_output[-1].shape[0]) - stage_events_for_req = self.stage_events.get(request_id, []) - if stage_events_for_req: - for stage_event in stage_events_for_req: - if stage_event.stage_id == stage_id: - stage_event.audio_generated_frames += nframes - break - else: - logger.warning( - "Failed to record audio generated frames for request %s at stage %s: no stage event found", - request_id, - stage_id, - ) + try: + if ( + output_to_yield.final_output_type == "audio" + and finished + and (multimodal_output := output_to_yield.multimodal_output.get("audio")) is not None + and len(multimodal_output) > 0 + ): + last = multimodal_output[-1] + nframes = int(last.shape[0]) if last.ndim > 0 else 1 + stage_events_for_req = self.stage_events.get(request_id, []) + if stage_events_for_req: + for stage_event in stage_events_for_req: + if stage_event.stage_id == stage_id: + stage_event.audio_generated_frames += nframes + break + else: + logger.warning( + "Failed to record audio generated frames for request %s at stage %s: no stage event found", + request_id, + stage_id, + ) + except Exception: + logger.debug( + "Failed to record audio frames for request %s", + request_id, + exc_info=True, + ) def _as_stage_request_stats( self, From 43e8747c007373cc671eeb10aac2480d09c84952 Mon Sep 17 00:00:00 2001 From: dongbo910220 <32610838+dongbo910220@users.noreply.github.com> Date: Tue, 10 Feb 2026 17:36:10 +0800 Subject: [PATCH 08/13] [Feat]: support VAE patch parallelism (#756) Signed-off-by: dongbo910220 <1275604947@qq.com> Co-authored-by: hsliuustc0106 --- .buildkite/pipeline.yml | 2 +- .buildkite/test-amd.yaml | 2 +- docs/contributing/features/tensor_parallel.md | 2 +- .../diffusion/parallelism_acceleration.md | 69 ++- docs/user_guide/diffusion_acceleration.md | 75 ++- .../text_to_image/text_to_image.py | 23 +- .../distributed/test_vae_patch_parallel.py | 222 ++++++++ ...parallel.py => test_zimage_parallelism.py} | 97 +++- vllm_omni/diffusion/data.py | 6 +- .../diffusion/distributed/parallel_state.py | 39 +- .../distributed/vae_patch_parallel.py | 477 ++++++++++++++++++ vllm_omni/diffusion/registry.py | 39 ++ vllm_omni/engine/input_processor.py | 12 +- 13 files changed, 983 insertions(+), 82 deletions(-) create mode 100644 tests/diffusion/distributed/test_vae_patch_parallel.py rename tests/e2e/offline_inference/{test_zimage_tensor_parallel.py => test_zimage_parallelism.py} (68%) create mode 100644 vllm_omni/diffusion/distributed/vae_patch_parallel.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 75cc818a6..9e7fc07a9 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -141,7 +141,7 @@ steps: timeout_in_minutes: 20 depends_on: image-build commands: - - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py + - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py agents: queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU plugins: diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 773748c17..5b4cedd3f 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -66,7 +66,7 @@ steps: - export GPU_ARCHS=gfx942 - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py + - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py - label: "Diffusion GPU Worker Test" timeout_in_minutes: 20 diff --git a/docs/contributing/features/tensor_parallel.md b/docs/contributing/features/tensor_parallel.md index aecc1aedc..6214dc783 100644 --- a/docs/contributing/features/tensor_parallel.md +++ b/docs/contributing/features/tensor_parallel.md @@ -264,7 +264,7 @@ Complete examples in the codebase: | **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Standard TP | Full implementation with validation | | **FLUX** | `vllm_omni/diffusion/models/flux/flux_transformer.py` | Dual-stream | Image + text streams | | **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Standard TP | With RoPE | -| **TP Tests** | `tests/e2e/offline_inference/test_zimage_tensor_parallel.py` | E2E testing | TP correctness and performance | +| **TP Tests** | `tests/e2e/offline_inference/test_zimage_parallelism.py` | E2E testing | TP correctness and performance | | **Constraint Tests** | `tests/diffusion/models/z_image/test_zimage_tp_constraints.py` | Unit testing | Validation logic | --- diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index e2dae5c84..7b6e9390e 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -14,23 +14,25 @@ The following parallelism methods are currently supported in vLLM-Omni: 4. [Tensor Parallelism](#tensor-parallelism): Tensor parallelism shards model weights across devices. This can reduce per-GPU memory usage. Note that for diffusion models we currently shard the majority of layers within the DiT. +5. [VAE Patch Parallelism](#vae-patch-parallelism): VAE patch parallelism shards VAE decode spatially across ranks. This can reduce the peak memory of VAE decode and (depending on resolution and communication overhead) speed up VAE decode. + The following table shows which models are currently supported by parallelism method: ### ImageGen -| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | -|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:| -| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | -| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | -| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | -| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | -| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | -| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ | -| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | -| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | +| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | +|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:| +| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | +| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | +| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | +| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | +| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ | ❌ | +| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ | +| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | !!! note "TP Limitations for Diffusion Models" We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP. @@ -47,7 +49,7 @@ The following table shows which models are currently supported by parallelism me ### VideoGen -| Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel | +| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | |-------|------------------|------------|---------|--------------------------| | **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | @@ -76,6 +78,45 @@ outputs = omni.generate( ) ``` +### VAE Patch Parallelism + +VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP). + +!!! note "Enablement and feature gate" + - VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`). + - If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.) + +#### Offline Inference + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig + +omni = Omni( + model="Tongyi-MAI/Z-Image-Turbo", + parallel_config=DiffusionParallelConfig( + tensor_parallel_size=2, + vae_patch_parallel_size=2, + ), + vae_use_tiling=True, +) + +outputs = omni.generate( + prompt="a cat reading a book", + num_inference_steps=9, + width=1024, + height=1024, +) +``` + +#### How it works (method selection) + +VAE patch parallelism automatically selects between two internal decode methods based on whether diffusers tiling would kick in: + +- `_distributed_tiled_decode`: Used when the latent spatial size exceeds `vae.tile_latent_min_size` (i.e., diffusers tiled decode). Each rank decodes a subset of tiles; rank0 gathers and runs the same overlap+blend+stitch logic as diffusers. This matches the single-rank diffusers tiled output. + +- `_distributed_patch_decode`: Used when diffusers tiling would not kick in. Each rank decodes a grid patch expanded with a latent-space halo; then rank0 gathers the cropped core patches and stitches them into the full image. This path has no blending and can introduce small numerical differences compared to the non-parallel decode. + ### Sequence Parallelism #### Ulysses-SP diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index ba122a1cb..7856ea960 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -22,6 +22,10 @@ vLLM-Omni also supports parallelism methods for diffusion models, including: 3. [CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel) - runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step. +4. [Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism) - shards DiT weights across devices to reduce per-GPU memory usage. + +5. [VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism) - shards VAE decode/encode spatially across ranks to reduce VAE peak memory (and can speed up VAE decode). + ## Quick Comparison ### Cache Methods @@ -37,20 +41,20 @@ The following table shows which models are currently supported by each accelerat ### ImageGen -| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | -|-------|------------------|:----------:|:-----------:|:-----------:|:----------------:|:----------------:| -| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | -| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | -| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | -| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | -| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ | -| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | -| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | -| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | +| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | +|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:---------------:|:------------------:| +| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | +| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | +| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | +| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ (TP=2 only) | ✅ | +| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | +| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ### VideoGen @@ -201,6 +205,47 @@ outputs = omni.generate( ) ``` +### Using Tensor Parallelism + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig + +omni = Omni( + model="Tongyi-MAI/Z-Image-Turbo", + parallel_config=DiffusionParallelConfig(tensor_parallel_size=2), +) + +outputs = omni.generate( + prompt="a cat reading a book", + num_inference_steps=9, + width=512, + height=512, +) +``` + +### Using VAE Patch Parallelism + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig + +omni = Omni( + model="Tongyi-MAI/Z-Image-Turbo", + parallel_config=DiffusionParallelConfig( + tensor_parallel_size=2, + vae_patch_parallel_size=2, + ), +) + +outputs = omni.generate( + prompt="a cat reading a book", + num_inference_steps=9, + width=1024, + height=1024, +) +``` + ### Using CFG-Parallel Run image-to-image: @@ -233,5 +278,7 @@ For detailed information on each acceleration method: - **[TeaCache Guide](diffusion/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices - **[Cache-DiT Acceleration Guide](diffusion/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters +- **[Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism)** - Guidance on how to enable TP for diffusion models. - **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration. - **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks. +- **[VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism)** - Guidance on how to reduce VAE memory via patch/tile parallelism. diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index a79e5d640..4b2817f8d 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -118,12 +118,6 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of ready layers (blocks) to keep on GPU during generation.", ) - parser.add_argument( - "--tensor_parallel_size", - type=int, - default=1, - help="Number of GPUs used for tensor parallelism (TP) inside the DiT.", - ) parser.add_argument( "--vae_use_slicing", action="store_true", @@ -134,6 +128,18 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable VAE tiling for memory optimization.", ) + parser.add_argument( + "--tensor_parallel_size", + type=int, + default=1, + help="Number of GPUs used for tensor parallelism (TP) inside the DiT.", + ) + parser.add_argument( + "--vae_patch_parallel_size", + type=int, + default=1, + help="Number of ranks used for VAE patch/tile parallelism (decode/encode).", + ) return parser.parse_args() @@ -176,6 +182,7 @@ def main(): ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size, tensor_parallel_size=args.tensor_parallel_size, + vae_patch_parallel_size=args.vae_patch_parallel_size, ) # Check if profiling is requested via environment variable @@ -207,8 +214,10 @@ def main(): print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}") print( f" Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, " - f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}" + f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, " + f"vae_patch_parallel_size={args.vae_patch_parallel_size}" ) + print(f" CPU offload: {args.enable_cpu_offload}") print(f" Image size: {args.width}x{args.height}") print(f"{'=' * 60}\n") diff --git a/tests/diffusion/distributed/test_vae_patch_parallel.py b/tests/diffusion/distributed/test_vae_patch_parallel.py new file mode 100644 index 000000000..43bbc71d5 --- /dev/null +++ b/tests/diffusion/distributed/test_vae_patch_parallel.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for VAE patch/tile parallelism helpers (CPU-only).""" + +import pytest +import torch + +from vllm_omni.diffusion.distributed import vae_patch_parallel as vae_patch_parallel + + +class _DummyConfig: + def __init__(self, **attrs): + for k, v in attrs.items(): + setattr(self, k, v) + + +class _DummyVae: + def __init__(self, *, config=None, **attrs): + self.config = config + for k, v in attrs.items(): + setattr(self, k, v) + + +def test_get_vae_spatial_scale_factor_uses_block_out_channels_len_minus_1(): + vae = _DummyVae(config=_DummyConfig(block_out_channels=[128, 256, 512, 512])) + assert vae_patch_parallel._get_vae_spatial_scale_factor(vae) == 8 + + vae = _DummyVae(config=_DummyConfig(block_out_channels=[1, 2, 3, 4, 5])) + assert vae_patch_parallel._get_vae_spatial_scale_factor(vae) == 16 + + +def test_get_vae_spatial_scale_factor_defaults_to_8_on_missing_or_empty(): + assert vae_patch_parallel._get_vae_spatial_scale_factor(_DummyVae(config=_DummyConfig())) == 8 + assert vae_patch_parallel._get_vae_spatial_scale_factor(_DummyVae(config=_DummyConfig(block_out_channels=[]))) == 8 + assert vae_patch_parallel._get_vae_spatial_scale_factor(_DummyVae(config=None)) == 8 + + +def test_get_vae_spatial_scale_factor_defaults_to_8_on_exception(): + class _BrokenConfig: + @property + def block_out_channels(self): + raise RuntimeError("boom") + + assert vae_patch_parallel._get_vae_spatial_scale_factor(_DummyVae(config=_BrokenConfig())) == 8 + + +@pytest.mark.parametrize( + ("pp_size", "expected"), + [ + (0, (1, 1)), + (1, (1, 1)), + (2, (1, 2)), + (3, (1, 3)), + (4, (2, 2)), + (6, (2, 3)), + (8, (2, 4)), + (12, (3, 4)), + (16, (4, 4)), + ], +) +def test_factor_pp_grid(pp_size: int, expected: tuple[int, int]): + assert vae_patch_parallel._factor_pp_grid(pp_size) == expected + + +def test_get_world_rank_pp_size(monkeypatch): + monkeypatch.setattr(vae_patch_parallel.dist, "get_world_size", lambda _: 8) + monkeypatch.setattr(vae_patch_parallel.dist, "get_rank", lambda _: 3) + + world_size, rank, pp_size = vae_patch_parallel._get_world_rank_pp_size(object(), 4) + assert (world_size, rank, pp_size) == (8, 3, 4) + + world_size, rank, pp_size = vae_patch_parallel._get_world_rank_pp_size(object(), 16) + assert (world_size, rank, pp_size) == (8, 3, 8) + + +def test_get_vae_out_channels_defaults_to_3(): + assert vae_patch_parallel._get_vae_out_channels(_DummyVae(config=None)) == 3 + assert vae_patch_parallel._get_vae_out_channels(_DummyVae(config=_DummyConfig())) == 3 + + +def test_get_vae_out_channels_reads_config(): + assert vae_patch_parallel._get_vae_out_channels(_DummyVae(config=_DummyConfig(out_channels=4))) == 4 + assert vae_patch_parallel._get_vae_out_channels(_DummyVae(config=_DummyConfig(out_channels="5"))) == 5 + + +def test_get_vae_tile_params_returns_none_if_missing(): + assert ( + vae_patch_parallel._get_vae_tile_params(_DummyVae(tile_latent_min_size=None, tile_overlap_factor=0.25)) is None + ) + assert ( + vae_patch_parallel._get_vae_tile_params(_DummyVae(tile_latent_min_size=128, tile_overlap_factor=None)) is None + ) + + +def test_get_vae_tile_params_parses_types(): + vae = _DummyVae(tile_latent_min_size="128", tile_overlap_factor="0.25") + assert vae_patch_parallel._get_vae_tile_params(vae) == (128, 0.25) + + +def test_get_vae_tiling_params_returns_none_if_missing(): + vae = _DummyVae(tile_latent_min_size=128, tile_overlap_factor=0.25, tile_sample_min_size=None) + assert vae_patch_parallel._get_vae_tiling_params(vae) is None + + vae = _DummyVae(tile_latent_min_size=None, tile_overlap_factor=0.25, tile_sample_min_size=1024) + assert vae_patch_parallel._get_vae_tiling_params(vae) is None + + +def test_get_vae_tiling_params_parses_types(): + vae = _DummyVae(tile_latent_min_size="128", tile_overlap_factor="0.25", tile_sample_min_size="1024") + assert vae_patch_parallel._get_vae_tiling_params(vae) == (128, 0.25, 1024) + + +def test_distributed_tiled_decode_stitches_tiles(monkeypatch): + class _TinyConfig: + def __init__(self): + self.out_channels = 1 + self.use_post_quant_conv = False + + class _TinyVae: + def __init__(self): + self.config = _TinyConfig() + self.tile_latent_min_size = 2 + self.tile_overlap_factor = 0.0 + self.tile_sample_min_size = 2 + + def decoder(self, x: torch.Tensor) -> torch.Tensor: + return x + + def blend_v(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor: + return b + + def blend_h(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor: + return b + + def _collect_local_tiles( + *, + vae: _TinyVae, + z: torch.Tensor, + rank: int, + pp_size: int, + ) -> tuple[list[torch.Tensor], list[tuple[int, int, int]]]: + tile_latent_min_size = vae.tile_latent_min_size + overlap_size = int(tile_latent_min_size * (1 - vae.tile_overlap_factor)) + h_starts = list(range(0, z.shape[2], overlap_size)) + w_starts = list(range(0, z.shape[3], overlap_size)) + + local_tiles: list[torch.Tensor] = [] + local_meta: list[tuple[int, int, int]] = [] + tile_id = 0 + for i in h_starts: + for j in w_starts: + tile_rank = (tile_id + 1) % pp_size + if tile_rank == rank: + tile = z[:, :, i : i + tile_latent_min_size, j : j + tile_latent_min_size] + decoded = vae.decoder(tile) + local_tiles.append(decoded) + local_meta.append((tile_id, int(decoded.shape[-2]), int(decoded.shape[-1]))) + tile_id += 1 + return local_tiles, local_meta + + vae = _TinyVae() + z = torch.arange(16, dtype=torch.float32).reshape(1, 1, 4, 4) + + rank0_tiles, rank0_meta = _collect_local_tiles(vae=vae, z=z, rank=0, pp_size=2) + rank1_tiles, rank1_meta = _collect_local_tiles(vae=vae, z=z, rank=1, pp_size=2) + max_count = max(len(rank0_tiles), len(rank1_tiles)) + + def _pack_meta_and_tiles( + tiles: list[torch.Tensor], + meta: list[tuple[int, int, int]], + max_count: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + meta_tensor = torch.full((max_count, 3), -1, dtype=torch.int64) + tile_tensor = torch.zeros( + (max_count, z.shape[0], vae.config.out_channels, vae.tile_sample_min_size, vae.tile_sample_min_size), + dtype=z.dtype, + ) + for idx, (tile_id, h, w) in enumerate(meta): + meta_tensor[idx, 0] = tile_id + meta_tensor[idx, 1] = h + meta_tensor[idx, 2] = w + tile_tensor[idx, :, :, :h, :w] = tiles[idx] + return meta_tensor, tile_tensor + + rank1_meta_tensor, rank1_tile_tensor = _pack_meta_and_tiles(rank1_tiles, rank1_meta, max_count) + rank1_count_tensor = torch.tensor([len(rank1_tiles)], dtype=torch.int64) + + def _fake_gather(tensor, gather_list=None, dst=0, group=None): + if gather_list is None: + return + if tensor.ndim == 1 and tensor.numel() == 1: + gather_list[0].copy_(tensor) + gather_list[1].copy_(rank1_count_tensor) + return + if tensor.ndim == 2 and tensor.shape[1] == 3: + gather_list[0].copy_(tensor) + gather_list[1].copy_(rank1_meta_tensor) + return + if tensor.ndim == 5: + gather_list[0].copy_(tensor) + gather_list[1].copy_(rank1_tile_tensor) + return + raise AssertionError("Unexpected gather payload for test.") + + def _fake_broadcast(_tensor, src=0, group=None): + return + + monkeypatch.setattr(vae_patch_parallel.dist, "get_world_size", lambda _group: 2) + monkeypatch.setattr(vae_patch_parallel.dist, "get_rank", lambda _group: 0) + monkeypatch.setattr(vae_patch_parallel.dist, "gather", _fake_gather) + monkeypatch.setattr(vae_patch_parallel.dist, "broadcast", _fake_broadcast) + + output = vae_patch_parallel._distributed_tiled_decode( + vae=vae, + orig_decode=lambda z, return_dict=False: (z,), + z=z, + group=object(), + vae_patch_parallel_size=2, + ) + + assert torch.equal(output, z) diff --git a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py b/tests/e2e/offline_inference/test_zimage_parallelism.py similarity index 68% rename from tests/e2e/offline_inference/test_zimage_tensor_parallel.py rename to tests/e2e/offline_inference/test_zimage_parallelism.py index 2d051e5aa..79fdb4699 100644 --- a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py +++ b/tests/e2e/offline_inference/test_zimage_parallelism.py @@ -1,6 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Z-Image end-to-end tests for diffusion parallelism. + +This file currently covers: +- DiT tensor parallelism (TP=2) vs TP=1. +- VAE patch parallelism (vae_patch_parallel_size=2) vs baseline on TP=2. + +Note: CUDA-only (>=2 GPUs). We use `enforce_eager=False` (default) to enable +`torch.compile`. +""" + import os import sys import time @@ -69,16 +79,32 @@ def _extract_single_image(outputs) -> Image.Image: def _run_zimage_generate( - *, tp_size: int, height: int, width: int, num_inference_steps: int, seed: int + *, + tp_size: int, + height: int, + width: int, + num_inference_steps: int, + seed: int, + enforce_eager: bool, + vae_use_tiling: bool = False, + vae_patch_parallel_size: int = 1, + num_requests: int = 4, ) -> tuple[Image.Image, float, float]: + if num_requests < 2: + raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)") + torch.cuda.empty_cache() device_index = torch.cuda.current_device() monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02) monitor.start() - m = Omni( model=_get_zimage_model(), - parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size), + parallel_config=DiffusionParallelConfig( + tensor_parallel_size=tp_size, + vae_patch_parallel_size=vae_patch_parallel_size, + ), + enforce_eager=enforce_eager, + vae_use_tiling=vae_use_tiling, ) try: # NOTE: Omni closes itself when a generate() call is exhausted. @@ -89,7 +115,6 @@ def _run_zimage_generate( # This also serves as a warmup: the first output may include extra # compilation/caching overhead, while later outputs are closer to # steady-state inference. - num_requests = 4 # 1 warmup + 3 timed gen = m.generate( [PROMPT] * num_requests, OmniDiffusionSamplingParams( @@ -104,6 +129,7 @@ def _run_zimage_generate( ) warmup_output = next(gen) + t_prev = time.perf_counter() per_request_times_s: list[float] = [] last_output = warmup_output @@ -124,6 +150,7 @@ def _run_zimage_generate( return _extract_single_image([last_output]), median_time_s, peak_memory_mb finally: monitor.stop() + m.close() cleanup_dist_env_and_memory() @@ -137,6 +164,8 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): if not torch.cuda.is_available() or torch.cuda.device_count() < 2: pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.") + enforce_eager = False + height = 512 width = 512 num_inference_steps = 2 @@ -148,6 +177,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): width=width, num_inference_steps=num_inference_steps, seed=seed, + enforce_eager=enforce_eager, ) tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate( tp_size=2, @@ -155,6 +185,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): width=width, num_inference_steps=num_inference_steps, seed=seed, + enforce_eager=enforce_eager, ) tp1_path = tmp_path / "zimage_tp1.png" @@ -186,3 +217,61 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): assert tp2_peak_mem < tp1_peak_mem, ( f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})" ) + + +@pytest.mark.integration +def test_zimage_vae_patch_parallel_tp2(tmp_path: Path): + if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): + pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA for now.") + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 CUDA devices.") + + enforce_eager = False + + # Use a larger image to ensure there are multiple VAE tiles. + height = 1152 + width = 1152 + num_inference_steps = 2 + seed = 42 + + baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate( + tp_size=2, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + enforce_eager=enforce_eager, + vae_use_tiling=True, + vae_patch_parallel_size=1, + num_requests=2, + ) + pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate( + tp_size=2, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + enforce_eager=enforce_eager, + vae_use_tiling=True, + vae_patch_parallel_size=2, + num_requests=2, + ) + + baseline_path = tmp_path / "zimage_tp2_vae_pp1.png" + pp2_path = tmp_path / "zimage_tp2_vae_pp2.png" + baseline_img.save(baseline_path) + pp2_img.save(pp2_path) + + mean_abs_diff, max_abs_diff = _diff_metrics(baseline_img, pp2_img) + mean_threshold = 5e-3 + max_threshold = 1e-1 + print( + "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): " + f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; " + f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; " + f"pp1_img={baseline_path}, pp2_img={pp2_path}" + ) + assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, ( + f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} " + f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e})" + ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index f884fb6f1..a4f0ba6fa 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -45,6 +45,9 @@ class DiffusionParallelConfig: cfg_parallel_size: int = 1 """Number of Classifier Free Guidance (CFG) parallel groups.""" + vae_patch_parallel_size: int = 1 + """Number of ranks used for VAE patch/tile parallelism (decode/encode).""" + @model_validator(mode="after") def _validate_parallel_config(self) -> Self: """Validates the config relationships among the parallel strategies.""" @@ -56,6 +59,7 @@ class DiffusionParallelConfig: assert self.ring_degree > 0, "Ring degree must be > 0" assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0" assert self.cfg_parallel_size in [1, 2], f"CFG parallel size must be 1 or 2, but got {self.cfg_parallel_size}" + assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0" assert self.sequence_parallel_size == self.ulysses_degree * self.ring_degree, ( "Sequence parallel size must be equal to the product of ulysses degree and ring degree," f" but got {self.sequence_parallel_size} != {self.ulysses_degree} * {self.ring_degree}" @@ -65,6 +69,7 @@ class DiffusionParallelConfig: def __post_init__(self) -> None: if self.sequence_parallel_size is None: self.sequence_parallel_size = self.ulysses_degree * self.ring_degree + self.world_size = ( self.pipeline_parallel_size * self.data_parallel_size @@ -359,7 +364,6 @@ class OmniDiffusionConfig: # support multi images input supports_multimodal_inputs: bool = False - # Logging log_level: str = "info" # Omni configuration (injected from stage config) diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 12da9c43b..26112af46 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -57,7 +57,6 @@ _PP: PipelineGroupCoordinator | None = None _CFG: GroupCoordinator | None = None _DP: GroupCoordinator | None = None _DIT: GroupCoordinator | None = None -_VAE: GroupCoordinator | None = None def generate_masked_orthogonal_rank_groups( @@ -344,7 +343,7 @@ def is_dp_last_group(): def get_dit_world_size(): - """Return world size for the DiT model (excluding VAE).""" + """Return world size for the DiT model.""" return ( get_data_parallel_world_size() * get_classifier_free_guidance_world_size() @@ -354,22 +353,6 @@ def get_dit_world_size(): ) -# Add VAE getter functions -def get_vae_parallel_group() -> GroupCoordinator: - assert _VAE is not None, "VAE parallel group is not initialized" - return _VAE - - -def get_vae_parallel_world_size(): - """Return world size for the VAE parallel group.""" - return get_vae_parallel_group().world_size - - -def get_vae_parallel_rank(): - """Return my rank for the VAE parallel group.""" - return get_vae_parallel_group().rank_in_group - - # * SET @@ -490,18 +473,6 @@ def get_dit_group(): return _DIT -def init_vae_group( - dit_parallel_size: int, - vae_parallel_size: int, - backend: str, -): - # Initialize VAE group first - global _VAE - assert _VAE is None, "VAE parallel group is already initialized" - vae_ranks = list(range(dit_parallel_size, dit_parallel_size + vae_parallel_size)) - _VAE = torch.distributed.new_group(ranks=vae_ranks, backend=backend) - - # adapted from https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py def set_seq_parallel_pg( sp_ulysses_degree: int, @@ -658,7 +629,6 @@ def initialize_model_parallel( ring_degree: int = 1, tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, - vae_parallel_size: int = 0, backend: str | None = None, ) -> None: if backend is None: @@ -802,8 +772,6 @@ def initialize_model_parallel( backend=backend, parallel_mode="tensor", ) - if vae_parallel_size > 0: - init_vae_group(dit_parallel_size, vae_parallel_size, backend) init_dit_group(dit_parallel_size, backend) @@ -833,11 +801,6 @@ def destroy_model_parallel(): _PP.destroy() _PP = None - global _VAE - if _VAE: - _VAE.destroy() - _VAE = None - def destroy_distributed_environment(): global _WORLD diff --git a/vllm_omni/diffusion/distributed/vae_patch_parallel.py b/vllm_omni/diffusion/distributed/vae_patch_parallel.py new file mode 100644 index 000000000..1a3edde3c --- /dev/null +++ b/vllm_omni/diffusion/distributed/vae_patch_parallel.py @@ -0,0 +1,477 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Distributed VAE patch/tile parallelism utilities.""" + +from __future__ import annotations + +import math +from collections.abc import Callable +from typing import Any + +import torch +import torch.distributed as dist +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def _get_vae_spatial_scale_factor(vae: Any) -> int: + try: + block_out_channels = getattr(getattr(vae, "config", None), "block_out_channels", None) + if block_out_channels: + return 2 ** (len(block_out_channels) - 1) + except Exception: + pass + return 8 + + +def _factor_pp_grid(pp_size: int) -> tuple[int, int]: + """Pick a (rows, cols) grid whose product equals `pp_size`.""" + if pp_size <= 1: + return (1, 1) + root = int(math.sqrt(pp_size)) + for rows in range(root, 0, -1): + if pp_size % rows == 0: + return (rows, pp_size // rows) + return (1, pp_size) + + +def _get_world_rank_pp_size( + group: dist.ProcessGroup, + vae_patch_parallel_size: int, +) -> tuple[int, int, int]: + world_size = dist.get_world_size(group) + rank = dist.get_rank(group) + pp_size = min(int(vae_patch_parallel_size), int(world_size)) + return world_size, rank, pp_size + + +def _get_vae_out_channels(vae: Any) -> int: + return int(getattr(getattr(vae, "config", None), "out_channels", 3)) + + +def _get_vae_tile_params(vae: Any) -> tuple[int, float] | None: + tile_latent_min_size = getattr(vae, "tile_latent_min_size", None) + tile_overlap_factor = getattr(vae, "tile_overlap_factor", None) + if tile_latent_min_size is None or tile_overlap_factor is None: + return None + return int(tile_latent_min_size), float(tile_overlap_factor) + + +def _get_vae_tiling_params(vae: Any) -> tuple[int, float, int] | None: + tile_sample_min_size = getattr(vae, "tile_sample_min_size", None) + tile_params = _get_vae_tile_params(vae) + if tile_params is None or tile_sample_min_size is None: + return None + tile_latent_min_size, tile_overlap_factor = tile_params + return tile_latent_min_size, tile_overlap_factor, int(tile_sample_min_size) + + +def _distributed_tiled_decode( + *, + vae: Any, + orig_decode: Callable[..., Any], + z: torch.Tensor, + group: dist.ProcessGroup, + vae_patch_parallel_size: int, +) -> torch.Tensor: + """Distributed version of diffusers AutoencoderKL.tiled_decode (decode only). + + Each rank decodes a subset of tiles; rank0 gathers all tiles and performs the + original blend + stitch logic. Non-rank0 ranks return an empty tensor; callers + can broadcast the stitched result if needed. + """ + world_size, rank, pp_size = _get_world_rank_pp_size(group, vae_patch_parallel_size) + if pp_size <= 1: + return orig_decode(z, return_dict=False)[0] + + tiling_params = _get_vae_tiling_params(vae) + if tiling_params is None: + return orig_decode(z, return_dict=False)[0] + tile_latent_min_size, tile_overlap_factor, tile_sample_min_size = tiling_params + + overlap_size = int(tile_latent_min_size * (1 - tile_overlap_factor)) + if overlap_size <= 0: + return orig_decode(z, return_dict=False)[0] + + h_starts = list(range(0, z.shape[2], overlap_size)) + w_starts = list(range(0, z.shape[3], overlap_size)) + num_rows = len(h_starts) + num_cols = len(w_starts) + num_tiles = num_rows * num_cols + + if num_tiles < 2: + return orig_decode(z, return_dict=False)[0] + + blend_extent = int(tile_sample_min_size * tile_overlap_factor) + row_limit = int(tile_sample_min_size - blend_extent) + + # Decide which ranks actively decode tiles. + active = rank < pp_size + + local_tiles: list[torch.Tensor] = [] + local_meta: list[tuple[int, int, int]] = [] + + tile_id = 0 + for i in h_starts: + for j in w_starts: + # Offset assignment by 1 so rank0 avoids decoding the largest (tile_id=0) tile. + tile_rank = (tile_id + 1) % pp_size + if active and (tile_rank == rank): + tile = z[:, :, i : i + tile_latent_min_size, j : j + tile_latent_min_size] + if getattr(getattr(vae, "config", None), "use_post_quant_conv", False): + tile = vae.post_quant_conv(tile) + decoded = vae.decoder(tile) + local_tiles.append(decoded) + local_meta.append((tile_id, int(decoded.shape[-2]), int(decoded.shape[-1]))) + tile_id += 1 + + # Gather per-rank tile counts. + count_tensor = torch.tensor([len(local_tiles)], device=z.device, dtype=torch.int64) + if rank == 0: + count_gather = [torch.empty_like(count_tensor) for _ in range(world_size)] + else: + count_gather = None + dist.gather(count_tensor, gather_list=count_gather, dst=0, group=group) + max_count = 0 + if rank == 0: + counts = [int(t.item()) for t in count_gather] # type: ignore[arg-type] + max_count = max(counts) if counts else 0 + max_count_tensor = torch.tensor([max_count], device=z.device, dtype=torch.int64) + dist.broadcast(max_count_tensor, src=0, group=group) + max_count = int(max_count_tensor.item()) + + out_channels = _get_vae_out_channels(vae) + + # Prepare padded metadata + tiles for gather. + meta_tensor = torch.full((max_count, 3), -1, device=z.device, dtype=torch.int64) + tile_tensor = torch.zeros( + (max_count, z.shape[0], out_channels, tile_sample_min_size, tile_sample_min_size), + device=z.device, + dtype=z.dtype, + ) + for idx, (tile_id, h, w) in enumerate(local_meta): + meta_tensor[idx, 0] = tile_id + meta_tensor[idx, 1] = h + meta_tensor[idx, 2] = w + tile_tensor[idx, :, :, :h, :w] = local_tiles[idx] + + if rank == 0: + meta_gather = [torch.empty_like(meta_tensor) for _ in range(world_size)] + tile_gather = [torch.empty_like(tile_tensor) for _ in range(world_size)] + else: + meta_gather = None + tile_gather = None + + dist.gather(meta_tensor, gather_list=meta_gather, dst=0, group=group) + dist.gather(tile_tensor, gather_list=tile_gather, dst=0, group=group) + + if rank != 0: + return torch.empty(0, device=z.device, dtype=z.dtype) + + # Reconstruct the full tile grid on rank0. + tile_map: dict[int, torch.Tensor] = {} + for src_rank in range(world_size): + meta_src = meta_gather[src_rank] # type: ignore[index] + tiles_src = tile_gather[src_rank] # type: ignore[index] + for idx in range(max_count): + tid = int(meta_src[idx, 0].item()) + if tid < 0: + continue + h = int(meta_src[idx, 1].item()) + w = int(meta_src[idx, 2].item()) + tile_map[tid] = tiles_src[idx, :, :, :h, :w] + + rows: list[list[torch.Tensor]] = [] + for r in range(num_rows): + row: list[torch.Tensor] = [] + for c in range(num_cols): + tid = r * num_cols + c + row.append(tile_map[tid]) + rows.append(row) + + result_rows: list[torch.Tensor] = [] + for i, row in enumerate(rows): + result_row: list[torch.Tensor] = [] + for j, tile in enumerate(row): + if i > 0: + tile = vae.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = vae.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=3)) + + return torch.cat(result_rows, dim=2) + + +def _distributed_patch_decode( + *, + vae: Any, + orig_decode: Callable[..., Any], + z: torch.Tensor, + group: dist.ProcessGroup, + vae_patch_parallel_size: int, + vae_scale_factor: int, +) -> torch.Tensor: + """Decode one spatial block per rank, then stitch RGB blocks on rank0. + + Intended for sizes where diffusers tiling would not kick in, so we can still + reduce the per-rank VAE decode activation peak. Each rank decodes a core + block with a small latent-space halo, then crops to the core and gathers the + RGB blocks to rank0 for final stitching. + """ + world_size, rank, pp_size = _get_world_rank_pp_size(group, vae_patch_parallel_size) + if pp_size <= 1: + return orig_decode(z, return_dict=False)[0] + + tile_params = _get_vae_tile_params(vae) + if tile_params is None: + return orig_decode(z, return_dict=False)[0] + tile_latent_min_size, tile_overlap_factor = tile_params + + overlap_latent = int(tile_latent_min_size * float(tile_overlap_factor)) + halo_base = max(0, overlap_latent // 2) + + # Only ranks < pp_size participate in decoding. Others send empty tensors. + active = rank < pp_size + + bsz, _, latent_h, latent_w = z.shape + scale = int(vae_scale_factor) + out_h = latent_h * scale + out_w = latent_w * scale + + out_channels = _get_vae_out_channels(vae) + + local_core = torch.empty(0, device=z.device, dtype=z.dtype) + local_h = 0 + local_w = 0 + + grid_rows, grid_cols = _factor_pp_grid(pp_size) + + if active: + patch_idx = rank + patch_row = patch_idx // grid_cols + patch_col = patch_idx % grid_cols + + h0 = (patch_row * latent_h) // grid_rows + h1 = ((patch_row + 1) * latent_h) // grid_rows + w0 = (patch_col * latent_w) // grid_cols + w1 = ((patch_col + 1) * latent_w) // grid_cols + + core_h = max(0, h1 - h0) + core_w = max(0, w1 - w0) + if core_h == 0 or core_w == 0: + local_core = torch.empty(0, device=z.device, dtype=z.dtype) + else: + halo = max(halo_base, min(core_h, core_w) // 2) + ph0 = max(0, h0 - halo) + ph1 = min(latent_h, h1 + halo) + pw0 = max(0, w0 - halo) + pw1 = min(latent_w, w1 + halo) + + tile = z[:, :, ph0:ph1, pw0:pw1] + if getattr(getattr(vae, "config", None), "use_post_quant_conv", False): + tile = vae.post_quant_conv(tile) + decoded = vae.decoder(tile) + + ch0 = (h0 - ph0) * scale + cw0 = (w0 - pw0) * scale + ch1 = ch0 + core_h * scale + cw1 = cw0 + core_w * scale + local_core = decoded[:, :, ch0:ch1, cw0:cw1] + + local_h = int(local_core.shape[-2]) if local_core.numel() else 0 + local_w = int(local_core.shape[-1]) if local_core.numel() else 0 + + # Gather block shapes. + shape_tensor = torch.tensor([local_h, local_w], device=z.device, dtype=torch.int64) + if rank == 0: + shape_gather = [torch.empty_like(shape_tensor) for _ in range(world_size)] + else: + shape_gather = None + dist.gather(shape_tensor, gather_list=shape_gather, dst=0, group=group) + + max_h = 0 + max_w = 0 + if rank == 0: + shapes = [tuple(int(x.item()) for x in t) for t in shape_gather] # type: ignore[arg-type] + max_h = max((h for h, _ in shapes), default=0) + max_w = max((w for _, w in shapes), default=0) + + max_hw_tensor = torch.tensor([max_h, max_w], device=z.device, dtype=torch.int64) + dist.broadcast(max_hw_tensor, src=0, group=group) + max_h = int(max_hw_tensor[0].item()) + max_w = int(max_hw_tensor[1].item()) + + # Pad local block for gather. + if max_h == 0 or max_w == 0: + padded = torch.empty(0, device=z.device, dtype=z.dtype) + else: + padded = torch.zeros((bsz, out_channels, max_h, max_w), device=z.device, dtype=z.dtype) + if local_h and local_w: + padded[:, :, :local_h, :local_w] = local_core + + if rank == 0: + block_gather = [torch.empty_like(padded) for _ in range(world_size)] + else: + block_gather = None + dist.gather(padded, gather_list=block_gather, dst=0, group=group) + + if rank != 0: + return torch.empty(0, device=z.device, dtype=z.dtype) + + # Stitch on rank0. + out = torch.empty((bsz, out_channels, out_h, out_w), device=z.device, dtype=z.dtype) + + for patch_idx in range(pp_size): + src_rank = patch_idx + patch_row = patch_idx // grid_cols + patch_col = patch_idx % grid_cols + + h0 = (patch_row * latent_h) // grid_rows + h1 = ((patch_row + 1) * latent_h) // grid_rows + w0 = (patch_col * latent_w) // grid_cols + w1 = ((patch_col + 1) * latent_w) // grid_cols + + ph = (h1 - h0) * scale + pw = (w1 - w0) * scale + if ph <= 0 or pw <= 0: + continue + + tile = block_gather[src_rank] # type: ignore[index] + out[:, :, h0 * scale : h1 * scale, w0 * scale : w1 * scale] = tile[:, :, :ph, :pw] + + return out + + +class VaePatchParallelism: + """Patch/tile-parallel VAE decode wrapper. + + This is meant to wrap `vae.decode` as an instance-level override + so pipelines don't need model-specific code paths. + """ + + def __init__( + self, + vae: Any, + *, + vae_patch_parallel_size: int, + group_getter: Callable[[], dist.ProcessGroup], + ) -> None: + self._vae = vae + self._vae_patch_parallel_size = int(vae_patch_parallel_size) + self._group_getter = group_getter + + self._vae_scale_factor = _get_vae_spatial_scale_factor(vae) + self._orig_decode = vae.decode + + def decode(self, z: torch.Tensor, return_dict: bool = True, *args: Any, **kwargs: Any): + # Keep the original path for unsupported VAE types / shapes. + if z.ndim != 4: + return self._orig_decode(z, return_dict=return_dict, *args, **kwargs) + + if self._vae_patch_parallel_size <= 1 or not dist.is_initialized(): + return self._orig_decode(z, return_dict=return_dict, *args, **kwargs) + + if not getattr(self._vae, "use_tiling", False): + return self._orig_decode(z, return_dict=return_dict, *args, **kwargs) + + try: + group = self._group_getter() + except Exception: + return self._orig_decode(z, return_dict=return_dict, *args, **kwargs) + + world_size = dist.get_world_size(group) + rank = dist.get_rank(group) + pp_size = min(int(self._vae_patch_parallel_size), int(world_size)) + if pp_size <= 1: + return self._orig_decode(z, return_dict=return_dict, *args, **kwargs) + + # Match diffusers' condition for when VAE tiling would be used. + tile_latent_min_size = getattr(self._vae, "tile_latent_min_size", None) + if tile_latent_min_size is None: + decoded = _distributed_tiled_decode( + vae=self._vae, + orig_decode=self._orig_decode, + z=z, + group=group, + vae_patch_parallel_size=pp_size, + ) + else: + should_tile = (z.shape[-1] > tile_latent_min_size) or (z.shape[-2] > tile_latent_min_size) + if should_tile: + decoded = _distributed_tiled_decode( + vae=self._vae, + orig_decode=self._orig_decode, + z=z, + group=group, + vae_patch_parallel_size=pp_size, + ) + else: + decoded = _distributed_patch_decode( + vae=self._vae, + orig_decode=self._orig_decode, + z=z, + group=group, + vae_patch_parallel_size=pp_size, + vae_scale_factor=self._vae_scale_factor, + ) + + if rank == 0 and decoded.numel() == 0: + logger.warning("VAE patch parallel decode produced empty output on rank0; falling back to vae.decode.") + decoded = self._orig_decode(z, return_dict=False, *args, **kwargs)[0] + if rank == 0 and decoded.dtype != z.dtype: + decoded = decoded.to(dtype=z.dtype) + if rank == 0 and not decoded.is_contiguous(): + decoded = decoded.contiguous() + + shape_tensor = torch.empty((4,), device=z.device, dtype=torch.int64) + if rank == 0: + shape_tensor.copy_(torch.tensor(tuple(decoded.shape), device=z.device, dtype=torch.int64)) + dist.broadcast(shape_tensor, src=0, group=group) + + if rank != 0: + decoded = torch.empty(tuple(int(x) for x in shape_tensor.tolist()), device=z.device, dtype=z.dtype) + dist.broadcast(decoded, src=0, group=group) + + if not return_dict: + return (decoded,) + + from diffusers.models.autoencoders.vae import DecoderOutput + + return DecoderOutput(sample=decoded) + + +def maybe_wrap_vae_decode_with_patch_parallelism( + pipeline: Any, + *, + vae_patch_parallel_size: int, + group_getter: Callable[[], dist.ProcessGroup], +) -> None: + """Wrap a diffusers-style pipeline's `vae.decode` with patch/tile parallel decode.""" + if vae_patch_parallel_size <= 1: + return + + vae = getattr(pipeline, "vae", None) + if vae is None or not hasattr(vae, "decode"): + return + try: + from diffusers.models.autoencoders import AutoencoderKL + except Exception: + return + if not isinstance(vae, AutoencoderKL): + return + + if getattr(vae, "_vllm_vae_patch_parallel_installed", False): + return + + wrapper = VaePatchParallelism( + vae, + vae_patch_parallel_size=vae_patch_parallel_size, + group_getter=group_getter, + ) + + vae._vllm_vae_patch_parallel_installed = True # type: ignore[attr-defined] + vae._vllm_vae_patch_parallel_original_decode = vae.decode # type: ignore[attr-defined] + vae.decode = wrapper.decode # type: ignore[assignment] diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 3d236ac31..173c576a2 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -113,6 +113,11 @@ DiffusionModelRegistry = _ModelRegistry( } ) +_VAE_PATCH_PARALLEL_ALLOWLIST = { + # Only enable for models we have validated end-to-end. + "ZImagePipeline", +} + def initialize_model( od_config: OmniDiffusionConfig, @@ -137,12 +142,46 @@ def initialize_model( model_class = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name) if model_class is not None: model = model_class(od_config=od_config) + + vae_pp_size = od_config.parallel_config.vae_patch_parallel_size + if vae_pp_size > 1 and od_config.model_class_name not in _VAE_PATCH_PARALLEL_ALLOWLIST: + logger.warning( + "vae_patch_parallel_size=%d is set but VAE patch parallelism is only enabled for %s; ignoring.", + vae_pp_size, + sorted(_VAE_PATCH_PARALLEL_ALLOWLIST), + ) + if ( + vae_pp_size > 1 + and od_config.model_class_name in _VAE_PATCH_PARALLEL_ALLOWLIST + and not od_config.vae_use_tiling + ): + logger.info( + "vae_patch_parallel_size=%d requires vae_use_tiling; automatically enabling it.", + vae_pp_size, + ) + od_config.vae_use_tiling = True + # Configure VAE memory optimization settings from config if hasattr(model.vae, "use_slicing"): model.vae.use_slicing = od_config.vae_use_slicing if hasattr(model.vae, "use_tiling"): model.vae.use_tiling = od_config.vae_use_tiling + if ( + vae_pp_size > 1 + and hasattr(model, "vae") + and od_config.model_class_name in _VAE_PATCH_PARALLEL_ALLOWLIST + and od_config.vae_use_tiling + ): + from vllm_omni.diffusion.distributed.parallel_state import get_dit_group + from vllm_omni.diffusion.distributed.vae_patch_parallel import maybe_wrap_vae_decode_with_patch_parallelism + + maybe_wrap_vae_decode_with_patch_parallelism( + model, + vae_patch_parallel_size=vae_pp_size, + group_getter=get_dit_group, + ) + # Apply sequence parallelism if enabled # This follows diffusers' pattern where enable_parallelism() is called # at model loading time, not inside individual model files diff --git a/vllm_omni/engine/input_processor.py b/vllm_omni/engine/input_processor.py index a1e467a88..64ffab3a1 100644 --- a/vllm_omni/engine/input_processor.py +++ b/vllm_omni/engine/input_processor.py @@ -11,7 +11,17 @@ from vllm.inputs.parse import split_enc_dec_inputs from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalUUIDDict -from vllm.multimodal.processing.context import set_request_id + +try: + from vllm.multimodal.processing import set_request_id +except ImportError: # vllm without set_request_id (older releases) + from contextlib import contextmanager + + @contextmanager + def set_request_id(_request_id: str): + yield + + from vllm.multimodal.utils import argsort_mm_positions from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams From a05aa1dd6a0a2e63019f417d0b1d733133454b9f Mon Sep 17 00:00:00 2001 From: Gao Han Date: Tue, 10 Feb 2026 17:54:12 +0800 Subject: [PATCH 09/13] [CI] Disable Qwen3-TTS E2E Test in pipeline.yml (#1306) Signed-off-by: Gao Han --- .buildkite/pipeline.yml | 38 +++++++++++++++++++------------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 9e7fc07a9..9865a4a4b 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -228,25 +228,25 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Qwen3-TTS E2E Test" - timeout_in_minutes: 10 - depends_on: image-build - commands: - - export VLLM_LOGGING_LEVEL=DEBUG - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py - agents: - queue: "gpu_4_queue" - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - shm-size: "8gb" - environment: - - "HF_HOME=/fsx/hf_cache" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" + # - label: "Qwen3-TTS E2E Test" + # timeout_in_minutes: 10 + # depends_on: image-build + # commands: + # - export VLLM_LOGGING_LEVEL=DEBUG + # - export VLLM_WORKER_MULTIPROC_METHOD=spawn + # - pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py + # agents: + # queue: "gpu_4_queue" + # plugins: + # - docker#v5.2.0: + # image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + # always-pull: true + # propagate-environment: true + # shm-size: "8gb" + # environment: + # - "HF_HOME=/fsx/hf_cache" + # volumes: + # - "/fsx/hf_cache:/fsx/hf_cache" # - label: "Omni Model Test with H100" # timeout_in_minutes: 30 From e192b0a6c5d6e9c90e1ad710cb6b4acb1d5ec95c Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Tue, 10 Feb 2026 18:47:08 +0800 Subject: [PATCH 10/13] [Misc] Add per-request generator_device to online image gen and edit (#1183) Signed-off-by: gcanlin --- vllm_omni/diffusion/worker/diffusion_model_runner.py | 8 +++++++- vllm_omni/entrypoints/openai/api_server.py | 3 +++ vllm_omni/entrypoints/openai/protocol/images.py | 4 ++++ vllm_omni/inputs/data.py | 1 + 4 files changed, 15 insertions(+), 1 deletion(-) diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 0e12460e8..1e00e4b0c 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -158,7 +158,13 @@ class DiffusionModelRunner: self.kv_transfer_manager.receive_kv_cache(req, target_device=getattr(self.pipeline, "device", None)) if req.sampling_params.generator is None and req.sampling_params.seed is not None: - req.sampling_params.generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed) + if req.sampling_params.generator_device is not None: + gen_device = req.sampling_params.generator_device + elif self.device.type == "cpu": + gen_device = "cpu" + else: + gen_device = self.device + req.sampling_params.generator = torch.Generator(device=gen_device).manual_seed(req.sampling_params.seed) # Refresh cache context if needed if ( diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index b5e48ab95..c4fc970c4 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -971,6 +971,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) _update_if_not_none( gen_params, "seed", request.seed if request.seed is not None else random.randint(0, 2**32 - 1) ) + _update_if_not_none(gen_params, "generator_device", request.generator_device) request_id = f"img_gen_{uuid.uuid4().hex}" @@ -1047,6 +1048,7 @@ async def edit_images( guidance_scale: float | None = Form(None), true_cfg_scale: float | None = Form(None), seed: int | None = Form(None), + generator_device: str | None = Form(None), # vllm-omni extension for per-request LoRA. lora: str | None = Form(None), # Json string ) -> ImageGenerationResponse: @@ -1129,6 +1131,7 @@ async def edit_images( # This fixes issues where using the default global generator # might produce blurry images in some environments. _update_if_not_none(gen_params, "seed", seed if seed is not None else random.randint(0, 2**32 - 1)) + _update_if_not_none(gen_params, "generator_device", generator_device) # 4. Generate images using AsyncOmni (multi-stage mode) request_id = f"img_edit_{int(time.time())}" diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index 63514b191..aed6536a6 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -88,6 +88,10 @@ class ImageGenerationRequest(BaseModel): description="True CFG scale (model-specific parameter, may be ignored if not supported)", ) seed: int | None = Field(default=None, description="Random seed for reproducibility") + generator_device: str | None = Field( + default=None, + description="Device for the seeded torch.Generator (e.g. 'cpu', 'cuda'). Defaults to the runner's device.", + ) # vllm-omni extension for per-request LoRA. # This mirrors the `extra_body.lora` convention in /v1/chat/completions. diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 3684744ba..bc9434817 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -166,6 +166,7 @@ class OmniDiffusionSamplingParams: num_outputs_per_prompt: int = 1 seed: int | None = None generator: torch.Generator | list[torch.Generator] | None = None + generator_device: str | None = None # layered info layers: int = 4 From 8228b5a8fe32546874687d74a8fb2a0a758098da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Tue, 10 Feb 2026 19:17:44 +0800 Subject: [PATCH 11/13] [Bagel]: Support TP (#1293) Signed-off-by: princepride --- examples/offline_inference/bagel/README.md | 18 + examples/online_serving/bagel/README.md | 21 +- .../models/bagel/bagel_transformer.py | 314 ++++++++++++++---- .../diffusion/models/bagel/pipeline_bagel.py | 28 +- 4 files changed, 314 insertions(+), 67 deletions(-) diff --git a/examples/offline_inference/bagel/README.md b/examples/offline_inference/bagel/README.md index 7bcfb42fa..c1850c27a 100644 --- a/examples/offline_inference/bagel/README.md +++ b/examples/offline_inference/bagel/README.md @@ -151,6 +151,24 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can ------ +#### Tensor Parallelism (TP) + +For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) by modifying the stage configuration (e.g., [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)). + +1. **Set `tensor_parallel_size`**: Increase this value (e.g., to `2` or `4`). +2. **Set `devices`**: Specify the comma-separated GPU IDs to be used for the stage (e.g., `"0,1"`). + +Example configuration for TP=2 on GPUs 0 and 1: +```yaml + engine_args: + tensor_parallel_size: 2 + ... + runtime: + devices: "0,1" +``` + +------ + #### 🔗 Runtime Configuration | Parameter | Value | Description | diff --git a/examples/online_serving/bagel/README.md b/examples/online_serving/bagel/README.md index 3fbea0550..9308acb77 100644 --- a/examples/online_serving/bagel/README.md +++ b/examples/online_serving/bagel/README.md @@ -22,12 +22,29 @@ cd /workspace/vllm-omni/examples/online_serving/bagel bash run_server.sh ``` -If you have a custom stage configs file, launch the server with the command below: - ```bash vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file ``` +#### 🚀 Tensor Parallelism (TP) + +For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) for the server. + +1. **Modify Stage Config**: Create or modify a stage configuration yaml (e.g., [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)). Set `tensor_parallel_size` to `2` (or more) and update `devices` to include multiple GPU IDs (e.g., `"0,1"`). + +```yaml + engine_args: + tensor_parallel_size: 2 + ... + runtime: + devices: "0,1" +``` + +2. **Launch Server**: +```bash +vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/your/custom_bagel.yaml +``` + ### Send Multi-modal Request Get into the bagel folder: diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 6ee81a4fd..839fc1d72 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -8,6 +8,7 @@ # available at https://github.com/huggingface/transformers/blob/main/LICENSE. import math +from collections.abc import Iterable from dataclasses import dataclass from typing import Any @@ -17,13 +18,18 @@ from torch import nn from torch.nn.attention.flex_attention import flex_attention from transformers.models.qwen2.configuration_qwen2 import Qwen2Config from transformers.models.qwen2.modeling_qwen2 import ( - Qwen2Attention, - Qwen2MLP, Qwen2PreTrainedModel, - Qwen2RMSNorm, - Qwen2RotaryEmbedding, ) from transformers.utils import ModelOutput +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.transformers_utils.configs.bagel import BagelConfig from vllm.vllm_flash_attn import flash_attn_varlen_func @@ -53,17 +59,105 @@ def patchify(imgs, p): class MLPconnector(nn.Module): def __init__(self, input_dim, output_dim, activation="gelu_pytorch_tanh"): super().__init__() - self.fc1 = nn.Linear(input_dim, output_dim) + self.fc1 = ColumnParallelLinear(input_dim, output_dim, bias=True, gather_output=False) if activation == "gelu": self.act = nn.GELU() elif activation == "gelu_pytorch_tanh": self.act = nn.GELU(approximate="tanh") else: self.act = nn.ReLU() - self.fc2 = nn.Linear(output_dim, output_dim) + self.fc2 = RowParallelLinear(output_dim, output_dim, bias=True, input_is_parallel=True) def forward(self, x): - return self.fc2(self.act(self.fc1(x))) + x_parallel, _ = self.fc1(x) + x_parallel = self.act(x_parallel) + return self.fc2(x_parallel)[0] + + +class BagelRotaryEmbedding(nn.Module): + """Standalone rotary embedding that generates cos/sin from position ids. + + Replaces HuggingFace's Qwen2RotaryEmbedding while preserving full + ``rope_scaling`` support. When ``config.rope_scaling`` is set (e.g. + linear, dynamic-NTK, YaRN, …), we delegate the ``inv_freq`` / + ``attention_scaling`` computation to HF's ``ROPE_INIT_FUNCTIONS`` so + that the frequency basis and scaling factor are identical to the + original checkpoint. This module has no learnable parameters. + """ + + def __init__(self, config): + super().__init__() + + if config.rope_scaling is not None: + # Delegate to HF's rope-scaling helpers for non-default types. + from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS + + rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type", "default")) + rope_init_fn = ROPE_INIT_FUNCTIONS[rope_type] + inv_freq, self.attention_scaling = rope_init_fn(config, device=None) + else: + dim = config.hidden_size // config.num_attention_heads + inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + self.attention_scaling = 1.0 + + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @torch.no_grad() + def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Generate cos/sin embeddings for given position ids. + + Args: + x: Input tensor (only used for dtype inference). + position_ids: Position indices, shape (batch_size, seq_len). + + Returns: + cos, sin: Rotary embeddings, each of shape (batch_size, seq_len, dim). + """ + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class BagelMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str = "silu", + ) -> None: + super().__init__() + self.gate_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + gather_output=False, + ) + self.up_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + gather_output=False, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + input_is_parallel=True, + bias=False, + ) + if hidden_act != "silu": + raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported.") + self.act_fn = nn.SiLU() + + def forward(self, x): + gate, _ = self.gate_proj(x) + up, _ = self.up_proj(x) + x = self.act_fn(gate) * up + x, _ = self.down_proj(x) + return x torch._dynamo.config.cache_size_limit = 512 @@ -157,24 +251,69 @@ class BaseNavitOutputWithPast(ModelOutput): past_key_values: NaiveCache | None = None -class PackedAttentionMoT(Qwen2Attention): +class PackedAttentionMoT(nn.Module): + """Packed attention with Mixture-of-Tokens routing for understanding/generation. + + Uses vLLM's QKVParallelLinear and RowParallelLinear for tensor parallelism + support, following the same pattern as vLLM's Qwen2Attention. + + The q/k/v projections are stacked into a single QKVParallelLinear: + - qkv_proj : stacks q_proj + k_proj + v_proj (understanding + gen text) + - qkv_proj_moe_gen : stacks q_proj_moe_gen + k_proj_moe_gen + v_proj_moe_gen (gen vae) + """ + def __init__(self, config, layer_idx: int | None = None): - super().__init__(config, layer_idx) - self.q_norm = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.q_norm_moe_gen = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm_moe_gen = Qwen2RMSNorm(self.head_dim, eps=config.rms_norm_eps) - + super().__init__() + self.layer_idx = layer_idx self.hidden_size = config.hidden_size - self.num_heads = config.num_attention_heads - self.num_key_value_heads = config.num_key_value_heads - self.head_dim = config.hidden_size // config.num_attention_heads - head_dim = self.head_dim - self.q_proj_moe_gen = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=True) - self.k_proj_moe_gen = nn.Linear(config.hidden_size, config.num_key_value_heads * head_dim, bias=True) - self.v_proj_moe_gen = nn.Linear(config.hidden_size, config.num_key_value_heads * head_dim, bias=True) - self.o_proj_moe_gen = nn.Linear(config.num_attention_heads * head_dim, config.hidden_size, bias=False) + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = self.hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + + # Understanding mode projections (stacked q/k/v) + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=True, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + ) + + # Generation mode MoE projections (stacked q/k/v) + self.qkv_proj_moe_gen = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=True, + ) + self.o_proj_moe_gen = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=False, + ) + + # QK normalization + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.q_norm_moe_gen = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm_moe_gen = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.rotary_op = RotaryEmbedding(is_neox_style=True) @@ -194,38 +333,43 @@ class PackedAttentionMoT(Qwen2Attention): packed_text_indexes=None, ): if mode == "und": - packed_query_states = self.q_proj(packed_query_sequence).view(-1, self.num_heads, self.head_dim) - packed_key_states = self.k_proj(packed_query_sequence).view(-1, self.num_key_value_heads, self.head_dim) - packed_value_states = self.v_proj(packed_query_sequence).view(-1, self.num_key_value_heads, self.head_dim) + qkv, _ = self.qkv_proj(packed_query_sequence) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + packed_query_states = q.view(-1, self.num_heads, self.head_dim) + packed_key_states = k.view(-1, self.num_kv_heads, self.head_dim) + packed_value_states = v.view(-1, self.num_kv_heads, self.head_dim) packed_query_states = self.q_norm(packed_query_states) packed_key_states = self.k_norm(packed_key_states) elif mode == "gen": packed_query_sequence = packed_query_sequence.to(torch.bfloat16) - packed_query_states = packed_query_sequence.new_zeros( - (packed_query_sequence.shape[0], self.num_heads * self.head_dim) - ) - packed_key_states = packed_query_sequence.new_zeros( - (packed_query_sequence.shape[0], self.num_key_value_heads * self.head_dim) - ) - packed_value_states = packed_query_sequence.new_zeros( - (packed_query_sequence.shape[0], self.num_key_value_heads * self.head_dim) - ) packed_text_query_sequence = packed_query_sequence[packed_text_indexes] packed_vae_query_sequence = packed_query_sequence[packed_vae_token_indexes] - packed_query_states[packed_text_indexes] = self.q_proj(packed_text_query_sequence) - packed_query_states[packed_vae_token_indexes] = self.q_proj_moe_gen(packed_vae_query_sequence) + # Project text tokens through base qkv + text_qkv, _ = self.qkv_proj(packed_text_query_sequence) + text_q, text_k, text_v = text_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - packed_key_states[packed_text_indexes] = self.k_proj(packed_text_query_sequence) - packed_key_states[packed_vae_token_indexes] = self.k_proj_moe_gen(packed_vae_query_sequence) + # Project vae tokens through moe_gen qkv + vae_qkv, _ = self.qkv_proj_moe_gen(packed_vae_query_sequence) + vae_q, vae_k, vae_v = vae_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - packed_value_states[packed_text_indexes] = self.v_proj(packed_text_query_sequence) - packed_value_states[packed_vae_token_indexes] = self.v_proj_moe_gen(packed_vae_query_sequence) + # Merge into packed tensors + total_len = packed_query_sequence.shape[0] + packed_query_states = packed_query_sequence.new_zeros((total_len, self.q_size)) + packed_key_states = packed_query_sequence.new_zeros((total_len, self.kv_size)) + packed_value_states = packed_query_sequence.new_zeros((total_len, self.kv_size)) + + packed_query_states[packed_text_indexes] = text_q + packed_query_states[packed_vae_token_indexes] = vae_q + packed_key_states[packed_text_indexes] = text_k + packed_key_states[packed_vae_token_indexes] = vae_k + packed_value_states[packed_text_indexes] = text_v + packed_value_states[packed_vae_token_indexes] = vae_v packed_query_states = packed_query_states.view(-1, self.num_heads, self.head_dim) - packed_key_states = packed_key_states.view(-1, self.num_key_value_heads, self.head_dim) - packed_value_states = packed_value_states.view(-1, self.num_key_value_heads, self.head_dim) + packed_key_states = packed_key_states.view(-1, self.num_kv_heads, self.head_dim) + packed_value_states = packed_value_states.view(-1, self.num_kv_heads, self.head_dim) packed_query_states = packed_query_states.to(torch.float32) packed_query_states[packed_text_indexes] = self.q_norm(packed_query_states[packed_text_indexes]) @@ -252,8 +396,8 @@ class PackedAttentionMoT(Qwen2Attention): past_value_states = past_key_values.value_cache[self.layer_idx] seqlens = sum(query_lens) + sum(key_values_lens) - merged_key_states = past_key_states.new_zeros(size=[seqlens, self.num_key_value_heads, self.head_dim]) - merged_value_states = past_key_states.new_zeros(size=[seqlens, self.num_key_value_heads, self.head_dim]) + merged_key_states = past_key_states.new_zeros(size=[seqlens, self.num_kv_heads, self.head_dim]) + merged_value_states = past_key_states.new_zeros(size=[seqlens, self.num_kv_heads, self.head_dim]) merged_key_states[packed_query_indexes] = packed_key_states merged_key_states[packed_key_value_indexes] = past_key_states merged_value_states[packed_query_indexes] = packed_value_states @@ -277,14 +421,16 @@ class PackedAttentionMoT(Qwen2Attention): max_seqlen_k=max(key_values_lens).item(), causal=is_causal, ) - packed_attn_output = packed_attn_output.reshape(-1, self.hidden_size) + packed_attn_output = packed_attn_output.reshape(-1, self.q_size) if mode == "und": - packed_attn_output = self.o_proj(packed_attn_output) + packed_attn_output, _ = self.o_proj(packed_attn_output) elif mode == "gen": - packed_attn_output[packed_text_indexes] = self.o_proj(packed_attn_output[packed_text_indexes]) - packed_attn_output[packed_vae_token_indexes] = self.o_proj_moe_gen( - packed_attn_output[packed_vae_token_indexes] - ) + text_out, _ = self.o_proj(packed_attn_output[packed_text_indexes]) + vae_out, _ = self.o_proj_moe_gen(packed_attn_output[packed_vae_token_indexes]) + full_output = text_out.new_zeros((packed_attn_output.shape[0], self.hidden_size)) + full_output[packed_text_indexes] = text_out + full_output[packed_vae_token_indexes] = vae_out + packed_attn_output = full_output if update_past_key_values: past_key_values.key_cache[self.layer_idx] = merged_key_states @@ -298,19 +444,19 @@ class Qwen2MoTDecoderLayer(nn.Module): self, config, layer_idx: int | None = None, - attn_module: Qwen2Attention | None = PackedAttentionMoT, + attn_module: type[nn.Module] | None = PackedAttentionMoT, ): super().__init__() self.hidden_size = config.hidden_size self.self_attn = attn_module(config, layer_idx) - self.mlp = Qwen2MLP(config) - self.mlp_moe_gen = Qwen2MLP(config) - self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.input_layernorm_moe_gen = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm_moe_gen = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.mlp = BagelMLP(config.hidden_size, config.intermediate_size, config.hidden_act) + self.mlp_moe_gen = BagelMLP(config.hidden_size, config.intermediate_size, config.hidden_act) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.input_layernorm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, @@ -391,7 +537,7 @@ class Qwen2MoTModel(Qwen2PreTrainedModel): self.vocab_size = config.vocab_size self.use_moe = "Mo" in config.layer_module - self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) self.layers = nn.ModuleList( [ Qwen2MoTDecoderLayer(config, layer_idx, attn_module=PackedAttentionMoT) @@ -399,10 +545,10 @@ class Qwen2MoTModel(Qwen2PreTrainedModel): ] ) - self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) if self.use_moe: - self.norm_moe_gen = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = Qwen2RotaryEmbedding(config=config) + self.norm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = BagelRotaryEmbedding(config=config) # Initialize weights and apply final processing self.post_init() @@ -528,6 +674,50 @@ class Qwen2MoTForCausalLM(Qwen2PreTrainedModel): return outputs + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights for vLLM parallel layers. + + Handles stacked parameter remapping for QKVParallelLinear: + - q_proj, k_proj, v_proj -> qkv_proj (shard ids: q, k, v) + - q_proj_moe_gen, k_proj_moe_gen, v_proj_moe_gen -> qkv_proj_moe_gen + Other parallel layers (gate_proj, up_proj, down_proj, embed_tokens, etc.) + keep HF checkpoint names and use weight_loader for TP sharding. + """ + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + # More specific _moe_gen patterns FIRST to avoid substring + # ambiguity (`.q_proj` is a substring of `.q_proj_moe_gen`). + (".qkv_proj_moe_gen", ".q_proj_moe_gen", "q"), + (".qkv_proj_moe_gen", ".k_proj_moe_gen", "k"), + (".qkv_proj_moe_gen", ".v_proj_moe_gen", "v"), + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + ] + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict.get(name) + if param is None: + break + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight, shard_id) + break + else: + param = params_dict.get(name) + if param is None: + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + return loaded_params + def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0): grid_h = np.arange(grid_size, dtype=np.float32) @@ -639,7 +829,7 @@ def get_flattened_position_ids_extrapolate(img_h, img_w, patch_size, max_num_pat return pos_ids -class Bagel(torch.nn.Module): +class Bagel(nn.Module): config_class = BagelConfig base_model_prefix = "bagel" diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index bdb9f1f5c..040f7ddec 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -402,7 +402,7 @@ class BagelPipeline(nn.Module): ) # Fail fast with a clear error instead of CUDA gather OOB. max_tid = int(generation_input["packed_text_ids"].max().item()) - emb_n = int(self.language_model.model.embed_tokens.weight.shape[0]) + emb_n = int(self.language_model.vocab_size) if max_tid >= emb_n: raise ValueError( "Tokenizer/model vocab mismatch: max token id " @@ -438,7 +438,7 @@ class BagelPipeline(nn.Module): ) # Fail fast for special tokens used by the image path as well. max_tid_img = int(generation_input["packed_text_ids"].max().item()) - emb_n = int(self.language_model.model.embed_tokens.weight.shape[0]) + emb_n = int(self.language_model.vocab_size) if max_tid_img >= emb_n: raise ValueError( "Tokenizer/model vocab mismatch (image path): max token id " @@ -483,6 +483,28 @@ class BagelPipeline(nn.Module): allowed = set(state.keys()) shapes = {k: tuple(v.shape) for k, v in state.items()} + tp_aware_params = {name for name, p in self.named_parameters() if hasattr(p, "weight_loader")} + + # Expand allowed/tp_aware_params with stacked param source names. + # QKVParallelLinear merges q_proj+k_proj+v_proj into qkv_proj; the + # checkpoint stores the original separate names. We must recognise + # those names so _filtered_weights does not drop them. + _stacked_expansions = [ + (".qkv_proj", ".q_proj"), + (".qkv_proj", ".k_proj"), + (".qkv_proj", ".v_proj"), + (".qkv_proj_moe_gen", ".q_proj_moe_gen"), + (".qkv_proj_moe_gen", ".k_proj_moe_gen"), + (".qkv_proj_moe_gen", ".v_proj_moe_gen"), + ] + stacked_source_names: set[str] = set() + for name in list(allowed): + for target_suffix, source_suffix in _stacked_expansions: + if target_suffix in name: + stacked_source_names.add(name.replace(target_suffix, source_suffix)) + allowed.update(stacked_source_names) + tp_aware_params.update(stacked_source_names) + def _normalize_name(name: str) -> str: # Common wrappers/prefixes in checkpoints. for pfx in ("module.", "model."): @@ -536,7 +558,7 @@ class BagelPipeline(nn.Module): for cand in _iter_candidate_names(name): if cand in allowed: # Only accept if tensor shape matches target param/buffer shape. - if tuple(tensor.shape) == shapes.get(cand): + if tuple(tensor.shape) == shapes.get(cand) or cand in tp_aware_params: picked = cand break else: From 8a9644c8213b6a7234f6cf781d19ef417e301e1a Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Wed, 11 Feb 2026 05:18:36 +0800 Subject: [PATCH 12/13] [Bugfix] Fix image edit RoPE crash when explicit height/width are provided (#1265) Signed-off-by: lishunyang Co-authored-by: Hongsheng Liu --- .../image_to_image/openai_chat_client.py | 4 ++++ .../image_to_image/run_curl_image_edit.sh | 17 ++++++++--------- .../qwen_image/pipeline_qwen_image_edit.py | 4 ++-- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/examples/online_serving/image_to_image/openai_chat_client.py b/examples/online_serving/image_to_image/openai_chat_client.py index 0fe4b0ede..7aac2b214 100644 --- a/examples/online_serving/image_to_image/openai_chat_client.py +++ b/examples/online_serving/image_to_image/openai_chat_client.py @@ -76,6 +76,10 @@ def edit_image( # Build extra_body with generation parameters extra_body = {} + if height is not None: + extra_body["height"] = height + if width is not None: + extra_body["width"] = width if steps is not None: extra_body["num_inference_steps"] = steps if guidance_scale is not None: diff --git a/examples/online_serving/image_to_image/run_curl_image_edit.sh b/examples/online_serving/image_to_image/run_curl_image_edit.sh index 748a0ebe5..1a5328685 100644 --- a/examples/online_serving/image_to_image/run_curl_image_edit.sh +++ b/examples/online_serving/image_to_image/run_curl_image_edit.sh @@ -19,26 +19,25 @@ if [[ ! -f "$INPUT_IMG" ]]; then exit 1 fi -IMG_B64=$(base64 -w0 "$INPUT_IMG") +REQUEST_JSON_FILE=$(mktemp) +trap 'rm -f "$REQUEST_JSON_FILE"' EXIT -REQUEST_JSON=$( - jq -n --arg prompt "$PROMPT" --arg img "$IMG_B64" '{ +# Pipe base64 into jq via stdin to avoid ARG_MAX limit on large images +base64 -w0 "$INPUT_IMG" \ + | jq -Rs --arg prompt "$PROMPT" '{ messages: [{ role: "user", content: [ {"type": "text", "text": $prompt}, - {"type": "image_url", "image_url": {"url": ("data:image/png;base64," + $img)}} + {"type": "image_url", "image_url": {"url": ("data:image/png;base64," + .)}} ] }], extra_body: { - height: 1024, - width: 1024, num_inference_steps: 50, guidance_scale: 1, seed: 42 } - }' -) + }' > "$REQUEST_JSON_FILE" echo "Generating edited image..." echo "Server: $SERVER" @@ -48,7 +47,7 @@ echo "Output: $OUTPUT" curl -s "$SERVER/v1/chat/completions" \ -H "Content-Type: application/json" \ - -d "$REQUEST_JSON" \ + -d @"$REQUEST_JSON_FILE" \ | jq -r '.choices[0].message.content[0].image_url.url' \ | cut -d',' -f2 \ | base64 -d > "$OUTPUT" diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index 082c9e754..bff248a44 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py @@ -110,9 +110,9 @@ def get_qwen_image_edit_pre_process_func( if image is not None and not ( isinstance(image, torch.Tensor) and len(image.shape) > 1 and image.shape[1] == latent_channels ): - image = image_processor.resize(image, height, width) + image = image_processor.resize(image, calculated_height, calculated_width) prompt_image = image - image = image_processor.preprocess(image, height, width) + image = image_processor.preprocess(image, calculated_height, calculated_width) image = image.unsqueeze(2) # Store preprocessed image and prompt image in request From dc8b28e0aa6d1b4df85a93d3bdcb30280ca31975 Mon Sep 17 00:00:00 2001 From: Alicia <115451386+congw729@users.noreply.github.com> Date: Wed, 11 Feb 2026 06:00:43 +0800 Subject: [PATCH 13/13] [Doc] Sync (#1216) Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com> --- .github/PULL_REQUEST_TEMPLATE.md | 2 +- docs/.nav.yml | 8 ++++---- docs/api/README.md | 7 ++----- .../examples/offline_inference/text_to_image.md | 3 +-- docs/user_guide/examples/online_serving/image_to_image.md | 3 +++ 5 files changed, 11 insertions(+), 12 deletions(-) diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index ea9b94c93..2466bc71e 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -14,7 +14,7 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT - [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)". - [ ] The test plan, such as providing test command. - [ ] The test results, such as pasting the results comparison before and after, or e2e results -- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. +- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.** - [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft. diff --git a/docs/.nav.yml b/docs/.nav.yml index e0eb6d674..ae4a58f5c 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -15,20 +15,20 @@ nav: - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md - - LoRA Inference Examples: user_guide/examples/offline_inference/lora_inference.md + - LoRA Inference(Diffusion): user_guide/examples/offline_inference/lora_inference.md - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md - - Qwen3-TTS Offline Inference: user_guide/examples/offline_inference/qwen3_tts.md + - Qwen3-TTS: user_guide/examples/offline_inference/qwen3_tts.md - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md - Online Serving: - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md - Image-To-Image: user_guide/examples/online_serving/image_to_image.md - - Online LoRA Inference (Diffusion): user_guide/examples/online_serving/lora_inference.md + - LoRA Inference(Diffusion): user_guide/examples/online_serving/lora_inference.md - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md - - Qwen3-TTS Online Serving: user_guide/examples/online_serving/qwen3_tts.md + - Qwen3-TTS: user_guide/examples/online_serving/qwen3_tts.md - Text-To-Image: user_guide/examples/online_serving/text_to_image.md - General: - usage/* diff --git a/docs/api/README.md b/docs/api/README.md index eb3a40950..b4ec398f2 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -104,15 +104,12 @@ Worker classes and model runners for distributed inference. - [vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker][] - [vllm_omni.platforms.xpu.worker.xpu_generation_model_runner.XPUGenerationModelRunner][] - [vllm_omni.platforms.xpu.worker.xpu_generation_worker.XPUGenerationWorker][] +- [vllm_omni.worker.base.OmniGPUWorkerBase][] - [vllm_omni.worker.gpu_ar_model_runner.ExecuteModelState][] - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][] - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][] - [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][] - [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][] +- [vllm_omni.worker.gpu_memory_utils.parse_cuda_visible_devices][] - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][] - [vllm_omni.worker.mixins.OmniWorkerMixin][] - - -## Metrics - -- [vllm_omni.metrics.OrchestratorAggregator][] diff --git a/docs/user_guide/examples/offline_inference/text_to_image.md b/docs/user_guide/examples/offline_inference/text_to_image.md index 2e8739949..c0dbfd141 100644 --- a/docs/user_guide/examples/offline_inference/text_to_image.md +++ b/docs/user_guide/examples/offline_inference/text_to_image.md @@ -51,8 +51,7 @@ if __name__ == "__main__": For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support - batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md - ) to let the model accept a longer batch of prompts. + batched inputs, you can [modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a longer batch of prompts. Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM. This is useful for models that support negative prompts. diff --git a/docs/user_guide/examples/online_serving/image_to_image.md b/docs/user_guide/examples/online_serving/image_to_image.md index b89f0fe82..6be2a4a7e 100644 --- a/docs/user_guide/examples/online_serving/image_to_image.md +++ b/docs/user_guide/examples/online_serving/image_to_image.md @@ -15,6 +15,9 @@ For **multi-image** input editing, use **Qwen-Image-Edit-2509** (QwenImageEditPl vllm serve Qwen/Qwen-Image-Edit --omni --port 8092 ``` +!!! note + If you encounter Out-of-Memory (OOM) issues or have limited GPU memory, you can enable VAE slicing and tiling to reduce memory usage, --vae-use-slicing --vae-use-tiling + ### Multi-Image Edit (Qwen-Image-Edit-2509) ```bash