Merge remote-tracking branch 'origin/main' into dev/vllm-align

2026-02-11 01:44:54 +00:00
commit 279b6b5050
@@ -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:
@@ -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
@@ -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
@@ -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.
 </details>

@@ -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/*

@@ -8,26 +8,27 @@ 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:
      - 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/*
@@ -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][]
@@ -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 |

 ---
@@ -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 <model> --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
+```
@@ -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
@@ -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.
@@ -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.
@@ -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
@@ -1,9 +1,11 @@
-# Qwen3-TTS Online Serving
+# Qwen3-TTS

 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_tts>.


-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

@@ -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                      |
@@ -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")

@@ -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:
@@ -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:
@@ -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"
@@ -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
@@ -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",
@@ -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 \
@@ -8,7 +8,7 @@ repo_url: https://github.com/vllm-project/vllm-omni
 edit_uri: edit/main/docs/

 # Copyright
-copyright: Copyright &copy; 2025 vLLM-Omni Team
+copyright: Copyright &copy; 2026 vLLM-Omni Team

 # Theme
 theme:
@@ -192,6 +192,6 @@ extend-ignore-identifiers-re = [
    ".*nothink.*",
    ".*NOTHINK.*",
    ".*nin.*",
-    "Ono_Anna",
+    ".*[Oo]no_[Aa]nna.*",
    ".*cann.*",
 ]
@@ -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)
@@ -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})"
+    )
@@ -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
@@ -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
 # =============================================================================
@@ -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)
@@ -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
@@ -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]
@@ -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"

@@ -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:
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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 (
@@ -10,7 +10,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
@@ -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:
@@ -973,6 +973,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}"

@@ -1049,6 +1050,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:
@@ -1131,6 +1133,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())}"
@@ -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.
@@ -500,9 +500,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",
    }
@@ -168,6 +168,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
@@ -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,
@@ -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(