Revert "[Quantization] Enable FP8 online quantization for Z-image text encode…"

This reverts commit ac66282746.

.buildkite/test-nightly.yml

@@ -297,7 +297,7 @@ steps:
       - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100"
         timeout_in_minutes: 120
         commands:
-          - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy
+          - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy
         agents:
           queue: "mithril-h100-pool"
         plugins:
@@ -332,10 +332,48 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Diffusion X2I(&A&T) · Bagel Function Test with H100"
+        timeout_in_minutes: 120
+        commands:
+          - pytest -sv tests/e2e/online_serving/test_bagel_expansion.py -m "full_model and diffusion and H100" --run-level "full_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 3
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
       - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
         timeout_in_minutes: 60
         commands:
-          - pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
+          - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
         agents:
           queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
         plugins:

.gitattributes

@@ -27,6 +27,7 @@
 *.mov binary
 *.mp4 binary
 *.mp3 binary
+*.wav binary
 *.flv binary
 *.fla binary
 *.swf binary

.gitignore

@@ -264,6 +264,8 @@ tmp_test
 vllm_omni/_version.py
 # output files
 *.wav
+# Vendored test reference audio (small, intentionally checked in)
+!tests/assets/**/*.wav
 .worktrees/
 # CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
 tests/.ci_generated/

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.19.0
+ARG VLLM_BASE_TAG=v0.20.0-cu130
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docs/.nav.yml

@@ -10,6 +10,7 @@ nav:
       - Image Generation: serving/image_generation_api.md
       - Image Edit: serving/image_edit_api.md
       - Text to Speech: serving/speech_api.md
+      - Audio Generation: serving/audio_generate_api.md
       - Streaming Video Input: serving/video_stream_api.md
   - Examples:
     - examples/README.md
@@ -62,11 +63,6 @@ nav:
       - Cache Acceleration:
         - TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
         - Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
-      - Quantization:
-        - Overview: user_guide/diffusion/quantization/overview.md
-        - FP8: user_guide/diffusion/quantization/fp8.md
-        - Int8: user_guide/diffusion/quantization/int8.md
-        - GGUF: user_guide/diffusion/quantization/gguf.md
       - Attention Backends: user_guide/diffusion/attention_backends.md
       - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
       - Parallelism:
@@ -81,6 +77,14 @@ nav:
       - LoRA: user_guide/diffusion/lora.md
       - Custom Pipeline: features/custom_pipeline.md
       - Step Execution: user_guide/diffusion/step_execution.md
+    - Quantization:
+      - Overview: user_guide/quantization/overview.md
+      - Online Quantization: user_guide/quantization/online.md
+      - FP8 W8A8: user_guide/quantization/fp8.md
+      - Int8 W8A8: user_guide/quantization/int8.md
+      - GGUF: user_guide/quantization/gguf.md
+      - AutoRound: user_guide/quantization/autoround.md
+      - msModelSlim: user_guide/quantization/msmodelslim.md
     - ComfyUI: features/comfyui.md
 - Developer Guide:
   - General:
@@ -92,7 +96,6 @@ nav:
     - contributing/model/adding_omni_model.md
     - contributing/model/adding_tts_model.md
     - contributing/model/adding_diffusion_model.md
-    - contributing/model/adding_quantization_model.md
   - CI: contributing/ci
   - Design Documents:
     - design/index.md

docs/contributing/ci/test_examples/l4_functionality_tests.inc.md

@@ -1,6 +1,6 @@
 **Scope**
 
-For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
+For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
 
 **Test Case Design**
 

docs/contributing/model/adding_quantization_model.md

@@ -1,478 +0,0 @@
-# Support FP8 Quantization
-
-This section describes how to add FP8 quantization support to a diffusion model. We use the Qwen-Image pipeline as the reference implementation.
-
----
-
-## Table of Contents
-
-- [Overview](#overview)
-- [Step-by-Step Implementation](#step-by-step-implementation)
-- [Testing](#testing)
-- [Troubleshooting](#troubleshooting)
-- [Reference Implementations](#reference-implementations)
-- [Summary](#summary)
-
----
-
-## Overview
-
-### What is FP8 Quantization?
-
-FP8 quantization reduces model memory footprint by storing weights in 8-bit floating point (`float8_e4m3fn`) instead of 16-bit (`bfloat16`/`float16`). This saves ~50% of weight memory.
-
-A diffusion pipeline typically has three quantizable components:
-
-| Component | Layer Types | Quantization Approach |
-|-----------|-----------|----------------------|
-| **DiT (transformer)** | `nn.Linear` | Replace with vLLM quantized linear layers |
-| **Text encoder** | `nn.Linear` | FP8 weight storage via hooks |
-| **VAE** | `nn.Conv2d`, `nn.Conv3d` | FP8 weight storage via hooks |
-
-### Architecture
-
-The quantization system has two mechanisms:
-
-**1. vLLM Quantized Linear Layers (for DiT)**
-
-vLLM provides drop-in replacements for `nn.Linear` that perform FP8 W8A8 computation on supported hardware. These are configured via `quant_config` passed to layer constructors.
-
-Key classes:
-
-| Class | Location | Purpose |
-|-------|----------|---------|
-| `build_quant_config` | `vllm_omni/quantization/__init__.py` | Build quantization config from string/dict/per-component spec |
-| `ComponentQuantizationConfig` | `vllm_omni/quantization/component_config.py` | Per-component routing by layer prefix |
-
-**2. FP8 Weight Storage via Hooks (for text encoder and VAE)**
-
-For models loaded via `from_pretrained()` (outside vLLM's weight pipeline), we use forward hooks to store weights in FP8 and dequantize on-the-fly:
-
-```
-At rest:     weight.data = fp8_tensor (half memory)
-Pre-hook:    weight.data = fp8_weight.to(bf16) * scale
-Forward:     Normal computation in BF16
-Post-hook:   weight.data = fp8_weight (back to FP8)
-```
-
-Key function:
-
-| Function | Location | Purpose |
-|----------|----------|---------|
-| `apply_fp8_weight_storage` | `vllm_omni/diffusion/models/utils.py` | Apply FP8 hooks to Linear/Conv layers (introduced in [#1412](https://github.com/vllm-project/vllm-omni/pull/1412)) |
-
----
-
-## Step-by-Step Implementation
-
-### Step 1: Add `quant_config` to the Transformer
-
-Pass `quant_config` to vLLM parallel linear layers in the transformer. This enables FP8 W8A8 for DiT linear layers.
-
-#### 1.1: Accept `quant_config` in Constructor
-
-```python
-from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
-
-class YourTransformer2DModel(nn.Module):
-    def __init__(
-        self,
-        *,
-        od_config: OmniDiffusionConfig,
-        quant_config: QuantizationConfig | None = None,
-        # ... other params
-    ):
-        self.quant_config = quant_config
-        # Pass to sub-modules
-        self.blocks = nn.ModuleList([
-            YourTransformerBlock(quant_config=quant_config, ...)
-            for _ in range(num_layers)
-        ])
-```
-
-#### 1.2: Pass `quant_config` to vLLM Linear Layers
-
-Replace `nn.Linear` with vLLM parallel linear layers that accept `quant_config`:
-
-```python
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    QKVParallelLinear,
-    RowParallelLinear,
-)
-
-class YourAttentionBlock(nn.Module):
-    def __init__(self, hidden_size, num_heads, quant_config=None, prefix=""):
-        # QKV projection
-        self.to_qkv = QKVParallelLinear(
-            hidden_size=hidden_size,
-            head_size=hidden_size // num_heads,
-            total_num_heads=num_heads,
-            quant_config=quant_config,
-            prefix=f"{prefix}.to_qkv",
-        )
-        # Output projection
-        self.to_out = RowParallelLinear(
-            input_size=hidden_size,
-            output_size=hidden_size,
-            quant_config=quant_config,
-            prefix=f"{prefix}.to_out",
-        )
-
-class YourFeedForward(nn.Module):
-    def __init__(self, hidden_size, intermediate_size, quant_config=None, prefix=""):
-        self.gate_up_proj = ColumnParallelLinear(
-            input_size=hidden_size,
-            output_size=intermediate_size * 2,
-            quant_config=quant_config,
-            prefix=f"{prefix}.gate_up_proj",
-        )
-        self.down_proj = RowParallelLinear(
-            input_size=intermediate_size,
-            output_size=hidden_size,
-            quant_config=quant_config,
-            prefix=f"{prefix}.down_proj",
-        )
-```
-
-!!! important
-    Always pass the `prefix` parameter. It is used by `ignored_layers` matching to skip sensitive layers.
-
-#### 1.3: Identify Sensitive Layers
-
-Some layers degrade output quality when quantized. Common sensitive layers:
-
-| Layer | Why Sensitive | Typical Models | Quant Methods Affected |
-|-------|-------------|---------------|----------------------|
-| `img_mlp` | Processes denoising latents with shifting dynamic range | Qwen-Image | FP8, Int8 |
-| `feed_forward` | FFN layers in DiT blocks, large dynamic range | Z-Image | Int8 |
-| `proj_out` | Final output projection, small quantization errors amplified | Various | FP8, Int8 |
-| `lm_head` | Output vocabulary projection, precision-critical | Omni (thinker) | NVFP4 |
-| `mlp.gate` | MoE router gates, routing decisions are precision-critical | MoE models | NVFP4 |
-
-**Quantization quality benchmarks** (LPIPS: lower = better, <0.01 imperceptible, >0.1 noticeable):
-
-| Model | Method | All Layers | With `ignored_layers` | Recommendation |
-|-------|--------|-----------|----------------------|----------------|
-| Qwen-Image-2512 | Int8 | LPIPS 0.0197 | 0.0027 (skip `img_mlp`) | Skip `img_mlp` |
-| Z-Image-Turbo | Int8 | LPIPS 0.1597 | 0.0290 (skip `feed_forward`) | Skip `feed_forward` |
-| Z-Image-Turbo | FP8 | LPIPS ~0.005 | — | All layers OK |
-
-To identify sensitive layers for a new model:
-
-1. Run inference with all layers quantized
-2. Compare output quality with BF16 baseline using LPIPS (see `benchmarks/diffusion/quantization_quality.py`)
-3. Selectively disable quantization per layer (via `ignored_layers`) until quality matches
-4. Document the recommended `ignored_layers` in the model's supported models table
-
-### Step 2: Wire Up Quantization in the Pipeline
-
-#### 2.1: Create `quant_config` and Pass to Transformer
-
-In the pipeline's `__init__`, extract the vLLM quant config and pass it to the transformer:
-
-```python
-from vllm_omni.quantization import build_quant_config
-
-class YourPipeline(nn.Module):
-    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
-        super().__init__()
-        # ... load scheduler, text_encoder, vae ...
-
-        # Build quant config from OmniDiffusionConfig
-        quant_config = build_quant_config(od_config.quantization_config)
-
-        # Pass to transformer
-        self.transformer = YourTransformer2DModel(
-            od_config=od_config,
-            quant_config=quant_config,
-            # ... other kwargs
-        )
-```
-
-#### 2.2: Apply FP8 Weight Storage to Text Encoder and VAE
-
-For models loaded via `from_pretrained()`, apply hook-based FP8 weight storage.
-
-!!! note
-    The `apply_fp8_weight_storage` function is provided by `vllm_omni/diffusion/models/utils.py`, introduced in [PR #1412](https://github.com/vllm-project/vllm-omni/pull/1412).
-
-```python
-from vllm_omni.diffusion.models.utils import apply_fp8_weight_storage
-
-class YourPipeline(nn.Module):
-    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
-        super().__init__()
-
-        self.text_encoder = SomeTextEncoder.from_pretrained(...)
-        self.vae = SomeVAE.from_pretrained(...).to(self.device)
-
-        # Apply FP8 weight storage to text encoder and VAE
-        if (
-            od_config.quantization_config is not None
-            and getattr(od_config.quantization_config, "quant_method", None) == "fp8"
-        ):
-            apply_fp8_weight_storage(self.vae)
-            apply_fp8_weight_storage(self.text_encoder)
-
-        # ... rest of init
-```
-
-`apply_fp8_weight_storage` walks all `nn.Linear`, `nn.Conv2d`, and `nn.Conv3d` modules in the model and applies FP8 hooks to each.
-
-#### 2.3: Fix `load_weights` to Include Pre-loaded Components
-
-Since text encoders and VAEs are loaded via `from_pretrained()` (not through vLLM's weight pipeline), their parameters won't appear in the `loaded_weights` set. Mark them explicitly:
-
-```python
-def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-    loader = AutoWeightsLoader(self)
-    loaded_weights = loader.load_weights(weights)
-    # Mark pre-loaded component weights as loaded
-    loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
-    loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
-    return loaded_weights
-```
-
-### Step 3: Add to Supported Models Documentation
-
-Update the FP8 supported models table in `docs/user_guide/diffusion/quantization/fp8.md`:
-
-```markdown
-| Your-Model | `org/Your-Model` | ✅ | ✅ | ✅ | `sensitive_layer` |
-```
-
----
-
-## Testing
-
-### Quick Verification
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# Test with FP8 enabled
-omni = Omni(model="your-model", quantization="fp8")
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-### Quality Gate Test (LPIPS)
-
-We provide a pytest-based quality gate at `tests/diffusion/quantization/test_quantization_quality.py`.
-It generates outputs with both BF16 and your quantized method using the same seed, computes
-[LPIPS](https://github.com/richzhang/PerceptualSimilarity) perceptual distance, and **fails if it
-exceeds a threshold**. This is the recommended way to validate that a quantization method does not
-unacceptably degrade output quality.
-
-**Step 1: Add your test config** — edit `QUALITY_CONFIGS` in the test file:
-
-```python
-QualityTestConfig(
-    id="int8_z_image",              # pytest ID
-    model="Tongyi-MAI/Z-Image-Turbo",
-    quantization="int8",            # your quantization method
-    task="t2i",                     # "t2i" or "t2v"
-    prompt="a cup of coffee on a wooden table, morning light",
-    max_lpips=0.15,                 # fail threshold (see guide below)
-    num_inference_steps=20,
-),
-```
-
-**Step 2: Run the test:**
-
-```bash
-# Install dependency
-pip install lpips
-
-# Run all quality tests
-pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""
-
-# Run only your method
-pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "int8"
-
-# Run a specific model
-pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
-```
-
-**Step 3: Read the output:**
-
-```
-Quantization Quality: int8_z_image
-  Model:         Tongyi-MAI/Z-Image-Turbo
-  Method:        int8
-  LPIPS:         0.0290  (threshold: 0.15)
-  BF16 memory:   19.15 GiB
-  Quant memory:  14.80 GiB  (23% reduction)
-  Result:        PASS
-```
-
-**LPIPS threshold guide:**
-
-| LPIPS Range | Meaning | Recommendation |
-|-------------|---------|----------------|
-| < 0.01 | Imperceptible | Ideal — no visible difference |
-| 0.01 – 0.05 | Minor differences | Good — acceptable for most use cases |
-| 0.05 – 0.15 | Noticeable on close inspection | Acceptable with `ignored_layers` tuning |
-| 0.15 – 0.30 | Clearly visible | Investigate — likely needs `ignored_layers` |
-| > 0.30 | Significant degradation | Quantization method may not suit this model |
-
-!!! tip
-    If your test fails, don't raise the threshold — instead, identify sensitive layers and add
-    them to `ignored_layers` (see next section). Then re-run the test with the lower LPIPS.
-
-### Quality Benchmark (Detailed)
-
-For a full multi-prompt benchmark with Markdown output (suitable for PR descriptions), use the
-standalone benchmark script:
-
-```bash
-python benchmarks/diffusion/quantization_quality.py \
-    --model Tongyi-MAI/Z-Image-Turbo \
-    --quantization fp8 \
-    --prompts \
-        "a cup of coffee on the table" \
-        "an aerial view of a coral reef" \
-    --height 1024 --width 1024 \
-    --num-inference-steps 50 --seed 42
-```
-
-This generates BF16 and quantized outputs, computes per-prompt LPIPS, and prints a Markdown
-summary table ready to paste into your PR.
-
-### Identifying Sensitive Layers with `ignored_layers`
-
-Not all layers tolerate quantization equally. **`ignored_layers`** lets you keep specific layers
-in BF16 while quantizing everything else. This is critical for achieving good quality with
-aggressive quantization methods (Int8, NVFP4).
-
-**Why some layers are sensitive:**
-
-| Layer Type | Why Sensitive | Typical Impact |
-|------------|--------------|----------------|
-| `img_mlp` | Processes denoising latents with shifting dynamic range across timesteps. No built-in normalization to absorb quantization error. | Color shifts, blurriness |
-| `feed_forward` | FFN layers with large dynamic range in DiT blocks. | Artifacts, loss of detail |
-| `proj_out` | Final output projection — small errors get amplified since there are no subsequent layers to correct them. | Overall quality loss |
-| `lm_head` | Vocabulary projection in LLMs — precision-critical for token selection. | Garbage text output |
-| `mlp.gate` | MoE router gates — routing decisions are binary and precision-critical. | Wrong expert selection, quality collapse |
-
-**How to find sensitive layers for your model:**
-
-1. Run the quality gate test with all layers quantized:
-    ```bash
-    pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "your_model"
-    ```
-
-2. If LPIPS exceeds your threshold, try skipping common sensitive layers one at a time:
-    ```python
-    QualityTestConfig(
-        id="fp8_your_model_skip_mlp",
-        model="your-org/Your-Model",
-        quantization="fp8",
-        task="t2i",
-        prompt="a cup of coffee on a wooden table",
-        max_lpips=0.05,
-    ),
-    ```
-
-    And in the test, modify the Omni call:
-    ```python
-    omni_qt = Omni(
-        model=config.model,
-        quantization_config={
-            "method": config.quantization,
-            "ignored_layers": ["img_mlp"],
-        },
-    )
-    ```
-
-3. Compare LPIPS with and without `ignored_layers` to isolate the problematic layers.
-
-4. Document the recommended `ignored_layers` in the supported models table.
-
-**Real-world examples:**
-
-| Model | Method | All Layers (LPIPS) | With `ignored_layers` | Which Layers to Skip |
-|-------|--------|-------------------|----------------------|---------------------|
-| Qwen-Image-2512 | Int8 | 0.0197 | 0.0027 | `img_mlp` |
-| Z-Image-Turbo | Int8 | 0.1597 | 0.0290 | `feed_forward` |
-| Z-Image-Turbo | FP8 | ~0.005 | — | None needed |
-
-### What to Verify (Checklist)
-
-1. **Model loads without errors** — Check logs for quantization initialization messages
-2. **Quality gate passes** — LPIPS within threshold vs BF16 baseline
-3. **Memory reduction** — Expect ~30-50% savings on quantized components
-4. **Sensitive layers identified** — Document which layers (if any) need `ignored_layers`
-5. **`ignored_layers` pattern works** — Verify that the layer name prefix matches your transformer's module names
-
----
-
-## Troubleshooting
-
-### Issue: Output Quality Degradation
-
-**Symptoms:** Generated images have visible artifacts, color shifts, or blurriness compared to BF16 baseline.
-
-**Causes & Solutions:**
-
-- **Sensitive layers being quantized:**
-
-    **Solution:** Identify and skip sensitive layers:
-    ```python
-    omni = Omni(
-        model="your-model",
-        quantization_config={
-            "method": "fp8",
-            "ignored_layers": ["img_mlp"],
-        },
-    )
-    ```
-
-- **All layers sensitive (rare):**
-
-    **Solution:** The model may not be suitable for FP8 quantization. Consider weight-only quantization or other methods.
-
-### Issue: Missing Weights Warning
-
-**Symptoms:** Warnings about unloaded weights for `vae.*` or `text_encoder.*` parameters.
-
-**Solution:** Ensure `load_weights` marks pre-loaded component weights. See [Step 2.3](#23-fix-load_weights-to-include-pre-loaded-components).
-
-### Issue: CUDA Out of Memory
-
-**Symptoms:** OOM errors during model loading despite FP8 being enabled.
-
-**Causes & Solutions:**
-
-- **Text encoder/VAE loaded in BF16 before FP8 conversion:**
-
-    FP8 weight storage is applied after `from_pretrained()`, so peak load memory is still BF16. This is a one-time cost during initialization.
-
----
-
-## Reference Implementations
-
-| Model | Pipeline File | Transformer File | Notes |
-|-------|-------------|-----------------|-------|
-| **Z-Image** | `models/z_image/pipeline_z_image.py` | `models/z_image/z_image_transformer.py` | DiT + text encoder FP8 |
-| **Qwen-Image** | `models/qwen_image/pipeline_qwen_image.py` | `models/qwen_image/qwen_image_transformer.py` | DiT + text encoder + VAE FP8 |
-
-All files are under `vllm_omni/diffusion/`.
-
----
-
-## Summary
-
-Adding FP8 quantization support to a new model:
-
-1. ✅ **Add `quant_config` to transformer** — Accept and pass to vLLM parallel linear layers with proper `prefix`
-2. ✅ **Identify sensitive layers** — Profile which layers need `ignored_layers` for quality
-3. ✅ **Apply FP8 weight storage** — Call `apply_fp8_weight_storage()` on text encoder and VAE in pipeline `__init__`
-4. ✅ **Fix `load_weights`** — Mark pre-loaded component weights as loaded
-5. ✅ **Update docs** — Add model to the supported models table in `docs/user_guide/diffusion/quantization/fp8.md`
-6. ✅ **Test** — Verify quality, memory savings, and sensitive layer behavior

docs/getting_started/installation/gpu/cuda.inc.md

@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library
 
 ```bash
-uv pip install vllm==0.19.0 --torch-backend=auto
+uv pip install vllm==0.20.0 --torch-backend=auto
 ```
 
-The 0.19.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.20.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.19.0
+git checkout v0.20.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.19.0/vllm-0.19.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl
+#For CUDA 13.0 (the default for v0.20.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.20.0/vllm-0.20.0-cp38-abi3-manylinux_2_35_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -80,7 +80,7 @@ uv pip install --editable .
 Install vllm with command below (If you already have PyTorch).
 ```bash
 python use_existing_torch.py
-uv pip install -r requirements/build.txt
+uv pip install -r requirements/build/cuda.txt
 uv pip install --no-build-isolation --editable .
 ```
 </details>

docs/getting_started/installation/gpu/rocm.inc.md

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
+uv pip install vllm==0.20.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm700
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library
 
 ```bash
-uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
+uv pip install vllm==0.20.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm700
 ```
 
-The pre-built 0.19.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.20.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.19.0
+git checkout v0.20.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.19.0 \
+  vllm/vllm-omni-rocm:v0.20.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.19.0
+  vllm/vllm-omni-rocm:v0.20.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.19.0 --torch-backend=auto
+uv pip install vllm==0.20.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
+uv pip install vllm==0.20.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm700
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni

docs/serving/audio_generate_api.md

@@ -0,0 +1,338 @@
+# Audio Generate API
+
+vLLM-Omni provides an API for text-to-audio generation using diffusion-based models such as Stable Audio.
+
+Unlike the [Speech API](speech_api.md) which targets text-to-speech synthesis, the Audio Generate API is designed for general-purpose audio generation from text descriptions (sound effects, music, ambient soundscapes, etc.).
+
+Each server instance runs a single model (specified at startup via `vllm-omni serve <model> --omni`).
+
+## Quick Start
+
+### Start the Server
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8091 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+### Generate Audio
+
+**Using curl:**
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+**Using Python:**
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/generate",
+    json={
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0,
+    },
+    timeout=300.0,
+)
+
+with open("cat.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+Content-Type: application/json
+```
+
+### Request Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
+
+#### Diffusion Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `audio_length` | float | null | Audio duration in seconds (default value is the max ~47s for `stable-audio-open-1.0`) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducible generation |
+
+### Response Format
+
+Returns binary audio data with the appropriate `Content-Type` header:
+
+| `response_format` | Content-Type |
+|--------------------|--------------|
+| `wav` | `audio/wav` |
+| `mp3` | `audio/mpeg` |
+| `flac` | `audio/flac` |
+| `pcm` | `audio/pcm` |
+| `aac` | `audio/aac` |
+| `opus` | `audio/opus` |
+
+## Examples
+
+### Basic Generation
+
+Generate audio with only a text prompt (model defaults for all other parameters):
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of ocean waves crashing on a beach"
+    }' --output ocean.wav
+```
+
+### Custom Duration
+
+Specify an explicit audio length in seconds:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+```
+
+### High Quality with Negative Prompt
+
+Use a negative prompt to steer generation away from undesired characteristics, and increase inference steps for higher quality:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy",
+        "guidance_scale": 8.0,
+        "num_inference_steps": 150
+    }' --output piano_hq.wav
+```
+
+### Reproducible Generation
+
+Set a `seed` to get deterministic results across runs:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "seed": 42
+    }' --output thunder.wav
+```
+
+### Full Control
+
+Combine all parameters for precise control over generation:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42
+    }' --output thunder_rain.wav
+```
+
+### Quick Generation (Fewer Steps)
+
+For faster generation with slightly lower quality:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Birds chirping in a forest",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output birds_quick.wav
+```
+
+### Python Client
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/generate",
+    json={
+        "input": "Thunder and rain",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42,
+        "response_format": "wav",
+    },
+    timeout=300.0,
+)
+
+with open("thunder.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Parameter Tuning Guide
+
+### `guidance_scale`
+
+Controls how closely the generated audio follows the text prompt.
+
+| Range | Behaviour |
+|-------|-----------|
+| 3 - 5 | More creative / varied output |
+| 7 (default) | Balanced adherence |
+| 10+ | Strict adherence to the prompt |
+
+### `num_inference_steps`
+
+Controls the number of denoising steps in the diffusion process.
+
+| Steps | Quality | Speed | Use Case |
+|-------|---------|-------|----------|
+| 50 | Good | Fast | Quick previews |
+| 100 | Very Good | Medium | General purpose |
+| 150+ | Excellent | Slow | Final / critical audio |
+
+### `audio_length`
+
+Duration of the generated audio clip. For `stable-audio-open-1.0`, the maximum is approximately 47 seconds. If omitted, the model uses its own default length.
+
+### `negative_prompt`
+
+Describes characteristics to avoid. Common negative prompts include:
+
+- `"Low quality, distorted, noisy"`
+- `"Silence, static"`
+- `"Music"` (when generating sound effects only)
+
+## Supported Models
+
+| Model | Description |
+|-------|-------------|
+| `stabilityai/stable-audio-open-1.0` | Open-source audio generation model, up to ~47 seconds, 44.1 kHz stereo |
+
+## Error Responses
+
+### 400 Bad Request
+
+Invalid or missing parameters:
+
+```json
+{
+    "error": {
+        "message": "Audio generation model did not produce audio output.",
+        "type": "BadRequestError",
+        "param": null,
+        "code": 400
+    }
+}
+```
+
+### 404 Not Found
+
+Model mismatch:
+
+```json
+{
+    "error": {
+        "message": "The model `xxx` does not exist.",
+        "type": "NotFoundError",
+        "param": "model",
+        "code": 404
+    }
+}
+```
+
+### 422 Unprocessable Entity
+
+Pydantic validation failure (e.g. invalid `response_format`, `speed` out of range):
+
+```json
+{
+    "detail": [
+        {
+            "type": "literal_error",
+            "msg": "Input should be 'wav', 'pcm', 'flac', 'mp3', 'aac' or 'opus'",
+            ...
+        }
+    ]
+}
+```
+
+## Troubleshooting
+
+### "Audio generation model did not produce audio output"
+
+The model finished but returned no audio data. Verify the server started successfully and the model loaded without errors.
+
+### Server Not Responding
+
+```bash
+# Check if the server is healthy
+curl http://localhost:8091/health
+```
+
+### Audio Quality Issues
+
+- Increase `num_inference_steps` (e.g. 150).
+- Add a negative prompt: `"Low quality, distorted, noisy"`.
+- Increase `guidance_scale` for stronger prompt adherence.
+
+### Generation Timeout
+
+- Reduce `num_inference_steps`.
+- Reduce `audio_length`.
+- Check GPU memory with `nvidia-smi`.
+
+### Out of Memory
+
+- Lower `--gpu-memory-utilization` (e.g. 0.8).
+- Reduce `audio_length`.
+
+## Development
+
+Enable debug logging:
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8091 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni \
+    --uvicorn-log-level debug
+```

docs/serving/speech_api.md

@@ -5,7 +5,6 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
 - **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
 - **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
 - **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
-- **VoxCPM2** (`openbmb/VoxCPM2`) -- Single-stage AR + DiT/CFM TTS. Three synthesis modes: Voice Design (text-only voice description), Controllable Cloning (reference audio + style instructions), and Hi-Fi Cloning (reference audio + transcript for tightest identity match; aka "Ultimate Cloning" in the Gradio demo). Output: 48 kHz.
 
 Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
 
@@ -27,9 +26,6 @@ vllm serve fishaudio/s2-pro --omni --port 8091
 
 # Voxtral TTS
 vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
-
-# VoxCPM2
-vllm serve openbmb/VoxCPM2 --omni --port 8091
 ```
 
 ### Generate Speech
@@ -123,16 +119,6 @@ Content-Type: application/json
 | `ref_text` | string | null | Transcript of reference audio |
 | `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
 
-#### VoxCPM2-specific Parameters
-
-VoxCPM2 reuses the shared `extra_params: dict` field (introduced in #2338 for Voxtral TTS) for model-specific knobs rather than adding new top-level fields:
-
-| Key (inside `extra_params`) | Type | Default | Description |
-|-----------------------------|------|---------|-------------|
-| `cfg_value` | float | 2.0 | Classifier-free guidance scale for the CFM decoder. Range `0.1`–`10.0`; typical useful range is `1.5`–`3.0`. Higher values track the text more strictly at the cost of naturalness. |
-
-Send as e.g. `{"input": "...", "extra_params": {"cfg_value": 2.5}}`. VoxCPM2 also reuses the standard `instructions`, `ref_audio`, and `ref_text` fields, see the mode table below for how they combine.
-
 ### Response Format
 
 Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
@@ -557,64 +543,6 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne
 |-------|-------------|
 | `mistralai/Voxtral-4B-TTS-2603` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |
 
-### VoxCPM2
-
-| Model | Description |
-|-------|-------------|
-| `openbmb/VoxCPM2` | Single-stage AR + DiT/CFM TTS. 48 kHz output, multilingual (primary EN/ZH, usable cross-lingual identity transfer via Ultimate Cloning). |
-
-VoxCPM2 exposes three synthesis modes through the same `/v1/audio/speech` endpoint — the mode is inferred from which fields are present in the request.
-
-| Mode | Required fields | Optional fields | Notes |
-|------|-----------------|-----------------|-------|
-| **Voice Design** | `input` | `instructions`, `extra_params.cfg_value` | Synthesize a voice from a text description alone. `instructions` is prepended to the target text as `(instructions)input` following `voxcpm.cli.build_final_text`. Both English and Chinese instructions are supported. |
-| **Controllable Cloning** | `input`, `ref_audio` | `instructions`, `extra_params.cfg_value` | Clone the reference speaker's timbre; use `instructions` to steer emotion, pace, or other style attributes while preserving identity. |
-| **Hi-Fi Cloning** | `input`, `ref_audio`, `ref_text` | `extra_params.cfg_value` | The model treats `ref_audio` + `ref_text` as a spoken prefix and continues from it. Reproduces every vocal nuance of the reference. Per the canonical VoxCPM2 docs, when Hi-Fi mode is enabled the control instruction is ignored; the server drops `instructions` in this mode automatically if present. |
-
-Reference audio guidelines: 16 kHz mono, at least 5 seconds recommended for a good clone (1 second hard minimum, 30 seconds hard maximum). Clean, non-reverberant speech; the same `ref_audio` can be reused across Controllable and Hi-Fi modes.
-
-Inline non-verbal tags in `input` are passed through as regular text and recognized by the model when present (for example `[laughing]`, `[sigh]`, `[Uhm]`, `[Question-ah]`). See the upstream [VoxCPM2 cookbook](https://voxcpm.readthedocs.io/en/latest/cookbook.html) for the full tag list and usage tips.
-
-**Examples**
-
-Voice Design (no reference, pure text description):
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Hello everyone. Today I want to talk about open source AI.",
-        "instructions": "A warm young woman with a gentle, friendly voice",
-        "extra_params": {"cfg_value": 2.5}
-    }' --output voice_design.wav
-```
-
-Controllable Cloning (reference audio + style instructions):
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Good morning, let me tell you a short story.",
-        "ref_audio": "https://example.com/reference.wav",
-        "instructions": "Warm storyteller, slow cadence, dramatic pauses",
-        "extra_params": {"cfg_value": 2.5}
-    }' --output controllable_cloning.wav
-```
-
-Hi-Fi Cloning (reference audio + transcript; tightest identity):
-
-```bash
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Welcome everyone. I will give a brief update on the project today.",
-        "ref_audio": "https://example.com/reference.wav",
-        "ref_text": "Transcript of what the reference speaker is saying in the reference clip.",
-        "extra_params": {"cfg_value": 2.5}
-    }' --output ultimate_cloning.wav
-```
-
 ## Error Responses
 
 ### 400 Bad Request

docs/user_guide/diffusion/quantization/autoround.md

@@ -1,93 +0,0 @@
-# AutoRound Quantization
-
-## Overview
-
-[AutoRound](https://github.com/intel/auto-round) is an advanced quantization toolkit designed for Large Language Models (LLMs), Vision-Language Models (VLMs), and diffusion models. It achieves high accuracy at ultra-low bit widths (2–4 bits) with minimal tuning by leveraging sign-gradient descent, while providing broad hardware compatibility with multi-datatype support.
-
-The quantization config is auto-detected from the checkpoint's `config.json` (`quantization_config.quant_method = "auto-round"`). No extra CLI flags are needed.
-
-### Supported Schemes
-
-| Scheme | Bits | Status |
-|--------|------|--------|
-| W4A16 | 4 | ✅ Supported |
-| W8A16 | 8 | Planned |
-
-W4A16 is the first supported scheme. Additional schemes will be added in future releases.
-
-## Configuration
-
-1. **Python API**: point `model` at a pre-quantized checkpoint. The quantization is detected automatically.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(model="vllm-project-org/FLUX.1-dev-AutoRound-w4a16")
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=28),
-)
-outputs[0].save_images("output.png")
-```
-
-2. **CLI**: pass the quantized model path directly.
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model vllm-project-org/FLUX.1-dev-AutoRound-w4a16 \
-  --prompt "A cat sitting on a windowsill" \
-  --num-inference-steps 28 \
-  --output outputs/flux_w4a16.png
-```
-
-No `--quantization` flag is needed — the quantization method is read from the checkpoint.
-
-## How It Works
-
-The checkpoint's `config.json` contains:
-
-```json
-{
-  "quantization_config": {
-    "quant_method": "auto-round",
-    "bits": 4,
-    "group_size": 128,
-    "sym": true,
-    "packing_format": "auto_round:auto_gptq",
-    "block_name_to_quantize": "transformer_blocks,single_transformer_blocks"
-  }
-}
-```
-
-At load time:
-
-1. `TransformerConfig.from_dict()` parses the `quantization_config` section and builds a vLLM `INCConfig` via `build_quant_config("auto-round", ...)`.
-2. `OmniDiffusionConfig.set_tf_model_config()` propagates the detected config to the engine.
-3. The appropriate compute kernel (e.g. GPTQ-Marlin for W4A16) is selected automatically based on the checkpoint's bit-width and packing format.
-
-## Supported Models
-
-| Model | HF Checkpoint | Scheme | Group Size | Backend |
-|-------|--------------|--------|------------|---------|
-| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | W4A16 | 128 | GPTQ-Marlin |
-| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | W4A16 | 128 | GPTQ-Marlin |
-| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | W4A16 | 128 | GPTQ-Marlin |
-
-## Creating a Quantized Checkpoint
-
-Use [AutoRound](https://github.com/intel/auto-round) to quantize a BF16 model. The `--scheme` flag selects the quantization scheme:
-
-```bash
-# W4A16 (4-bit weight, 16-bit activation)
-auto-round \
-    --model black-forest-labs/FLUX.1-dev \
-    --scheme W4A16 \
-    --batch_size 1 \
-    --disable_opt_rtn \
-    --dataset coco2014 \
-    --iters 0
-```
-
-The output directory can be used directly as the `model` argument. See the [AutoRound documentation](https://github.com/intel/auto-round) for all available schemes and options.

docs/user_guide/diffusion/quantization/fp8.md

@@ -1,82 +0,0 @@
-# FP8 Quantization
-
-## Overview
-
-FP8 quantization converts BF16/FP16 weights to FP8 at model load time. No calibration or pre-quantized checkpoint needed.
-
-Depending on the model, either all layers can be quantized, or some sensitive layers should stay in BF16. See the [per-model table](#supported-models) for which case applies.
-
-Common sensitive layers in DiT-based diffusion models include **image-stream MLPs** (`img_mlp`). These are particularly vulnerable to FP8 precision loss because they process denoising latents whose dynamic range shifts significantly across timesteps, and unlike attention projections (which benefit from QK-Norm stabilization), MLPs have no built-in normalization to absorb quantization error. In deep architectures (e.g., 60+ residual blocks), small per-layer errors compound and degrade output quality. Other layers such as **attention projections** (`to_qkv`, `to_out`) and **text-stream MLPs** (`txt_mlp`) are generally more robust due to normalization or more stable input statistics.
-
-## Configuration
-
-1. **Python API**: set `quantization="fp8"`. To skip sensitive layers, use `quantization_config` with `ignored_layers`.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# All layers quantized
-omni = Omni(model="<your-model>", quantization="fp8")
-
-# Skip sensitive layers
-omni = Omni(
-    model="<your-model>",
-    quantization_config={
-        "method": "fp8",
-        "ignored_layers": ["<layer-name>"],
-    },
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-2. **CLI**: pass `--quantization fp8` and optionally `--ignored-layers`.
-
-```bash
-# All layers
-python text_to_image.py --model <your-model> --quantization fp8
-
-# Skip sensitive layers
-python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers "img_mlp"
-
-# Online serving
-vllm serve <your-model> --omni --quantization fp8
-```
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `method` | str | — | Quantization method (`"fp8"`) |
-| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16 |
-| `activation_scheme` | str | `"dynamic"` | `"dynamic"` (no calibration) or `"static"` |
-| `weight_block_size` | list[int] \| None | `None` | Block size for block-wise weight quantization |
-
-The available `ignored_layers` names depend on the model architecture (e.g., `to_qkv`, `to_out`, `img_mlp`, `txt_mlp`). Consult the transformer source for your target model.
-
-## Supported Models
-
-| Model | HF Models | Recommendation | `ignored_layers` |
-|-------|-----------|---------------|------------------|
-| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
-| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
-| Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
-| HunyuanImage-3 | `tencent/HunyuanImage3` | All layers | None |
-| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
-| GLM-Image | `zai-org/GLM-Image` | All layers | None |
-| Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | All layers | None |
-
-## Combining with Other Features
-
-FP8 quantization can be combined with cache acceleration:
-
-```python
-omni = Omni(
-    model="<your-model>",
-    quantization="fp8",
-    cache_backend="tea_cache",
-    cache_config={"rel_l1_thresh": 0.2},
-)
-```

docs/user_guide/diffusion/quantization/gguf.md

@@ -1,227 +0,0 @@
-# GGUF Quantization
-
-## Goals
-1. Reuse vLLM quantization configs and weight loaders as much as possible.
-2. Add native GGUF support to diffusion transformers without changing model definitions.
-3. Keep user-facing knobs minimal and consistent across offline and online flows.
-
-## Scope
-1. Models: Z-Image, Flux2-klein, and Qwen-Image family.
-2. Components: diffusion transformer weights, loader paths, and quantization configs.
-3. Modes: native GGUF (transformer-only weights).
-
-## Architecture Overview
-1. `OmniDiffusionConfig` accepts `quantization` or `quantization_config`.
-2. Diffusion quantization wrapper (`DiffusionGGUFConfig`) produces vLLM `QuantizationConfig` objects for linear layers.
-3. `DiffusersPipelineLoader` branches on quantization method and loads either HF weights or GGUF weights for the transformer.
-4. GGUF transformer loading is routed through model-specific adapters (e.g., Flux2Klein).
-5. vLLM GGUF path uses `GGUFConfig` and `GGUFLinearMethod` for matmul.
-
-## Call Chain (Offline)
-```
-CLI (examples/offline_inference/text_to_image/text_to_image.py)
-  |
-  v
-Omni (vllm_omni/entrypoints/__init__.py)
-  |
-  v
-AsyncOmniEngine
-  |
-  v
-Orchestrator
-  |
-  v
-StageDiffusionClient
-  |
-  v
-DiffusionWorker
-  |
-  v
-DiffusionModelRunner
-  |
-  v
-DiffusersPipelineLoader
-  |
-  v
-Pipeline.forward (Flux2/Qwen/Z-Image)
-  |
-  v
-DiffusionEngine
-  |
-  v
-OmniRequestOutput
-  |
-  v
-Client (saved PNG)
-```
-
-## Call Chain (Online)
-```
-Client
-  |
-  | POST /v1/images/generations
-  v
-APIServer (vllm_omni/entrypoints/openai/api_server.py)
-  |
-  v
-_generate_with_async_omni
-  |
-  v
-AsyncOmni
-  |
-  v
-AsyncOmniEngine
-  |
-  v
-Orchestrator
-  |
-  v
-StageDiffusionClient
-  |
-  v
-OmniRequestOutput
-  |
-  v
-encode_image_base64
-  |
-  v
-ImageGenerationResponse
-  |
-  v
-Client
-```
-
-## Call Chain (GGUF Operator Path)
-```
-Pipeline.forward (Flux2/Qwen/Z-Image)
-  |
-  v
-Transformer blocks
-  |
-  v
-QKVParallelLinear / ColumnParallelLinear / RowParallelLinear
-  |
-  v
-LinearBase.forward
-  |
-  v
-QuantMethod.apply (GGUFLinearMethod.apply)
-  |
-  v
-fused_mul_mat_gguf
-  |
-  v
-_fused_mul_mat_gguf (custom op)
-  |
-  v
-ops.ggml_dequantize
-  |
-  v
-x @ weight.T
-```
-
-## GGUF Weight Loading Path (Transformer-Only)
-1. `DiffusersPipelineLoader.load_model` detects `quantization_config.method == "gguf"`.
-2. `gguf_model` is resolved as one of: local file, `repo/file.gguf`, or `repo:quant_type`.
-3. GGUF weights are routed through adapters in `vllm_omni/diffusion/model_loader/gguf_adapters/`.
-4. Name mapping is applied per-architecture (Z-Image, Flux2Klein, Qwen-Image).
-5. GGUF weights are loaded into transformer modules, remaining non-transformer weights come from the HF checkpoint.
-
-## GGUF Adapter Design
-1. `GGUFAdapter` is an abstract base class for model-specific adapters.
-2. `Flux2KleinGGUFAdapter` implements Flux2-Klein remapping + qkv split + adaLN swap.
-3. `ZImageGGUFAdapter` implements Z-Image qkv + ffn shard handling and linear qweight routing.
-4. `QwenImageGGUFAdapter` implements Qwen-Image fused attention projection remapping.
-5. `get_gguf_adapter(...)` strictly selects by model class/config; unsupported models raise an error (no fallback adapter).
-
-Adapter paths:
-- Base: `vllm_omni/diffusion/model_loader/gguf_adapters/base.py`
-- Z-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py`
-- Flux2-Klein: `vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py`
-- Qwen-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py`
-
-## User Usage (Offline)
-
-### Baseline BF16
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
-  --prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
-  --height 768 \
-  --width 1360 \
-  --seed 42 \
-  --cfg-scale 4.0 \
-  --num-images-per-prompt 1 \
-  --num-inference-steps 4 \
-  --output outputs/flux2_klein_4b.png
-```
-
-### Native GGUF (Transformer Only)
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
-  --gguf-model "/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf" \
-  --quantization gguf \
-  --prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
-  --height 768 \
-  --width 1360 \
-  --seed 42 \
-  --cfg-scale 4.0 \
-  --num-images-per-prompt 1 \
-  --num-inference-steps 4 \
-  --output outputs/flux2_klein_4b_gguf.png
-```
-
-Notes for GGUF:
-1. Many GGUF repos do not ship `model_index.json` and configs. Use the base repo for `--model` and only pass the GGUF file via `--gguf-model`.
-2. `gguf_model` supports local path, `repo/file.gguf`, or `repo:quant_type`.
-
-## User Usage (Online)
-
-### Start Server (Native GGUF via CLI)
-```bash
-vllm serve /workspace/models/black-forest-labs/FLUX.2-klein-4B \
-  --omni \
-  --port 8000 \
-  --quantization-config '{"method":"gguf","gguf_model":"/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf"}'
-```
-
-### Online Request (Images API)
-```bash
-curl -X POST http://localhost:8000/v1/images/generations \
-  -H "Content-Type: application/json" \
-  -d '{
-    "prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
-    "size": "1024x1024",
-    "seed": 42,
-    "num_inference_steps": 4
-  }'
-```
-
-## Qwen-Image Example
-
-### Offline
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model Qwen/Qwen-Image \
-  --gguf-model QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf \
-  --quantization gguf \
-  --prompt "a red paper kite hanging from a pine tree in a winter courtyard" \
-  --height 1024 \
-  --width 1024 \
-  --seed 42 \
-  --num_inference_steps 20 \
-  --output outputs/qwen_image_q4km.png
-```
-
-### Online
-```bash
-vllm serve Qwen/Qwen-Image \
-  --omni \
-  --port 8000 \
-  --quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
-```
-
-Notes:
-1. The GGUF repo provides transformer weights only; the base model repo still provides tokenizer, text encoder, scheduler, and VAE.
-2. The same adapter path covers `QwenImagePipeline`, `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, and `QwenImageLayeredPipeline`.

docs/user_guide/diffusion/quantization/int8.md

@@ -1,75 +0,0 @@
-# Int8 Quantization
-
-## Overview
-
-Int8 quantization converts BF16/FP16 weights to Int8 at model load time. No calibration or pre-quantized checkpoint needed.
-
-Depending on the model, either all layers can be quantized, or some sensitive layers should stay in BF16/FP16. See the [per-model table](#supported-models) for which case applies.
-
-## Configuration
-
-1. **Python API**: set `quantization="int8"`. To skip sensitive layers, use `quantization_config` with `ignored_layers`.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# All layers quantized
-omni = Omni(model="<your-model>", quantization="int8")
-
-# Skip sensitive layers
-omni = Omni(
-    model="<your-model>",
-    quantization_config={
-        "method": "int8",
-        "ignored_layers": ["<layer-name>"],
-    },
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-2. **CLI**: pass `--quantization int8` and optionally `--ignored-layers`.
-
-```bash
-# All layers
-python text_to_image.py --model <your-model> --quantization int8
-
-# Skip sensitive layers
-python text_to_image.py --model <your-model> --quantization int8 --ignored-layers "img_mlp"
-
-# Online serving
-vllm serve <your-model> --omni --quantization int8
-```
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `method` | str | — | Quantization method (`"int8"`) |
-| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 |
-| `activation_scheme` | str | `"dynamic"` | `"dynamic"` (no calibration) |
-
-
-The available `ignored_layers` names depend on the model architecture (e.g., `to_qkv`, `to_out`, `img_mlp`, `txt_mlp`). Consult the transformer source for your target model.
-
-## Supported Models
-
-| Model | HF Models | Recommendation | `ignored_layers` |
-|-------|-----------|---------------|------------------|
-| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
-| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | All layers | None |
-
-## Combining with Other Features
-
-Int8 quantization can be combined with cache acceleration:
-
-```python
-omni = Omni(
-    model="<your-model>",
-    quantization="int8",
-    cache_backend="tea_cache",
-    cache_config={"rel_l1_thresh": 0.2},
-)
-```

docs/user_guide/diffusion/quantization/msmodelslim.md

@@ -1,56 +0,0 @@
-# msModelSlim Quantization
-
-## Overview
-
-[msModelSlim](https://github.com/Ascend/msmodelslim) is an Ascend-friendly compression tool focused on acceleration, using compression techniques, and built for Ascend hardware. It includes a series of inference optimization technologies such as quantization and compression, aiming to accelerate large language dense models, MoE models, multimodal understanding models, multimodal generation models, etc.
-
-Once you have a quantized model which is generated by **msModelSlim**, you can use vLLM Omni for inference by specifying the --quantization ascend parameter to enable quantization features.
-
-### Supported Schemes
-
-| Scheme | Bits | Status |
-|--------|------|--------|
-| W8A8 | 8 | ✅ Supported |
-| W4A4 | 4 | Planned |
-
-W8A8 is the first supported scheme. Additional schemes will be added in future releases.
-
-## Model Quantization
-
-The following example shows how to generate W8A8 quantized weights for the [Wan2_2 model](https://gitcode.com/Ascend/msmodelslim/blob/master/example/multimodal_sd/Wan2_2/README.md).
-
-**Quantization Script:**
-
-```bash
-msmodelslim quant \
-  --model_path /path/to/wan2_2_t2v_float_weights \
-  --save_path /path/to/wan2_2_t2v_quantized_weights \
-  --device npu \
-  --model_type Wan2_2 \
-  --config_path /lab_practice/wan2_2/wan2_2_w8a8f8_mxfp_t2v.yaml \
-  --trust_remote_code True
-```
-
-After quantization completes, the output directory will contain the quantized model files.
-
-For more examples, refer to the [official examples](https://gitcode.com/Ascend/msit/tree/master/msmodelslim/example).
-
-## Configuration
-
-1. **CLI**: pass `--quantization ascend`.
-
-```bash
-# Offline inference
-python text_to_image.py --model <your-model> --quantization ascend
-
-# Online serving
-vllm serve <your-model> --omni --quantization ascend
-```
-
-## Supported Models
-
-| Model | HF Models | Recommendation | `ignored_layers` |
-|-------|-----------|---------------|------------------|
-| HunyuanImage-3.0 | - | All layers | None |
-
-Currently, quantized HunyuanImage-3.0 weights have not been uploaded to public model platforms such as Hugging Face. You can use a [HunyuanImage-3.0-adapted msModelSlim version](https://gitcode.com/betta18/msmodelslim/tree/hyimage3_mxfp8) to generate the quantized weights manually. We will upload the quantized weights as soon as possible.

docs/user_guide/diffusion/quantization/overview.md

@@ -1,153 +0,0 @@
-# Quantization
-
-vLLM-Omni provides a unified quantization framework that supports both diffusion models (DiT, text encoders, VAEs) and multi-stage omni models (e.g., Qwen3-Omni thinker). Quantization reduces memory usage and accelerates inference.
-
-## Supported Methods
-
-### Diffusion Models
-
-| Method | Guide | Description | Tested Models | Min GPU |
-|--------|-------|-------------|---------------|---------|
-| FP8 | [FP8](fp8.md) | FP8 W8A8, dynamic or static | Z-Image, Qwen-Image, Flux, Bagel | SM 89 (Ada) |
-| Int8 | [Int8](int8.md) | Int8 W8A8 | Z-Image, Qwen-Image | SM 89 (Ada) / Ascend NPU |
-| GGUF | [GGUF](gguf.md) | GGUF format, dequant+GEMM for N-D tensors | Z-Image, Flux | SM 60 |
-| AutoRound | [AutoRound](autoround.md) | W4A16 (pre-quantized) | Flux | SM 80 (Ampere) |
-
-### Multi-stage Omni Models (Pre-quantized Checkpoints)
-
-| Method | Description | Tested Models | Min GPU |
-|--------|-------------|---------------|---------|
-| ModelOpt FP8 | Pre-quantized FP8 via NVIDIA ModelOpt | Qwen3-Omni (thinker) | SM 89 (Ada/Hopper) |
-| ModelOpt NVFP4 | Pre-quantized NVFP4 via NVIDIA ModelOpt | Qwen3-Omni (experimental, quality issues) | SM 100 (Blackwell) |
-
-!!! note
-    AWQ, GPTQ, and BitsAndBytes are supported by vLLM's upstream quantization registry but have **not been tested** with vLLM-Omni pipelines. They may work via `build_quant_config()` but are not validated.
-
-### Pre-quantized LLM Checkpoints (Multi-stage Models)
-
-For multi-stage models like Qwen3-Omni, the unified quantization framework auto-detects
-pre-quantized checkpoints via `quantization_config` in the HF config. Quantization is
-automatically scoped to the thinker's `language_model` — audio encoder, vision encoder,
-talker, and code2wav remain in BF16.
-
-| Format | `quant_algo` | Hardware | Example | Status |
-|--------|-------------|----------|---------|--------|
-| ModelOpt FP8 | `FP8` | Ada/Hopper (SM 89+) | `asdazd/Qwen3-Omni-30B-A3B-Instruct_modelopt_FP8` | Tested — 47% memory reduction, comparable throughput |
-| ModelOpt NVFP4 | `NVFP4` | Blackwell (SM 100+) | `shunyang90/Qwen3-Omni-30B-A3B-Instruct-NVFP4` | Experimental — loads and runs but output quality is unacceptable |
-
-**Tested FP8 results (Qwen3-Omni, 1×H200):**
-
-| Config | Model Memory (GiB) | Mem Reduction | Decode (tok/s) |
-|--------|-------------------|---------------|----------------|
-| BF16 baseline | 59.26 | — | 41.6 |
-| FP8 (ModelOpt) | 31.41 | 47% | 39.9 |
-
-FP8 enables the full 3-stage pipeline (thinker + talker + code2wav) on a single 64GB GPU,
-which is impossible with BF16 (thinker alone requires 59.26 GiB).
-
-## Quantization Scope
-
-### Diffusion Models
-
-When `--quantization fp8` is enabled for diffusion models:
-
-| Component | What Gets Quantized | Mechanism |
-|-----------|-------------------|-----------|
-| **DiT (transformer)** | `nn.Linear` layers | vLLM W8A8 FP8 compute (Ada/Hopper) or weight-only (older GPUs) |
-| **Text encoder** | `nn.Linear` layers | FP8 weight storage, BF16 compute |
-| **VAE** | `nn.Conv2d`, `nn.Conv3d` layers | FP8 weight storage, BF16 compute |
-
-### Multi-stage Omni Models
-
-For pre-quantized omni model checkpoints, quantization is scoped per component:
-
-| Component | Quantized? | Notes |
-|-----------|-----------|-------|
-| **Thinker (language_model)** | Yes | FP8 or NVFP4 via ModelOpt |
-| **Audio encoder** | No | Stays BF16 |
-| **Vision encoder** | No | Stays BF16 |
-| **Talker** | No | Stays BF16 |
-| **Code2Wav** | No | Stays BF16 |
-
-!!! note
-    Not all models support all components. See the [FP8 supported models table](fp8.md#supported-models) for per-model details.
-
-## Device Compatibility for FP8
-
-| Device | Example Hardware | FP8 | NVFP4 |
-|--------|-----------------|-----|-------|
-| Blackwell GPU (SM 100+) | RTX 5090, B100, B200 | Yes | Yes (native FP4 HW) |
-| Ada/Hopper GPU (SM 89+) | RTX 4090, H100, H200 | Yes (W8A8 native) | No |
-| Turing/Ampere GPU (SM 75-86) | RTX 3090, A100 | Yes (weight-only Marlin) | No |
-| Ascend NPU | Atlas 800T A2 (910B) | Not yet supported | No |
-
-Kernel selection is automatic on CUDA GPUs.
-
-## Device Compatibility for Int8
-
-| Device Type | Generation | Example | Int8 Mode |
-|-------------|---------------|-------------------|----------|
-| NVIDIA GPU | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
-| Ascend NPU | Atlas A2/Atlas A3 | Atlas 800T A2/Atlas 900 A3 | Full W8A8 with native hardware |
-
-## Python API
-
-The `build_quant_config()` factory accepts multiple input formats:
-
-```python
-from vllm_omni.quantization import build_quant_config
-
-# String
-config = build_quant_config("fp8")
-
-# Dict with parameters
-config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
-
-# AutoRound / INC (auto-detected from checkpoint, or explicit)
-config = build_quant_config("auto-round", bits=4, group_size=128)
-
-# Per-component dict
-config = build_quant_config({
-    "transformer": {"method": "fp8"},
-    "vae": None,
-})
-
-# None / "none"
-config = build_quant_config(None)  # returns None
-```
-
-## Architecture
-
-```
-build_quant_config(spec)
-    ├── str  → _build_single(method) → vLLM registry / _OVERRIDES
-    ├── dict with "method" → _build_single(method, **kwargs)
-    ├── per-component dict → ComponentQuantizationConfig
-    │       routes get_quant_method() by longest-prefix match
-    ├── QuantizationConfig → passthrough
-    └── None / "none" → None
-```
-
-## Migration Guide
-
-### Before (v0.14.x)
-
-```python
-# Old diffusion-specific API (removed)
-from vllm_omni.diffusion.quantization import get_diffusion_quant_config
-config = get_diffusion_quant_config("fp8", activation_scheme="static")
-```
-
-### After (v0.16.0+)
-
-```python
-# Unified API — delegates to vLLM's registry
-from vllm_omni.quantization import build_quant_config
-config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
-
-# Per-component (new)
-config = build_quant_config({
-    "transformer": {"method": "fp8"},
-    "vae": None,
-})
-```

docs/user_guide/diffusion_features.md

@@ -59,7 +59,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on
 | Method | Description | Best For |
 |--------|-------------|----------|
 | **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
-| **[Quantization](diffusion/quantization/overview.md)** | Reduces DiT layers from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss    |
+| **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss    |
 | **[VAE Patch Parallelism](diffusion/parallelism/vae_patch_parallel.md)** | Distributes VAE decode tiling across GPUs | High-resolution generation with reduced VAE memory peak |
 
 ### Extensions
@@ -86,9 +86,9 @@ Execution modes control how the diffusion pipeline processes denoise steps.
 
 | Method | Configuration | Description | Best For |
 |--------|--------------|-------------|----------|
-| **[FP8](diffusion/quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
-| **[INT8](diffusion/quantization/int8.md)** | `quantization="int8"` | INT8 weight-only, no calibration or pre-quantized checkpoint needed | Memory reduction, broad GPU compatibility |
-| **[GGUF](diffusion/quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs |
+| **[FP8](quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on validated transformer stages | Memory reduction, inference speedup |
+| **[INT8](quantization/int8.md)** | `quantization="int8"` | INT8 W8A8 on validated transformer stages | Memory reduction, broad GPU compatibility |
+| **[GGUF](quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs |
 
 ## Supported Models
 
@@ -101,7 +101,7 @@ The following tables show which models support each feature:
 > Notes:
 
 > 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only.
-> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](diffusion/quantization/overview.md) for per-method (FP8, GGUF, …) and per-model support details.
+> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details.
 
 ### ImageGen
 
@@ -261,7 +261,7 @@ Measured on NVIDIA H800:
 
 - **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage
 - **[VAE Patch Parallelism Guide](diffusion/parallelism/vae_patch_parallel.md)** - Distribute VAE decode tiling across GPUs for high-resolution images
-- **[Quantization Overview](diffusion/quantization/overview.md)** - Overview of quantization methods for diffusion models
+- **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
 
 **Extensions:**
 

docs/user_guide/examples/online_serving/text_to_audio.md

@@ -0,0 +1,193 @@
+# Text-To-Audio
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/stable_audio>.
+
+This example demonstrates how to deploy Stable Audio models for online text-to-audio generation using vLLM-Omni.
+
+## Supported Models
+
+| Model | Description |
+|-------|-------------|
+| `stabilityai/stable-audio-open-1.0` | Open-source audio generation, up to ~47 seconds, 44.1 kHz stereo |
+
+## Start Server
+
+### Basic Start
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8091 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+## API Calls
+
+### Method 1: Using curl
+
+```bash
+# Run all curl examples
+bash curl_examples.sh
+
+# Or execute directly
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+### Method 2: Using Python Client
+
+```bash
+cd examples/online_serving/stable_audio
+
+# Simple generation
+python stable_audio_client.py \
+    --text "The sound of a cat purring"
+
+# With custom duration
+python stable_audio_client.py \
+    --text "A dog barking" \
+    --audio_length 5.0
+
+# With all parameters
+python stable_audio_client.py \
+    --text "Thunder and rain" \
+    --audio_length 15.0 \
+    --negative_prompt "Low quality" \
+    --guidance_scale 7.0 \
+    --num_inference_steps 100 \
+    --seed 42 \
+    --output thunder.wav
+```
+
+The Python client supports the following command-line arguments:
+
+- `--api_url`: API endpoint URL (default: `http://localhost:8091/v1/audio/generate`)
+- `--text`: Text prompt for audio generation (default: `"The sound of a cat purring"`)
+- `--audio_length`: Audio length in seconds (default: `10.0`, max ~47s for `stable-audio-open-1.0`)
+- `--audio_start`: Audio start time in seconds (default: `0.0`)
+- `--negative_prompt`: Negative prompt for classifier-free guidance (default: `"Low quality"`)
+- `--guidance_scale`: Guidance scale for diffusion (default: `7.0`)
+- `--num_inference_steps`: Number of inference steps (default: `100`)
+- `--seed`: Random seed for reproducibility (default: `None`)
+- `--response_format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`
+- `--output`: Output file path (default: `stable_audio_output.wav`)
+
+### Method 3: Using Python httpx
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/generate",
+    json={
+        "input": "The sound of ocean waves crashing on a beach",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+    },
+    timeout=300.0,
+)
+
+with open("ocean.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Request Format
+
+### Simple Generation
+
+```json
+{
+    "input": "The sound of ocean waves"
+}
+```
+
+### Generation with Parameters
+
+```json
+{
+    "input": "A piano playing a gentle melody",
+    "audio_length": 10.0,
+    "negative_prompt": "Low quality, distorted, noisy",
+    "guidance_scale": 8.0,
+    "num_inference_steps": 150,
+    "seed": 42,
+    "response_format": "wav"
+}
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+Content-Type: application/json
+```
+
+### Generation Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
+| `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducible generation |
+
+### Response Format
+
+Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
+
+## Tuning Tips
+
+1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`.
+2. **Quality vs Speed**:
+   - 50 steps: Fast, decent quality (quick previews)
+   - 100 steps: Good balance (general purpose)
+   - 150+ steps: High quality, slower (final / critical audio)
+3. **Guidance Scale**:
+   - Lower (3 - 5): More creative / varied output
+   - Default (7): Good balance
+   - Higher (10+): Strict adherence to the prompt
+4. **Negative Prompts**: Use to avoid unwanted characteristics such as `"Low quality"`, `"distorted"`, `"noisy"`.
+5. **Seeds**: Set a fixed seed to get deterministic, reproducible results.
+
+## File Description
+
+| File | Description |
+|------|-------------|
+| `curl_examples.sh` | Curl examples covering common use cases |
+| `stable_audio_client.py` | Python client with full CLI argument support |
+
+## Troubleshooting
+
+1. **Audio generation model did not produce audio output**: Verify the server started successfully and the model loaded without errors.
+2. **Connection refused**: Make sure the server is running on the correct port.
+3. **Generation timeout**: Reduce `num_inference_steps` or `audio_length`, and check GPU memory with `nvidia-smi`.
+4. **Out of memory**: Lower `--gpu-memory-utilization` or reduce `audio_length`.
+5. **Audio quality issues**: Increase `num_inference_steps`, add a negative prompt, or raise `guidance_scale`.
+
+## Example materials
+
+??? abstract "stable_audio_client.py"
+    ``````py
+    --8<-- "examples/online_serving/stable_audio/stable_audio_client.py"
+    ``````
+??? abstract "curl_examples.sh"
+    ``````sh
+    --8<-- "examples/online_serving/stable_audio/curl_examples.sh"
+    ``````

docs/user_guide/quantization/autoround.md

@@ -0,0 +1,127 @@
+# AutoRound Quantization
+
+## Overview
+
+[AutoRound](https://github.com/intel/auto-round) produces pre-quantized
+checkpoints for LLMs, VLMs, and diffusion models. vLLM-Omni reads the
+checkpoint's `config.json` and auto-detects
+`quantization_config.quant_method = "auto-round"`.
+
+AutoRound is static quantization: no `--quantization` flag is needed at
+inference time when the checkpoint already contains the quantization config.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ |
+| AMD ROCm | ⭕ |
+| Intel XPU | ✅ |
+| Ascend NPU | ❌ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide. AutoRound is Intel-supported.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | Checkpoint | Scope | Scheme | Backend |
+|-------|------------|-------|--------|---------|
+| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
+| Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated |
+| Wan2.2 | Not listed | Diffusion transformer | W4A16 | Not validated |
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Checkpoint | Scope | Scheme | Backend |
+|-------|------------|-------|--------|---------|
+| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | Language-model stage | W4A16 | AutoRound |
+| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | Thinker language-model stage | W4A16 | AutoRound |
+| Qwen3-TTS | Not listed | TTS language-model stage | W4A16 | Not validated |
+
+AutoRound support is checkpoint-driven. A model is supported when its
+checkpoint uses a compatible INC/AutoRound config and the target stage maps to
+vLLM-Omni's runtime module names.
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
+| GLM-Image | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
+
+## Configuration
+
+Python API:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="vllm-project-org/FLUX.1-dev-AutoRound-w4a16")
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=28),
+)
+outputs[0].save_images("output.png")
+```
+
+CLI:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model vllm-project-org/FLUX.1-dev-AutoRound-w4a16 \
+  --prompt "A cat sitting on a windowsill" \
+  --num-inference-steps 28 \
+  --output outputs/flux_w4a16.png
+```
+
+## Parameters
+
+| Field | Type | Description |
+|-------|------|-------------|
+| `quant_method` | str | Must be `"auto-round"` |
+| `bits` | int | Quantized weight bit width, usually `4` |
+| `group_size` | int | Quantization group size |
+| `packing_format` | str | AutoRound packing format, for example `auto_round:auto_gptq` |
+| `block_name_to_quantize` | str | Checkpoint block names that should map to runtime module names |
+
+The checkpoint should contain a config like:
+
+```json
+{
+  "quantization_config": {
+    "quant_method": "auto-round",
+    "bits": 4,
+    "group_size": 128,
+    "sym": true,
+    "packing_format": "auto_round:auto_gptq",
+    "block_name_to_quantize": "transformer_blocks,single_transformer_blocks"
+  }
+}
+```
+
+## Validation and Notes
+
+At load time, vLLM-Omni builds an `OmniINCConfig`, remaps checkpoint block names
+to runtime module names, and selects the matching vLLM compute backend.
+
+Example checkpoint creation:
+
+```bash
+auto-round \
+  --model black-forest-labs/FLUX.1-dev \
+  --scheme W4A16 \
+  --batch_size 1 \
+  --disable_opt_rtn \
+  --dataset coco2014 \
+  --iters 0
+```
+
+Use the generated output directory directly as the `model` argument. See the
+[AutoRound documentation](https://github.com/intel/auto-round) for all
+available schemes and options.

docs/user_guide/quantization/fp8.md

@@ -0,0 +1,121 @@
+# FP8 Quantization
+
+## Overview
+
+FP8 quantization converts BF16/FP16 weights to FP8 at model load time, or loads
+a checkpoint whose target stage already declares an FP8 quantization config.
+Online activation scaling is the default and does not require calibration.
+Static activation scaling is supported when calibrated scale information is
+available.
+
+Some architectures can quantize all linear layers. Others have
+quality-sensitive layers that should stay in BF16 through `ignored_layers`.
+Image-stream MLPs (`img_mlp`) are a common sensitive target because denoising
+latent ranges shift across timesteps and small per-layer errors can compound
+in deep DiT blocks.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ |
+| AMD ROCm | ⭕ |
+| Intel XPU | ⭕ |
+| Ascend NPU | ❌ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide. FP8 on Ampere may use a weight-only path where available.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` |
+|-------|-----------|:-------:|:------:|----------------|------------------|
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` |
+| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD |
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None |
+| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None |
+| FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None |
+| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None |
+| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None |
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Scope | Format | Status |
+|-------|-------|--------|--------|
+| Qwen3-Omni | Thinker language-model stage | ModelOpt `quant_algo=FP8` | Tested for thinker memory reduction |
+| Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated |
+
+Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a
+model-specific guide says otherwise.
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Stage-specific transformer or DiT module | Not validated | Route FP8 to the intended stage before enabling |
+| GLM-Image | Stage-specific transformer or DiT module | Not validated | Validate quality against BF16 baseline |
+
+## Configuration
+
+Python API:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="<your-model>", quantization="fp8")
+
+omni_with_skips = Omni(
+    model="<your-model>",
+    quantization_config={
+        "method": "fp8",
+        "ignored_layers": ["img_mlp"],
+    },
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+CLI:
+
+```bash
+python text_to_image.py --model <your-model> --quantization fp8
+python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers "img_mlp"
+vllm serve <your-model> --omni --quantization fp8
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | - | Quantization method (`"fp8"`) |
+| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16 |
+| `activation_scheme` | str | `"dynamic"` | `"dynamic"` selects online activation scaling, or `"static"` when scales are available |
+| `weight_block_size` | list[int] \| None | `None` | Block size for block-wise weight quantization |
+
+The available `ignored_layers` names depend on the model architecture, for
+example `to_qkv`, `to_out`, `img_mlp`, or `txt_mlp`.
+
+## Validation and Notes
+
+FP8 quantization can be combined with cache acceleration:
+
+```python
+omni = Omni(
+    model="<your-model>",
+    quantization="fp8",
+    cache_backend="tea_cache",
+    cache_config={"rel_l1_thresh": 0.2},
+)
+```
+
+Compare generated outputs with a BF16 baseline before adding a new model to the
+supported table. GLM-Image and Helios are not listed as FP8-supported diffusion
+models until they have method-specific validation.

docs/user_guide/quantization/gguf.md

@@ -0,0 +1,107 @@
+# GGUF Quantization
+
+## Overview
+
+GGUF support loads pre-quantized diffusion transformer weights while keeping
+the rest of the pipeline on the base Hugging Face checkpoint. Use the base
+model for tokenizer, text encoder, scheduler, and VAE, then pass the GGUF file
+for the transformer.
+
+GGUF is static quantization: the quantized weights are produced before serving.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ |
+| AMD ROCm | ⭕ |
+| Intel XPU | ⭕ |
+| Ascend NPU | ❌ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | HF base model | GGUF input | Scope | Adapter |
+|-------|---------------|------------|-------|---------|
+| Qwen-Image family | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512`, edit and layered Qwen-Image pipelines | Local `.gguf`, `repo/file.gguf`, or `repo:quant_type` | Transformer only | `QwenImageGGUFAdapter` |
+| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Transformer only | No validated adapter listed |
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Local `.gguf`, `repo/file.gguf`, or `repo:quant_type` | Transformer only | `ZImageGGUFAdapter` |
+| FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Local `.gguf`, `repo/file.gguf`, or `repo:quant_type` | Transformer only | `Flux2KleinGGUFAdapter` |
+
+Generic FLUX.1 GGUF checkpoints are not listed here; the implemented adapter is
+for the FLUX.2-klein path.
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| Qwen3-Omni | Thinker language-model stage | Not validated | GGUF is not documented for omni/TTS AR stages |
+| Qwen3-TTS | TTS language-model stage | Not validated | GGUF is not documented for TTS stages |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Stage-specific transformer weights | Not validated | Requires a model-specific GGUF adapter |
+| GLM-Image | Stage-specific transformer weights | Not validated | Requires a model-specific GGUF adapter |
+
+## Configuration
+
+Offline:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model Qwen/Qwen-Image \
+  --gguf-model QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf \
+  --quantization gguf \
+  --prompt "a red paper kite hanging from a pine tree in a winter courtyard" \
+  --height 1024 \
+  --width 1024 \
+  --seed 42 \
+  --num_inference_steps 20 \
+  --output outputs/qwen_image_q4km.png
+```
+
+Online:
+
+```bash
+vllm serve Qwen/Qwen-Image \
+  --omni \
+  --port 8000 \
+  --quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | - | Quantization method (`"gguf"`) |
+| `gguf_model` | str | - | Local GGUF file, explicit Hugging Face file, or `repo:quant_type` selector |
+
+`gguf_model` accepts:
+
+| Form | Example |
+|------|---------|
+| Local file | `/models/z-image-Q4_K_M.gguf` |
+| Explicit HF file | `QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf` |
+| HF repo plus quant type | `owner/repo:Q4_K_M` |
+
+## Validation and Notes
+
+1. `OmniDiffusionConfig` receives `{"method": "gguf", "gguf_model": ...}`.
+2. `DiffusersPipelineLoader` resolves the GGUF file.
+3. A model-specific adapter remaps GGUF tensor names to vLLM-Omni transformer
+   names.
+4. Only transformer weights are loaded from GGUF. Missing non-transformer
+   weights are loaded from the base model repository.
+5. vLLM's GGUF linear method performs dequantization and GEMM at runtime.
+
+Unsupported models fail fast with a clear "No GGUF adapter matched" error
+instead of falling back to a generic mapper. Many GGUF repositories do not
+include `model_index.json`; always pass the normal base model through `--model`.

docs/user_guide/quantization/int8.md

@@ -0,0 +1,107 @@
+# Int8 Quantization
+
+## Overview
+
+Int8 quantization supports W8A8 diffusion transformer inference on CUDA and
+Ascend NPU. It can quantize BF16/FP16 weights at load time, or load serialized
+Int8 checkpoints that already contain quantized weights and scales.
+
+Only online activation scaling is currently supported.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ |
+| AMD ROCm | ⭕ |
+| Intel XPU | ⭕ |
+| Ascend NPU | ✅ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | HF models | CUDA | Ascend NPU | Mode | Recommendation |
+|-------|-----------|:----:|:----------:|------|----------------|
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Online W8A8 | All layers |
+| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Online W8A8 | Validate before enabling in docs |
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | Online W8A8 | All layers |
+
+Other diffusion models may work if their transformer uses supported linear
+layers, but they are not validated in this guide.
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| Qwen3-Omni | Thinker language-model stage | Not validated | Prefer checkpoint-supported ModelOpt FP8 or AutoRound paths |
+| Qwen3-TTS | TTS language-model stage | Not validated | No Int8 TTS stage support is documented |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Stage-specific transformer or DiT module | Not validated | Requires explicit stage routing |
+| GLM-Image | Stage-specific transformer or DiT module | Not validated | Requires quality comparison with BF16 |
+
+## Configuration
+
+Python API:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(model="<your-model>", quantization="int8")
+
+omni_with_skips = Omni(
+    model="<your-model>",
+    quantization_config={
+        "method": "int8",
+        "ignored_layers": ["<layer-name>"],
+    },
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+CLI:
+
+```bash
+python text_to_image.py --model <your-model> --quantization int8
+python text_to_image.py --model <your-model> --quantization int8 --ignored-layers "img_mlp"
+vllm serve <your-model> --omni --quantization int8
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | - | Quantization method (`"int8"`) |
+| `activation_scheme` | str | `"dynamic"` | `"dynamic"` selects online activation scaling; static is not supported |
+| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 |
+| `is_checkpoint_int8_serialized` | bool | `False` | Set by checkpoint config when loading serialized Int8 weights |
+
+## Validation and Notes
+
+Int8 quantization can be combined with cache acceleration:
+
+```python
+omni = Omni(
+    model="<your-model>",
+    quantization="int8",
+    cache_backend="tea_cache",
+    cache_config={"rel_l1_thresh": 0.2},
+)
+```
+
+Only add a new model to the supported table after comparing the Int8 output
+against a BF16 baseline and documenting any required `ignored_layers`.

docs/user_guide/quantization/msmodelslim.md

@@ -0,0 +1,95 @@
+# msModelSlim Quantization
+
+## Overview
+
+[msModelSlim](https://github.com/Ascend/msmodelslim) is an Ascend compression
+toolkit for producing pre-quantized model checkpoints. In vLLM-Omni, these
+checkpoints run through the Ascend/NPU path with `--quantization ascend`.
+
+msModelSlim is static quantization: quantized weights are generated offline
+before vLLM-Omni inference starts.
+
+## Hardware Support
+
+| Device | Support |
+|--------|---------|
+| NVIDIA Blackwell GPU (SM 100+) | ❌ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ❌ |
+| NVIDIA Ampere GPU (SM 80+) | ❌ |
+| AMD ROCm | ❌ |
+| Intel XPU | ❌ |
+| Ascend NPU | ✅ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Model | Base model | Scope | Hardware | Notes |
+|-------|------------|-------|----------|-------|
+| Wan2.2 | Wan2.2 diffusion weights | DiT or diffusion stage | Ascend NPU | Upstream msModelSlim provides a Wan2.2 quantization recipe; vLLM-Omni inference validation is not listed |
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | DiT or diffusion stage | Ascend NPU | Not validated in this guide |
+| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | DiT or diffusion stage | Ascend A2/A3 NPU | Generate quantized weights with the HunyuanImage-3.0 msModelSlim adaptation |
+
+Public Hugging Face quantized weights are not available yet. Use the
+[HunyuanImage-3.0 msModelSlim adaptation](https://gitcode.com/betta18/msmodelslim/tree/hyimage3_mxfp8)
+to generate the checkpoint manually.
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| Qwen3-Omni | Thinker or language-model stage | Not validated | No msModelSlim omni checkpoint path is documented |
+| Qwen3-TTS | TTS language-model stage | Not validated | No msModelSlim TTS checkpoint path is documented |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Model | Scope | Status | Notes |
+|-------|-------|--------|-------|
+| BAGEL | Stage-specific diffusion or transformer weights | Not validated | Requires a model-specific Ascend adaptation |
+| GLM-Image | Stage-specific diffusion or transformer weights | Not validated | Requires a model-specific Ascend adaptation |
+
+## Configuration
+
+Offline inference:
+
+```bash
+python text_to_image.py --model <quantized-model-path> --quantization ascend
+```
+
+Online serving:
+
+```bash
+vllm serve <quantized-model-path> --omni --quantization ascend
+```
+
+## Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `--quantization` | str | - | Use `ascend` for msModelSlim-produced checkpoints |
+| `model` | str | - | Path to the quantized checkpoint generated by Ascend tooling |
+
+Example msModelSlim command for a Wan2.2 W8A8 checkpoint:
+
+```bash
+msmodelslim quant \
+  --model_path /path/to/wan2_2_t2v_float_weights \
+  --save_path /path/to/wan2_2_t2v_quantized_weights \
+  --device npu \
+  --model_type Wan2_2 \
+  --config_path /path/to/wan2_2_w8a8f8_mxfp_t2v.yaml \
+  --trust_remote_code True
+```
+
+For HunyuanImage-3.0, use the Hunyuan-specific adaptation linked above.
+
+## Validation and Notes
+
+1. Run with the Ascend/NPU installation and environment.
+2. The `ascend` quantization method expects weights produced by the Ascend
+   tooling; it is not a load-time CUDA quantizer.
+3. Keep the quantized checkpoint aligned with the same model architecture and
+   stage config used for BF16 inference.

docs/user_guide/quantization/online.md

@@ -0,0 +1,93 @@
+# Online Quantization
+
+## Overview
+
+Online quantization means vLLM-Omni computes quantized weights and scales while
+loading the model. Use it when you want memory savings without preparing a
+separate quantized checkpoint.
+
+This mode is different from pre-quantized checkpoint formats such as GGUF,
+AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are
+prepared before serving and are documented in their method-specific guides.
+
+## Hardware Support
+
+| Device | FP8 W8A8 | Int8 W8A8 |
+|--------|----------|-----------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ |
+| AMD ROCm | ⭕ | ⭕ |
+| Intel XPU | ⭕ | ⭕ |
+| Ascend NPU | ❌ | ✅ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide. FP8 on Ampere may use a weight-only path where available.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+| Method | Guide | Example models | Status |
+|--------|-------|----------------|--------|
+| FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
+| Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+Online quantization is not currently validated for the omni/TTS stages. For
+Qwen3-Omni and related models, prefer checkpoint-declared ModelOpt or
+AutoRound paths when available.
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+Online quantization must be routed to the intended stage. BAGEL and GLM-Image
+need model-specific validation before they are listed as supported targets.
+
+## Configuration
+
+Python API:
+
+```python
+from vllm_omni import Omni
+
+omni_fp8 = Omni(model="<your-model>", quantization="fp8")
+omni_int8 = Omni(model="<your-model>", quantization="int8")
+```
+
+CLI:
+
+```bash
+vllm serve <your-model> --omni --quantization fp8
+vllm serve <your-model> --omni --quantization int8
+```
+
+Per-component routing:
+
+```python
+from vllm_omni.quantization import build_quant_config
+
+config = build_quant_config({
+    "transformer": {"method": "fp8"},
+    "vae": None,
+})
+```
+
+## Parameters
+
+| Parameter | Methods | Description |
+|-----------|---------|-------------|
+| `method` | FP8, Int8 | Quantization method, usually `"fp8"` or `"int8"` |
+| `ignored_layers` | FP8, Int8 | Layer name patterns to keep in BF16/FP16 |
+| `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling |
+| `weight_block_size` | FP8 | Optional block-wise FP8 weight quantization size |
+
+## Validation and Notes
+
+1. Compare the online-quantized output against a BF16 baseline with the same
+   seed and generation parameters.
+2. Use `ignored_layers` for quality-sensitive MLPs or output projections.
+3. Document any required skipped layers in the method page before marking a new
+   model as supported.
+4. If a model already ships quantized weights, use the matching pre-quantized
+   method guide instead of online quantization.

docs/user_guide/quantization/overview.md

@@ -0,0 +1,134 @@
+# Quantization
+
+vLLM-Omni exposes quantization through the unified `quantization_config`
+path. The same configuration entrypoint is used across diffusion-only models,
+multi-stage omni/TTS models, and multi-stage diffusion models, but each model
+type has a different quantization scope.
+
+## Quantization Modes
+
+| Mode | Guide | Description | Methods |
+|------|-------|-------------|---------|
+| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8 |
+| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | GGUF, AutoRound, msModelSlim, serialized Int8 |
+
+## Hardware Support
+
+| Device | FP8 W8A8 | Int8 W8A8 | GGUF | AutoRound | msModelSlim |
+|--------|----------|-----------|------|-----------|-------------|
+| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ❌ |
+| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ❌ |
+| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ✅ | ❌ |
+| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
+| Intel XPU | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
+| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ |
+
+Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
+guide. FP8 on Ampere may use a weight-only path where available.
+
+## Model Type Support
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+These models run a diffusion transformer as the primary inference module. The
+default quantization target is the transformer; tokenizer, scheduler, text
+encoder, and VAE stay on the base checkpoint unless a method guide says
+otherwise.
+
+| Method | Guide | Mode | Example models | Status |
+|--------|-------|------|----------------|--------|
+| FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
+| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
+| GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
+| AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
+| msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+These models combine an AR language model with audio, vision, talker, or TTS
+stages. Quantization is scoped to the AR language-model stage when the
+checkpoint contains a supported `quantization_config`; the non-AR stages stay
+in BF16 unless the model guide explicitly adds support.
+
+| Method | Guide | Scope | Example models | Status |
+|--------|-------|-------|----------------|--------|
+| FP8 | [FP8](fp8.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
+| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
+| GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
+| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
+| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+These models split generation across multiple stages. Quantization must be
+attached to the intended stage rather than applied globally.
+
+| Method | Guide | Scope | Example models | Status |
+|--------|-------|-------|----------------|--------|
+| FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
+| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
+| GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
+| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
+| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
+
+!!! note
+    "Online quantization" means vLLM-Omni computes the quantization data while
+    loading the model. "Pre-quantized" means the checkpoint or external
+    quantizer provides the required quantized weights and scales.
+
+## Quantization Scope
+
+### Diffusion Model (Qwen-Image, Wan2.2)
+
+The default target is the diffusion transformer. Component routing is available
+through `build_quant_config()`:
+
+```python
+from vllm_omni.quantization import build_quant_config
+
+config = build_quant_config({
+    "transformer": {"method": "fp8"},
+    "vae": None,
+})
+```
+
+| Component | Default quantized? | Notes |
+|-----------|--------------------|-------|
+| Diffusion transformer | Yes | Primary target for FP8, Int8, GGUF, AutoRound, and msModelSlim |
+| Text encoder | No | Keep BF16 unless a method-specific guide documents support |
+| VAE | No | Keep BF16; storage-only paths are method-specific |
+| Scheduler/tokenizer | No | Loaded from the base model repository |
+
+### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
+
+| Component | Default quantized? | Notes |
+|-----------|--------------------|-------|
+| Thinker or AR language model | Yes, when checkpoint config is supported | ModelOpt FP8/NVFP4 or AutoRound checkpoint config |
+| Audio encoder | No | BF16 |
+| Vision encoder | No | BF16 |
+| Talker or TTS stage | No | BF16 unless model-specific support is documented |
+| Code2Wav | No | BF16 |
+
+### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
+
+| Component | Default quantized? | Notes |
+|-----------|--------------------|-------|
+| Selected diffusion or transformer stage | Method-specific | Must be routed to the intended stage |
+| Other generation stages | No | Keep BF16 unless separately validated |
+| VAE, tokenizer, scheduler | No | Loaded from the base checkpoint |
+
+## Python API
+
+`build_quant_config()` accepts strings, dictionaries, per-component
+dictionaries, existing `QuantizationConfig` objects, or `None`.
+
+```python
+from vllm_omni.quantization import build_quant_config
+
+build_quant_config("fp8")
+build_quant_config({"method": "fp8", "activation_scheme": "static"})
+build_quant_config("auto-round", bits=4, group_size=128)
+build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"})
+build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
+build_quant_config(None)
+```

examples/offline_inference/qwen3_tts/end2end.py

@@ -83,12 +83,29 @@ def _estimate_prompt_len(
             if not isinstance(audio_path, str) or not audio_path.strip():
                 return None
             try:
-                from vllm.multimodal.media import MediaConnector
+                from urllib.parse import urlparse
 
-                connector = MediaConnector(allowed_local_media_path="/")
-                audio, sr = connector.fetch_audio(audio_path)
                 import numpy as np
 
+                def _is_url(path: str) -> bool:
+                    try:
+                        parsed = urlparse(path)
+                        if parsed.scheme in ("http", "https"):
+                            return bool(parsed.netloc)
+                        return parsed.scheme in ("file", "data")
+                    except Exception:
+                        return False
+
+                if _is_url(audio_path):
+                    from vllm.multimodal.media import MediaConnector
+
+                    connector = MediaConnector(allowed_local_media_path="/")
+                    audio, sr = connector.fetch_audio(audio_path)
+                else:
+                    from vllm.multimodal.media.audio import load_audio
+
+                    audio, sr = load_audio(audio_path, sr=None, mono=True)
+
                 wav_np = np.asarray(audio, dtype=np.float32)
 
                 if speech_tok is not None:

examples/online_serving/stable_audio/README.md

@@ -0,0 +1,234 @@
+# Stable Audio Online Serving
+
+Generate audio from text prompts using Stable Audio models via an OpenAI-compatible API endpoint.
+
+## Features
+
+- **OpenAI-compatible API**: Use `/v1/audio/generate` endpoint
+- **Flexible control**: Adjust audio length, guidance scale, inference steps
+- **Quality control**: Use negative prompts to avoid unwanted characteristics
+- **Reproducible**: Set random seed for deterministic generation
+
+## Quick Start
+
+### 1. Start the Server
+
+```bash
+vllm-omni serve stabilityai/stable-audio-open-1.0 \
+    --host 0.0.0.0 \
+    --port 8091 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
+```
+
+### 2. Generate Audio
+
+#### Using curl
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a cat purring",
+        "audio_length": 10.0
+    }' --output cat.wav
+```
+
+#### Using Python Client
+
+```bash
+python stable_audio_client.py \
+    --text "The sound of a cat purring" \
+    --audio_length 10.0 \
+    --output cat.wav
+```
+
+#### Using Bash Script
+
+```bash
+bash curl_examples.sh
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/generate
+```
+
+### Request Body
+
+```json
+{
+    "input": "Text description of the audio",
+    "audio_length": 10.0,
+    "audio_start": 0.0,
+    "negative_prompt": "Low quality",
+    "guidance_scale": 7.0,
+    "num_inference_steps": 100,
+    "seed": 42,
+    "response_format": "wav"
+}
+```
+
+### Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text prompt describing the audio to generate |
+| `audio_length` | float | ~47s | Audio duration in seconds (max ~47s for stable-audio-open-1.0) |
+| `audio_start` | float | 0.0 | Audio start time in seconds |
+| `negative_prompt` | string | null | Text describing what to avoid in generation |
+| `guidance_scale` | float | 7.0 | Classifier-free guidance scale (higher = more adherence to prompt) |
+| `num_inference_steps` | int | 50 | Number of denoising steps (higher = better quality, slower) |
+| `seed` | int | null | Random seed for reproducibility |
+| `response_format` | string | "wav" | Output format: wav, mp3, flac, pcm |
+
+### Response
+
+Returns audio data in the requested format (default: WAV).
+
+## Usage Examples
+
+### Basic Generation
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of ocean waves"
+    }' --output ocean.wav
+```
+
+### Custom Duration
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+```
+
+### High Quality with Negative Prompt
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy",
+        "guidance_scale": 8.0,
+        "num_inference_steps": 150
+    }' --output piano_hq.wav
+```
+
+### Reproducible Generation
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "seed": 42
+    }' --output thunder.wav
+```
+
+### Quick Generation (Fewer Steps)
+
+For faster generation with slightly lower quality:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Birds chirping in a forest",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output birds_quick.wav
+```
+
+## Python Client Examples
+
+### Simple Generation
+
+```bash
+python stable_audio_client.py \
+    --text "The sound of a cat purring"
+```
+
+### Custom Parameters
+
+```bash
+python stable_audio_client.py \
+    --text "Thunder and rain" \
+    --audio_length 15.0 \
+    --negative_prompt "Low quality" \
+    --guidance_scale 7.0 \
+    --num_inference_steps 100 \
+    --seed 42 \
+    --output thunder.wav
+```
+
+### Different Output Format
+
+```bash
+python stable_audio_client.py \
+    --text "Guitar playing" \
+    --response_format mp3 \
+    --output guitar.mp3
+```
+
+## Tips
+
+1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`
+2. **Quality vs Speed**:
+   - 50 steps: Fast, decent quality
+   - 100 steps: Good balance (default)
+   - 150+ steps: High quality, slower
+3. **Guidance Scale**:
+   - Lower (3-5): More creative/varied
+   - Default (7): Good balance
+   - Higher (10+): More literal to prompt
+4. **Negative Prompts**: Use to avoid "Low quality", "distorted", "noisy", etc.
+5. **Seeds**: Use same seed for reproducible results
+
+## Performance
+
+| Inference Steps | Quality | Speed | Use Case |
+|----------------|---------|-------|----------|
+| 50 | Good | Fast | Quick previews |
+| 100 (default) | Very Good | Medium | Production |
+| 150+ | Excellent | Slow | Final/critical audio |
+
+## Troubleshooting
+
+### Server not responding
+- Check if server is running: `curl http://localhost:8091/health`
+- Check server logs for errors
+
+### Audio quality issues
+- Increase `num_inference_steps` (e.g., 150)
+- Add negative prompts: `"Low quality, distorted, noisy"`
+- Increase `guidance_scale` for more prompt adherence
+
+### Generation timeout
+- Reduce `num_inference_steps`
+- Reduce `audio_length`
+- Check GPU memory with `nvidia-smi`
+
+### Wrong audio length
+- Ensure `audio_length` is within model limits (~47s max)
+- Adjust `audio_start` if trimming is needed
+
+## See Also
+
+- [Offline Inference Example](../../offline_inference/text_to_audio/README.md)
+- [Stable Audio Model Card](https://huggingface.co/stabilityai/stable-audio-open-1.0)
+- [vLLM-Omni Documentation](https://github.com/vllm-project/vllm-omni)

examples/online_serving/stable_audio/curl_examples.sh

@@ -0,0 +1,54 @@
+#!/bin/bash
+# Examples for using Stable Audio with curl via /v1/audio/generate endpoint
+
+# Example 1: Simple request with default parameters
+echo "Example 1: Simple request with default parameters"
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound audience clapping and cheering in a stadium"
+    }' --output stadium.wav
+
+# Example 2: Request with custom audio_length
+echo "Example 2: Custom audio length (5 seconds)"
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "The sound of a dog barking",
+        "audio_length": 5.0
+    }' --output dog_5s.wav
+
+# Example 3: Request with negative prompt for quality control
+echo "Example 3: With negative prompt"
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "A piano playing a gentle melody",
+        "audio_length": 10.0,
+        "negative_prompt": "Low quality, distorted, noisy"
+    }' --output piano.wav
+
+# Example 4: Full control with all parameters
+echo "Example 4: Full control (custom length, guidance, steps, seed)"
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Thunder and rain sounds",
+        "audio_length": 15.0,
+        "negative_prompt": "Low quality",
+        "guidance_scale": 7.0,
+        "num_inference_steps": 100,
+        "seed": 42
+    }' --output thunder_rain.wav
+
+# Example 5: Quick generation with fewer steps (faster but lower quality)
+echo "Example 5: Quick generation (fewer steps)"
+curl -X POST http://localhost:8091/v1/audio/generate \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Ocean waves crashing on a beach",
+        "audio_length": 8.0,
+        "num_inference_steps": 50
+    }' --output ocean.wav
+
+echo "All examples completed!"

examples/online_serving/stable_audio/stable_audio_client.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""
+OpenAI-compatible client for Stable Audio via /v1/audio/generate endpoint.
+
+This script demonstrates how to use the OpenAI-compatible speech API
+to generate audio from text using Stable Audio models.
+
+Examples:
+    # Simple generation
+    python stable_audio_client.py --text "The sound of a cat purring"
+
+    # With custom duration
+    python stable_audio_client.py --text "A dog barking" --audio_length 5.0
+
+    # With all parameters
+    python stable_audio_client.py --text "Thunder and rain" \
+        --audio_length 15.0 \
+        --negative_prompt "Low quality" \
+        --guidance_scale 7.0 \
+        --num_inference_steps 100 \
+        --seed 42 \
+        --output thunder.wav
+"""
+
+import argparse
+import sys
+
+import requests
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Generate audio with Stable Audio via OpenAI-compatible API")
+    parser.add_argument(
+        "--api_url",
+        default="http://localhost:8091/v1/audio/generate",
+        help="API endpoint URL",
+    )
+    parser.add_argument(
+        "--text",
+        default="The sound of a cat purring",
+        help="Text prompt for audio generation",
+    )
+    parser.add_argument(
+        "--audio_length",
+        type=float,
+        default=10.0,
+        help="Audio length in seconds (max ~47s for stable-audio-open-1.0)",
+    )
+    parser.add_argument(
+        "--audio_start",
+        type=float,
+        default=0.0,
+        help="Audio start time in seconds",
+    )
+    parser.add_argument(
+        "--negative_prompt",
+        default="Low quality",
+        help="Negative prompt for classifier-free guidance",
+    )
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=7.0,
+        help="Guidance scale for diffusion (higher = more adherence to prompt)",
+    )
+    parser.add_argument(
+        "--num_inference_steps",
+        type=int,
+        default=100,
+        help="Number of inference steps (higher = better quality, slower)",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Random seed for reproducibility",
+    )
+    parser.add_argument(
+        "--output",
+        default="stable_audio_output.wav",
+        help="Output file path",
+    )
+    parser.add_argument(
+        "--response_format",
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm"],
+        help="Audio output format",
+    )
+    return parser.parse_args()
+
+
+def generate_audio(args):
+    """Generate audio using the API."""
+
+    # Build request payload
+    payload = {
+        "input": args.text,
+        "audio_length": args.audio_length,
+        "audio_start": args.audio_start,
+        "response_format": args.response_format,
+    }
+
+    # Add optional parameters
+    if args.negative_prompt:
+        payload["negative_prompt"] = args.negative_prompt
+    if args.guidance_scale:
+        payload["guidance_scale"] = args.guidance_scale
+    if args.num_inference_steps:
+        payload["num_inference_steps"] = args.num_inference_steps
+    if args.seed is not None:
+        payload["seed"] = args.seed
+
+    print(f"\n{'=' * 60}")
+    print("Stable Audio - Text-to-Audio Generation")
+    print(f"{'=' * 60}")
+    print(f"API URL: {args.api_url}")
+    print(f"Prompt: {args.text}")
+    print(f"Audio length: {args.audio_length}s")
+    print(f"Negative prompt: {args.negative_prompt}")
+    print(f"Guidance scale: {args.guidance_scale}")
+    print(f"Inference steps: {args.num_inference_steps}")
+    if args.seed is not None:
+        print(f"Seed: {args.seed}")
+    print(f"Output: {args.output}")
+    print(f"{'=' * 60}\n")
+
+    try:
+        # Make the API request
+        print("Generating audio...")
+        response = requests.post(
+            args.api_url,
+            json=payload,
+            headers={"Content-Type": "application/json"},
+            timeout=300,  # 5 minute timeout for long generations
+        )
+
+        # Check for errors
+        if response.status_code != 200:
+            print(f"Error: API returned status code {response.status_code}")
+            print(f"Response: {response.text}")
+            return False
+
+        # Save the audio
+        with open(args.output, "wb") as f:
+            f.write(response.content)
+
+        print(f"✓ Audio saved to {args.output}")
+        print(f"  File size: {len(response.content) / 1024:.1f} KB")
+        return True
+
+    except requests.exceptions.Timeout:
+        print("Error: Request timed out. Try reducing inference steps or audio length.")
+        return False
+    except requests.exceptions.ConnectionError:
+        print(f"Error: Could not connect to {args.api_url}")
+        print("Make sure the server is running.")
+        return False
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+
+
+def main():
+    args = parse_args()
+    success = generate_audio(args)
+    sys.exit(0 if success else 1)
+
+
+if __name__ == "__main__":
+    main()

mkdocs.yml

@@ -97,7 +97,6 @@ plugins:
       exclude:
         - "re:vllm_omni\\._.*"  # Internal modules
         - "vllm_omni.diffusion.models.qwen_image"  # avoid importing vllm in mkdocs building
-        - "vllm_omni.diffusion.models.dreamid_omni.wan2_2"  # docstring signature warnings break strict docs
         - "vllm_omni.diffusion.quantization"  # avoid importing vllm in mkdocs building
         - "vllm_omni.quantization"  # avoid importing vllm in mkdocs building
         - "vllm_omni.entrypoints.async_diffusion"  # avoid importing vllm in mkdocs building

pyproject.toml

@@ -49,7 +49,7 @@ dev = [
     "soundfile>=0.13.1",
     "imageio[ffmpeg]>=0.6.0",
     "opencv-python>=4.12.0.88",
-    "mooncake-transfer-engine==0.3.8.post1",
+    "mooncake-transfer-engine-cuda13>=0.3.9",
     "av", # for ComfyUI tests
     "openpyxl>=3.0.0", # for nightly CI
     "pyttsx3>=2.99",

requirements/cuda.txt

@@ -1,5 +1,5 @@
 -r common.txt
 onnxruntime>=1.23.2
-fa3-fwd==0.0.2
+fa3-fwd==0.0.3
 # nvidia-cublas-cu12 is pinned to 12.9.1.4 via force-reinstall in Dockerfile.ci
 # because torch pins its own cublas version, making it unsatisfiable at resolve time.

tests/benchmarks/test_seed_tts_dataset_variants.py

@@ -8,8 +8,10 @@ from __future__ import annotations
 
 import importlib.util
 import sys
+import types
 from pathlib import Path
 
+import numpy as np
 import pytest
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -167,3 +169,58 @@ def test_attach_sets_seed_tts_row_even_without_extra_body():
     row_pos = src.index("seed_tts_row")
     not_ex_pos = src.index("if not ex:")
     assert row_pos < not_ex_pos, "seed_tts_row must be set before 'if not ex: return'"
+
+
+def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
+    from vllm_omni.benchmarks.data_modules import seed_tts_eval
+
+    calls = {}
+
+    class FakeTensor:
+        def __init__(self, name: str):
+            self.name = name
+            self.device = None
+
+        def to(self, device):
+            self.device = device
+            return self
+
+    class FakeProcessor:
+        def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False):
+            calls["return_attention_mask"] = return_attention_mask
+            assert sampling_rate == 16000
+            assert return_tensors == "pt"
+            assert len(wav) > 0
+            return types.SimpleNamespace(
+                input_features=FakeTensor("features"),
+                attention_mask=FakeTensor("mask") if return_attention_mask else None,
+            )
+
+        def get_decoder_prompt_ids(self, *, language, task):
+            assert language == "english"
+            assert task == "transcribe"
+            return [(1, 2)]
+
+        def batch_decode(self, predicted_ids, *, skip_special_tokens):
+            assert skip_special_tokens
+            assert predicted_ids == [[42]]
+            return ["hello"]
+
+    class FakeModel:
+        def generate(self, input_features, **kwargs):
+            calls["input_features"] = input_features
+            calls["generate_kwargs"] = kwargs
+            return [[42]]
+
+    monkeypatch.setattr(seed_tts_eval, "_ensure_en_asr", lambda: None)
+    monkeypatch.setattr(seed_tts_eval, "_en_processor", FakeProcessor())
+    monkeypatch.setattr(seed_tts_eval, "_en_model", FakeModel())
+    monkeypatch.setattr(seed_tts_eval, "_device", "cuda:1")
+
+    text = seed_tts_eval._transcribe_en_f32_16k(np.ones(1600, dtype=np.float32))
+
+    assert text == "hello"
+    assert calls["return_attention_mask"] is True
+    assert calls["input_features"].device == "cuda:1"
+    assert calls["generate_kwargs"]["attention_mask"].device == "cuda:1"
+    assert calls["generate_kwargs"]["forced_decoder_ids"] == [(1, 2)]

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -456,8 +456,8 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
         client_index=0,
         take_events=lambda: [],
         trace_headers=None,
-        num_cached_tokens=0,
-        num_external_computed_tokens=0,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
         num_nans_in_logits=0,
         get_finished_reason=lambda: "stop",
     )
@@ -537,8 +537,8 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
         client_index=0,
         take_events=lambda: [],
         trace_headers=None,
-        num_cached_tokens=0,
-        num_external_computed_tokens=0,
+        has_encoder_inputs=False,
+        take_prefill_stats=lambda: None,
         num_nans_in_logits=0,
         get_finished_reason=lambda: "stop",
     )

tests/e2e/offline_inference/test_cosyvoice3.py

@@ -13,10 +13,8 @@ verifies a stable reference recipe:
 from __future__ import annotations
 
 import functools
-import io
 import os
 from pathlib import Path
-from urllib.request import urlopen
 
 import numpy as np
 import pytest
@@ -25,6 +23,7 @@ from huggingface_hub import snapshot_download
 from vllm.sampling_params import SamplingParams
 
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import test_asset_path
 from tests.helpers.runtime import OmniRunner
 from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
@@ -33,7 +32,9 @@ from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokeni
 MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
 MODEL_DIR_ENV = "VLLM_OMNI_COSYVOICE3_MODEL_DIR"
 
-REFERENCE_PROMPT_WAV_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
+# Vendored under tests/assets/cosyvoice3/ so the test does not depend on
+# raw.githubusercontent.com being reachable from CI runners.
+REFERENCE_PROMPT_WAV_PATH = test_asset_path("cosyvoice3/zero_shot_prompt.wav")
 REFERENCE_PROMPT_TEXT = "You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。"
 REFERENCE_SYNTH_TEXT = (
     "CosyVoice is undergoing a comprehensive upgrade, providing more accurate, "
@@ -51,9 +52,7 @@ ASYNC_CHUNK_MODES = [
 
 @functools.lru_cache(maxsize=1)
 def _load_reference_prompt_wav() -> tuple[np.ndarray, int]:
-    with urlopen(REFERENCE_PROMPT_WAV_URL, timeout=30) as resp:
-        data = resp.read()
-    audio, sr = sf.read(io.BytesIO(data), dtype="float32", always_2d=False)
+    audio, sr = sf.read(str(REFERENCE_PROMPT_WAV_PATH), dtype="float32", always_2d=False)
     if isinstance(audio, np.ndarray) and audio.ndim > 1:
         audio = np.mean(audio, axis=-1)
     return np.asarray(audio, dtype=np.float32), int(sr)

tests/e2e/offline_inference/test_qwen3_tts_base.py

@@ -16,10 +16,12 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 import pytest
 
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
 from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
-REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
+# See tests/e2e/online_serving/test_qwen3_tts_base.py for the vendored-asset rationale.
+REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
 REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
 
 

tests/e2e/online_serving/test_bagel_expansion.py

@@ -15,18 +15,57 @@ assert_diffusion_response validates successful generation and the expected
 512x512 resolution.
 """
 
+import json
+
 import pytest
 
 from tests.helpers.mark import hardware_marks
 from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data
+from tests.helpers.stage_config import get_deploy_config_path
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
 
+_BAGEL_DEFAULT_YAML = get_deploy_config_path("ci/bagel.yaml")
+
 PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution."
 NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
 
 SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
 PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+TP_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=3)
+
+
+def _make_tp_cases(model: str, tp_size: int):
+    """Build Bagel TP test cases with devices auto-derived from tp_size.
+    Devices can not be set through CLI args, so we set them in the YAML.
+    """
+    # Stage 0 uses GPU 0, so place DiT TP ranks on GPU 1..N to
+    # avoid AR/DiT memory contention on one device.
+    devices = ",".join(str(i + 1) for i in range(tp_size))
+    stage_overrides = json.dumps(
+        {
+            "0": {"tensor_parallel_size": 1},
+            "1": {"devices": devices},
+        }
+    )
+    return [
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                stage_config_path=_BAGEL_DEFAULT_YAML,
+                server_args=[
+                    "--stage-overrides",
+                    stage_overrides,
+                    "--cache-backend",
+                    "cache_dit",
+                    "--tensor-parallel-size",
+                    str(tp_size),
+                ],
+            ),
+            id=f"parallel_tp_{tp_size}",
+            marks=TP_FEATURE_MARKS,
+        ),
+    ]
 
 
 def _get_diffusion_feature_cases(model: str):
@@ -75,19 +114,9 @@ def _get_diffusion_feature_cases(model: str):
             marks=PARALLEL_FEATURE_MARKS,
         ),
         # Tensor-Parallel size 2 (2 GPUs, Cache-DiT backend)
-        pytest.param(
-            OmniServerParams(
-                model=model,
-                server_args=[
-                    "--cache-backend",
-                    "cache_dit",
-                    "--tensor-parallel-size",
-                    "2",
-                ],
-            ),
-            id="parallel_tp_2",
-            marks=[*PARALLEL_FEATURE_MARKS, pytest.mark.skip(reason="issue: #2862")],
-        ),
+        # Stage 1 (DiT) needs visible GPUs matching TP size; the default YAML
+        # only exposes device "0", so we patch it here.
+        *_make_tp_cases(model, tp_size=2),
         # Ulysses-SP degree=2 (2 GPUs)
         pytest.param(
             OmniServerParams(

tests/e2e/online_serving/test_cosyvoice3_tts.py

@@ -15,13 +15,16 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 import pytest
 
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
 
-# Official CosyVoice zero-shot prompt audio and its transcript
-REF_AUDIO_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
+# Official CosyVoice zero-shot prompt audio and its transcript. Vendored under
+# tests/assets/ so the server does not depend on raw.githubusercontent.com being
+# reachable at request time (same rationale as issue #3263 for Qwen3-TTS).
+REF_AUDIO_URL = load_test_audio_data_url("cosyvoice3/zero_shot_prompt.wav")
 REF_TEXT = "希望你以后能够做的比我还好呦。"
 
 

tests/e2e/online_serving/test_qwen3_omni.py

@@ -180,6 +180,7 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", prefix_test_params, indirect=True)
+@pytest.mark.skip(reason="Flaky: identical greedy requests diverge on text_content (prefix caching path).")
 def test_thinker_prefix_caching(omni_server, openai_client) -> None:
     """
     Test thinker prefix caching by sending identical requests with an image (i.e.,

tests/e2e/online_serving/test_qwen3_tts_base.py

@@ -15,12 +15,16 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 import pytest
 
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
 
-REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
+# Vendored under tests/assets/qwen3_tts/clone_2.wav so the server does not need
+# to fetch the reference audio over HTTPS at request time (CI runners in some
+# regions cannot reach the original Aliyun OSS host; see issue #3263).
+REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
 REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
 
 

tests/e2e/online_serving/test_qwen3_tts_base_expansion.py

@@ -12,6 +12,7 @@ import os
 import pytest
 
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
 
@@ -22,7 +23,8 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base"
 
-REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
+# See tests/e2e/online_serving/test_qwen3_tts_base.py for the vendored-asset rationale.
+REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
 REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
 
 

tests/engine/test_arg_utils.py

@@ -175,23 +175,29 @@ def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch):
 
 
 def test_stage_specific_text_config_override():
-    """Ensure dependent attributes are updated when using stage-specific config."""
+    """Stage swap must refresh hf_text_config, dependent attrs, and model_arch_config."""
     vllm_config = EngineArgs().create_model_config()
     vllm_config.disable_sliding_window = True
+    thinker_mac = vllm_config.model_arch_config
+
+    talker_num_heads = max(2, thinker_mac.total_num_attention_heads // 2)
+    talker_num_kv_heads = max(1, talker_num_heads // 8)
+    talker_head_dim = 128
+    stage_text_config = SimpleNamespace(
+        sliding_window=4096,
+        attention_chunk_size=2048,
+        max_position_embeddings=4096,
+        num_attention_heads=talker_num_heads,
+        num_key_value_heads=talker_num_kv_heads,
+        head_dim=talker_head_dim,
+        hidden_size=talker_num_heads * talker_head_dim,
+        vocab_size=thinker_mac.vocab_size,
+        num_hidden_layers=4,
+    )
 
-    # Switch the created hf text config with a mock whose
-    # values we want to pull through the text config helper
-    stage_text_config = vllm_config.hf_text_config
     vllm_config.hf_text_config = SimpleNamespace()
-    stage_text_config.sliding_window = 4096
-    stage_text_config.attention_chunk_size = 2048
+    vllm_config.hf_config.thinker_config = SimpleNamespace(get_text_config=lambda: stage_text_config)
 
-    # Move the stage config's text config getter & thinker config
-    mock_stage_config = SimpleNamespace(get_text_config=lambda: stage_text_config)
-    vllm_config.hf_config.thinker_config = mock_stage_config
-
-    # Ensure that create from a vLLM config correctly pulls the
-    # expected values off of the thinker config & swaps the text config
     omni_config = OmniModelConfig.from_vllm_model_config(
         vllm_config,
         hf_config_name="thinker_config",
@@ -202,6 +208,21 @@ def test_stage_specific_text_config_override():
     assert omni_config.max_model_len == 4096
     assert omni_config.hf_text_config.sliding_window is None
 
+    stage_mac = omni_config.model_arch_config
+    assert stage_mac is not thinker_mac
+    assert stage_mac.total_num_attention_heads == talker_num_heads
+    assert stage_mac.total_num_kv_heads == talker_num_kv_heads
+    assert stage_mac.head_size == talker_head_dim
+
+    parallel_config = SimpleNamespace(
+        tensor_parallel_size=1,
+        pipeline_parallel_size=1,
+        decode_context_parallel_size=1,
+    )
+    assert omni_config.get_num_attention_heads(parallel_config) == talker_num_heads
+    assert omni_config.get_num_kv_heads(parallel_config) == talker_num_kv_heads
+    assert omni_config.get_head_size() == talker_head_dim
+
 
 def test_stage_configs_path_field():
     """OmniEngineArgs with stage_configs_path should construct without error."""
@@ -230,7 +251,7 @@ def test_voxcpm_model_arch_injects_model_type_override(mocker):
 def test_strip_single_engine_args():
     """_strip_single_engine_args should remove EngineArgs fields but keep omni fields."""
     kwargs = {
-        # Parent EngineArgs fields — should be stripped
+        # Parent EngineArgs fields — stripped unless explicitly allowlisted
         "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
         "tensor_parallel_size": 4,
         "gpu_memory_utilization": 0.9,
@@ -248,7 +269,7 @@ def test_strip_single_engine_args():
 
     # Stripped — parent EngineArgs fields
     assert "compilation_config" not in filtered
-    assert "tensor_parallel_size" not in filtered
+    assert filtered["tensor_parallel_size"] == 4
     assert "gpu_memory_utilization" not in filtered
     assert "model" not in filtered
 
@@ -278,15 +299,18 @@ def test_strip_single_engine_args_model_does_not_trigger_warning(mocker):
     mock_warn.assert_not_called()
 
     # When there *are* genuinely surprising overrides alongside model,
-    # the warning should mention them but not model.
+    # the warning should mention them but not model. Keep-listed fields such as
+    # tensor_parallel_size are intentionally passed through and should not warn.
     AsyncOmniEngine._strip_single_engine_args(
         {
             "model": "some/model",
+            "compilation_config": '{"cudagraph_mode": "FULL_AND_PIECEWISE"}',
             "tensor_parallel_size": 4,
             "custom_pipeline_args": {"pipeline_class": "my.Pipeline"},
         }
     )
     mock_warn.assert_called_once()
     warned_args = mock_warn.call_args[0][-1]  # the formatted arg list
-    assert "tensor_parallel_size" in warned_args
+    assert "compilation_config" in warned_args
+    assert "tensor_parallel_size" not in warned_args
     assert "model" not in warned_args

tests/engine/test_async_omni_engine_abort.py

@@ -9,11 +9,13 @@ from vllm.inputs import PromptType
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.config.stage_config import _PIPELINE_REGISTRY
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 SEED = 42
+pipeline = _PIPELINE_REGISTRY["qwen2_5_omni_thinker_only"]
 
 # Single-stage thinker-only deploy, materialized from tests.helpers.stage_config._CI_OVERLAYS.
 stage_config = get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml")

tests/engine/test_single_stage_mode.py

@@ -1137,13 +1137,20 @@ class TestLaunchDiffusionStage:
             "vllm_omni.engine.async_omni_engine.StageDiffusionClient.from_addresses",
             return_value=diffusion_client,
         )
+        mock_acquire_locks = mocker.patch(
+            "vllm_omni.engine.async_omni_engine.acquire_diffusion_device_locks",
+            return_value=[101],
+        )
+        mock_release_locks = mocker.patch("vllm_omni.engine.async_omni_engine.release_device_locks")
 
         result = engine._launch_diffusion_stage(
             stage_cfg=stage_cfg,
             metadata=metadata,
             omni_master_server=omni_master_server,
+            stage_init_timeout=60,
         )
 
+        mock_acquire_locks.assert_called_once_with(5, "diffusion-config", 60)
         mock_register.assert_called_once_with(
             omni_master_address="127.0.0.1",
             omni_master_port=25000,
@@ -1158,7 +1165,8 @@ class TestLaunchDiffusionStage:
             request_address="tcp://127.0.0.1:25002",
             response_address="tcp://127.0.0.1:25003",
         )
-        mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:25001")
+        mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:25001", 60)
+        mock_release_locks.assert_called_once_with([101])
         mock_from_addresses.assert_called_once_with(
             metadata,
             request_address="tcp://127.0.0.1:25002",
@@ -1205,7 +1213,7 @@ class TestCreateRemoteLlmStage:
 
         @contextmanager
         def fake_connect_cm(*args, **kwargs):
-            yield eng_mgr, coordinator, fake_addresses
+            yield eng_mgr, coordinator, fake_addresses, None
 
         mocker.patch(
             "vllm_omni.engine.async_omni_engine.build_engine_args_dict",
@@ -1265,7 +1273,7 @@ class TestCreateRemoteLlmStage:
 
         @contextmanager
         def fake_connect_cm(*args, **kwargs):
-            yield mocker.Mock(), mocker.Mock(), fake_addresses
+            yield mocker.Mock(), mocker.Mock(), fake_addresses, None
 
         mocker.patch(
             "vllm_omni.engine.async_omni_engine.build_engine_args_dict",
@@ -1329,7 +1337,7 @@ class TestCreateRemoteLlmStage:
 
         @contextmanager
         def boom(*args, **kwargs):
-            yield mocker.Mock(), mocker.Mock(), mocker.Mock()
+            yield mocker.Mock(), mocker.Mock(), mocker.Mock(), None
             raise RuntimeError("handshake failed")
 
         mocker.patch(
@@ -1402,7 +1410,7 @@ class TestConnectRemoteEngineCoresCoordinator:
             vllm_config=vllm_config,
             omni_master_server=omni_master_server,
             stage_id=7,
-        ) as (_, yielded_coordinator, yielded_addresses):
+        ) as (_, yielded_coordinator, yielded_addresses, _tensor_queue):
             assert yielded_coordinator is None
             assert yielded_addresses.coordinator_input == "tcp://coord-in"
             assert yielded_addresses.coordinator_output == "tcp://coord-out"
@@ -1439,7 +1447,7 @@ class TestConnectRemoteEngineCoresCoordinator:
             vllm_config=vllm_config,
             omni_master_server=omni_master_server,
             stage_id=7,
-        ) as (_, yielded_coordinator, yielded_addresses):
+        ) as (_, yielded_coordinator, yielded_addresses, _tensor_queue):
             assert yielded_coordinator is None
             assert yielded_addresses.coordinator_input is None
             assert yielded_addresses.coordinator_output is None

tests/entrypoints/openai_api/test_serving_audio_generate.py

@@ -0,0 +1,509 @@
+# tests/entrypoints/openai_api/test_serving_audio_generate.py
+import logging
+from inspect import Signature, signature
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from vllm_omni.entrypoints.openai.protocol.audio import (
+    CreateAudio,
+    OpenAICreateAudioGenerateRequest,
+)
+from vllm_omni.entrypoints.openai.serving_audio_generate import (
+    OmniOpenAIServingAudioGenerate,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+logger = logging.getLogger(__name__)
+
+
+# Helper: create a mock audio output for endpoint tests
+def create_mock_audio_output(
+    request_id: str = "audiogen-mock-123",
+    sample_rate: int = 44100,
+    num_samples: int = 44100,
+    audio_key: str = "audio",
+) -> OmniRequestOutput:
+    """Return an OmniRequestOutput mimicking diffusion audio model output."""
+
+    audio_tensor = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, num_samples))
+
+    return OmniRequestOutput.from_diffusion(
+        request_id=request_id,
+        images=[],
+        prompt=None,
+        metrics={},
+        multimodal_output={audio_key: audio_tensor, "sr": sample_rate},
+    )
+
+
+def _make_engine_client(*, audio_key: str = "audio", sample_rate: int = 44100):
+    """Build a mock engine client producing audio output."""
+    mock_engine_client = MagicMock()
+    mock_engine_client.errored = False
+    mock_engine_client.model_type = "StableAudioPipeline"
+    mock_engine_client.default_sampling_params_list = [{}]
+
+    async def mock_generate_fn(*args, **kwargs):
+        yield create_mock_audio_output(
+            request_id=kwargs.get("request_id", "audiogen-mock"),
+            sample_rate=sample_rate,
+            audio_key=audio_key,
+        )
+
+    mock_engine_client.generate = MagicMock(side_effect=mock_generate_fn)
+    return mock_engine_client
+
+
+def _make_server(engine_client=None):
+    """Build an OmniOpenAIServingAudioGenerate with mocks."""
+    if engine_client is None:
+        engine_client = _make_engine_client()
+
+    mock_models = MagicMock()
+    mock_models.is_base_model.return_value = True
+
+    return OmniOpenAIServingAudioGenerate(
+        engine_client=engine_client,
+        models=mock_models,
+        request_logger=MagicMock(),
+    )
+
+
+@pytest.fixture
+def test_app():
+    server = _make_server()
+
+    original_fn = server.create_audio_generate
+    sig = signature(original_fn)
+    new_params = [p for name, p in sig.parameters.items() if name != "raw_request"]
+    new_sig = Signature(parameters=new_params, return_annotation=sig.return_annotation)
+
+    async def patched_create_audio_generate(*args, **kwargs):
+        return await original_fn(*args, **kwargs)
+
+    patched_create_audio_generate.__signature__ = new_sig
+    server.create_audio_generate = patched_create_audio_generate
+
+    app = FastAPI()
+    app.add_api_route(
+        "/v1/audio/generate",
+        server.create_audio_generate,
+        methods=["POST"],
+        response_model=None,
+    )
+    return app
+
+
+@pytest.fixture
+def client(test_app):
+    return TestClient(test_app)
+
+
+# Request Validation (Pydantic model)
+class TestRequestValidation:
+    """Validate OpenAICreateAudioGenerateRequest pydantic constraints."""
+
+    def test_valid_minimal_request(self):
+        req = OpenAICreateAudioGenerateRequest(input="A calm piano melody")
+        assert req.input == "A calm piano melody"
+        assert req.response_format == "wav"
+        assert req.speed == 1.0
+
+    def test_fields_are_wired_correctly(self):
+        req = OpenAICreateAudioGenerateRequest(
+            input="rain sounds",
+            model="stable-audio",
+            response_format="flac",
+            speed=1.5,
+            audio_length=10.0,
+            audio_start=2.0,
+            negative_prompt="noise",
+            guidance_scale=7.5,
+            num_inference_steps=100,
+            seed=42,
+        )
+        assert req.input == "rain sounds"
+        assert req.model == "stable-audio"
+        assert req.response_format == "flac"
+        assert req.speed == 1.5
+        assert req.audio_length == 10.0
+        assert req.audio_start == 2.0
+        assert req.negative_prompt == "noise"
+        assert req.guidance_scale == 7.5
+        assert req.num_inference_steps == 100
+        assert req.seed == 42
+
+    def test_invalid_response_format(self):
+        with pytest.raises(Exception):
+            OpenAICreateAudioGenerateRequest(input="test", response_format="invalid_format")
+
+    def test_speed_lower_bound(self):
+        with pytest.raises(Exception):
+            OpenAICreateAudioGenerateRequest(input="test", speed=0.1)
+
+    def test_speed_upper_bound(self):
+        with pytest.raises(Exception):
+            OpenAICreateAudioGenerateRequest(input="test", speed=5.0)
+
+    def test_speed_at_boundaries(self):
+        req_low = OpenAICreateAudioGenerateRequest(input="test", speed=0.25)
+        assert req_low.speed == 0.25
+        req_high = OpenAICreateAudioGenerateRequest(input="test", speed=4.0)
+        assert req_high.speed == 4.0
+
+    def test_stream_format_sse_rejected(self):
+        with pytest.raises(Exception):
+            OpenAICreateAudioGenerateRequest(input="test", stream_format="sse")
+
+    def test_stream_format_audio_accepted(self):
+        req = OpenAICreateAudioGenerateRequest(input="test", stream_format="audio")
+        assert req.stream_format == "audio"
+
+
+# Constructor & Class Methods
+class TestConstructor:
+    def test_default_init(self):
+        server = _make_server()
+        assert server.diffusion_mode is False
+
+    def test_for_diffusion_factory(self):
+        engine_client = _make_engine_client()
+        mock_models = MagicMock()
+        mock_models.is_base_model.return_value = True
+
+        server = OmniOpenAIServingAudioGenerate.for_diffusion(
+            engine_client=engine_client,
+            models=mock_models,
+            request_logger=MagicMock(),
+        )
+        assert server.diffusion_mode is True
+
+    def test_is_stable_audio_model_true(self):
+        server = _make_server()
+        assert server._is_stable_audio_model() is True
+
+    def test_is_stable_audio_model_false(self):
+        engine = _make_engine_client()
+        engine.model_type = "SomeOtherModel"
+        server = _make_server(engine_client=engine)
+        assert server._is_stable_audio_model() is False
+
+
+# Parameter Wiring — verify request params reach the engine
+class TestParameterWiring:
+    """Ensure request parameters are correctly forwarded to the engine."""
+
+    @pytest.fixture
+    def server_and_engine(self):
+        engine = _make_engine_client()
+        server = _make_server(engine_client=engine)
+        return server, engine
+
+    @pytest.mark.asyncio
+    async def test_prompt_wiring(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="birds chirping")
+        await server.create_audio_generate(req)
+
+        engine.generate.assert_called_once()
+        call_kwargs = engine.generate.call_args[1]
+        assert call_kwargs["prompt"]["prompt"] == "birds chirping"
+        assert call_kwargs["output_modalities"] == ["audio"]
+
+    @pytest.mark.asyncio
+    async def test_negative_prompt_wiring(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="a calm ocean", negative_prompt="noise distortion")
+        await server.create_audio_generate(req)
+
+        call_kwargs = engine.generate.call_args[1]
+        assert call_kwargs["prompt"]["negative_prompt"] == "noise distortion"
+
+    @pytest.mark.asyncio
+    async def test_negative_prompt_absent(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="a calm ocean")
+        await server.create_audio_generate(req)
+
+        call_kwargs = engine.generate.call_args[1]
+        assert "negative_prompt" not in call_kwargs["prompt"]
+
+    @pytest.mark.asyncio
+    async def test_guidance_scale_wiring(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test", guidance_scale=12.0)
+        await server.create_audio_generate(req)
+
+        call_kwargs = engine.generate.call_args[1]
+        sp = call_kwargs["sampling_params_list"][0]
+        assert isinstance(sp, OmniDiffusionSamplingParams)
+        assert sp.guidance_scale == 12.0
+
+    @pytest.mark.asyncio
+    async def test_num_inference_steps_wiring(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test", num_inference_steps=200)
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        assert sp.num_inference_steps == 200
+
+    @pytest.mark.asyncio
+    async def test_seed_creates_generator(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test", seed=42)
+
+        with patch("vllm_omni.entrypoints.openai.serving_audio_generate.torch") as mock_torch:
+            mock_gen = MagicMock()
+            mock_gen.manual_seed.return_value = mock_gen
+            mock_torch.Generator.return_value = mock_gen
+
+            await server.create_audio_generate(req)
+
+            mock_torch.Generator.assert_called_once()
+            mock_gen.manual_seed.assert_called_once_with(42)
+
+    @pytest.mark.asyncio
+    async def test_seed_none_skips_generator(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test")
+
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        assert sp.generator is None
+
+    @pytest.mark.asyncio
+    async def test_audio_length_wiring(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test", audio_length=10.0, audio_start=2.0)
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        assert sp.extra_args["audio_start_in_s"] == 2.0
+        assert sp.extra_args["audio_end_in_s"] == 12.0  # start + length
+
+    @pytest.mark.asyncio
+    async def test_audio_length_default_start(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test", audio_length=5.0)
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        assert sp.extra_args["audio_start_in_s"] == 0.0
+        assert sp.extra_args["audio_end_in_s"] == 5.0
+
+    @pytest.mark.asyncio
+    async def test_no_audio_length_skips_extra_args(self, server_and_engine):
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        assert sp.extra_args == {}
+
+    @pytest.mark.asyncio
+    async def test_defaults_not_set_when_omitted(self, server_and_engine):
+        """Guidance scale and num_inference_steps keep dataclass defaults when not in request."""
+        server, engine = server_and_engine
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        await server.create_audio_generate(req)
+
+        sp = engine.generate.call_args[1]["sampling_params_list"][0]
+        defaults = OmniDiffusionSamplingParams()
+        assert sp.guidance_scale == defaults.guidance_scale
+        assert sp.num_inference_steps == defaults.num_inference_steps
+
+
+# Audio Response Format
+class TestAudioResponseFormat:
+    def test_wav_response(self, client):
+        payload = {"input": "a gentle rain", "response_format": "wav"}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 200
+        assert response.headers["content-type"] == "audio/wav"
+        assert len(response.content) > 0
+
+    def test_mp3_response(self, client):
+        payload = {"input": "a gentle rain", "response_format": "mp3"}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 200
+        assert response.headers["content-type"] == "audio/mpeg"
+        assert len(response.content) > 0
+
+    def test_flac_response(self, client):
+        payload = {"input": "a gentle rain", "response_format": "flac"}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 200
+        assert response.headers["content-type"] == "audio/flac"
+        assert len(response.content) > 0
+
+    def test_invalid_format_rejected(self, client):
+        payload = {"input": "test", "response_format": "banana"}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 422
+
+    @patch("vllm_omni.entrypoints.openai.serving_audio_generate.OmniOpenAIServingAudioGenerate.create_audio")
+    def test_speed_parameter_forwarded(self, mock_create_audio, test_app):
+        mock_audio_response = MagicMock()
+        mock_audio_response.audio_data = b"dummy_audio"
+        mock_audio_response.media_type = "audio/wav"
+        mock_create_audio.return_value = mock_audio_response
+
+        c = TestClient(test_app)
+        payload = {"input": "test", "response_format": "wav", "speed": 2.5}
+        c.post("/v1/audio/generate", json=payload)
+
+        mock_create_audio.assert_called_once()
+        audio_obj = mock_create_audio.call_args[0][0]
+        assert isinstance(audio_obj, CreateAudio)
+        assert audio_obj.speed == 2.5
+
+    @patch("vllm_omni.entrypoints.openai.serving_audio_generate.OmniOpenAIServingAudioGenerate.create_audio")
+    def test_sample_rate_from_output(self, mock_create_audio, test_app):
+        mock_audio_response = MagicMock()
+        mock_audio_response.audio_data = b"dummy"
+        mock_audio_response.media_type = "audio/wav"
+        mock_create_audio.return_value = mock_audio_response
+
+        c = TestClient(test_app)
+        payload = {"input": "test"}
+        c.post("/v1/audio/generate", json=payload)
+
+        audio_obj = mock_create_audio.call_args[0][0]
+        assert audio_obj.sample_rate == 44100  # Stable Audio default
+
+
+# Error Handling
+class TestErrorHandling:
+    @pytest.mark.asyncio
+    async def test_no_output_returns_error(self):
+        engine = _make_engine_client()
+
+        async def empty_gen(*args, **kwargs):
+            return
+            yield  # noqa: unreachable – makes this an async generator
+
+        engine.generate = MagicMock(side_effect=empty_gen)
+        server = _make_server(engine_client=engine)
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        resp = await server.create_audio_generate(req)
+
+        # create_error_response returns an ErrorResponse with .error.message
+        assert "No output generated" in resp.error.message
+
+    @pytest.mark.asyncio
+    async def test_no_audio_in_output_returns_error(self):
+        engine = _make_engine_client()
+
+        async def gen_without_audio(*args, **kwargs):
+            yield OmniRequestOutput.from_diffusion(
+                request_id="test",
+                images=[],
+                prompt=None,
+                metrics={},
+                multimodal_output={},  # no audio key
+            )
+
+        engine.generate = MagicMock(side_effect=gen_without_audio)
+        server = _make_server(engine_client=engine)
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        resp = await server.create_audio_generate(req)
+
+        assert "did not produce audio" in resp.error.message
+
+    @pytest.mark.asyncio
+    async def test_engine_errored_raises(self):
+        engine = _make_engine_client()
+        engine.errored = True
+        engine.dead_error = RuntimeError("engine is dead")
+        server = _make_server(engine_client=engine)
+
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        with pytest.raises(RuntimeError, match="engine is dead"):
+            await server.create_audio_generate(req)
+
+    @pytest.mark.asyncio
+    async def test_model_outputs_key_fallback(self):
+        """Audio data under 'model_outputs' key should be accepted."""
+        engine = _make_engine_client(audio_key="model_outputs")
+        server = _make_server(engine_client=engine)
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        resp = await server.create_audio_generate(req)
+
+        # Should succeed and return a Response with audio bytes
+        assert hasattr(resp, "body")
+        assert len(resp.body) > 0
+
+    @pytest.mark.asyncio
+    async def test_value_error_returns_error_response(self):
+        engine = _make_engine_client()
+
+        async def gen_value_error(*args, **kwargs):
+            raise ValueError("bad value")
+            yield  # noqa: unreachable
+
+        engine.generate = MagicMock(side_effect=gen_value_error)
+        server = _make_server(engine_client=engine)
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        resp = await server.create_audio_generate(req)
+
+        assert "bad value" in resp.error.message
+
+    @pytest.mark.asyncio
+    async def test_generic_exception_returns_error_response(self):
+        engine = _make_engine_client()
+
+        async def gen_runtime_error(*args, **kwargs):
+            raise RuntimeError("something went wrong")
+            yield  # noqa: unreachable
+
+        engine.generate = MagicMock(side_effect=gen_runtime_error)
+        server = _make_server(engine_client=engine)
+        req = OpenAICreateAudioGenerateRequest(input="test")
+        resp = await server.create_audio_generate(req)
+
+        assert "Audio generation failed" in resp.error.message
+
+
+# End-to-End via TestClient
+class TestAudioGenerateAPI:
+    def test_basic_success(self, client):
+        payload = {"input": "ambient forest sounds"}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 200
+        assert len(response.content) > 0
+
+    def test_with_all_params(self, client):
+        payload = {
+            "input": "gentle piano",
+            "response_format": "wav",
+            "speed": 1.0,
+            "audio_length": 5.0,
+            "audio_start": 0.0,
+            "negative_prompt": "noise",
+            "guidance_scale": 7.0,
+            "num_inference_steps": 50,
+            "seed": 123,
+        }
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 200
+        assert response.headers["content-type"] == "audio/wav"
+
+    def test_missing_input_rejected(self, client):
+        payload = {}
+        response = client.post("/v1/audio/generate", json=payload)
+        assert response.status_code == 422
+
+    def test_extra_unknown_fields_ignored(self, client):
+        payload = {"input": "test", "unknown_field": "value"}
+        response = client.post("/v1/audio/generate", json=payload)
+        # Pydantic v2 ignores extra fields by default
+        assert response.status_code == 200

tests/entrypoints/openai_api/test_serving_speech_voxcpm2.py

@@ -1,330 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""UTs for VoxCPM2 OpenAI speech serving behavior (`instructions` + `cfg_value`)."""
-
-import asyncio
-from types import SimpleNamespace
-from unittest.mock import AsyncMock
-
-import pytest
-from pytest_mock import MockerFixture
-
-from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest
-from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.omni]
-
-
-@pytest.fixture
-def voxcpm2_server(mocker: MockerFixture):
-    """Mock an OmniOpenAIServingSpeech instance detected as `voxcpm2`.
-
-    Single `latent_generator` stage with `model_arch=VoxCPM2...` so
-    `_detect_tts_model_type` returns ``voxcpm2``.
-    """
-    mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set())
-    mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None)
-
-    mock_engine_client = mocker.MagicMock()
-    mock_engine_client.errored = False
-    mock_engine_client.model_config = mocker.MagicMock(
-        model="openbmb/VoxCPM2",
-        hf_config=SimpleNamespace(
-            audio_vae_config={"sample_rate": 16000, "encoder_rates": [2, 5, 8, 8]},
-            patch_size=4,
-        ),
-    )
-    mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=4096)]
-    mock_engine_client.tts_batch_max_items = 32
-    # Explicit None so `_compute_max_instructions_length` skips the CLI-override
-    # branch (MagicMock attributes are truthy by default) and falls through to
-    # the `_TTS_MAX_INSTRUCTIONS_LENGTH = 500` default.
-    mock_engine_client.tts_max_instructions_length = None
-    mock_engine_client.generate = mocker.MagicMock(return_value="generator")
-    mock_engine_client.stage_configs = [
-        SimpleNamespace(
-            engine_args=SimpleNamespace(
-                model_stage="latent_generator",
-                model_arch="VoxCPM2TalkerForConditionalGeneration",
-            ),
-            tts_args={},
-        ),
-    ]
-
-    mock_models = mocker.MagicMock()
-    mock_models.is_base_model.return_value = True
-
-    server = OmniOpenAIServingSpeech(
-        engine_client=mock_engine_client,
-        models=mock_models,
-        request_logger=mocker.MagicMock(),
-    )
-    # Lazy tokenizer path is not exercised by these tests; provide a no-op encoder
-    # and sentinel tokenizer / split_map so `_build_voxcpm2_prompt` can run.
-    server._voxcpm2_tokenizer = mocker.MagicMock(name="voxcpm2_tokenizer")
-    server._voxcpm2_split_map = {}
-    mocker.patch.object(server, "_voxcpm2_encode", return_value=[])
-    return server
-
-
-@pytest.fixture
-def mock_build_voxcpm2_prompt(mocker: MockerFixture):
-    """Patch `build_voxcpm2_prompt` so we can inspect its call args without
-    touching the real VoxCPM2 tokenizer / audio VAE.
-
-    Returns the prompt dict shape the real helper would produce; the test
-    can then also assert `cfg_value` stashing into `additional_information`.
-    """
-    return mocker.patch(
-        "vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker.build_voxcpm2_prompt",
-        return_value={"prompt_token_ids": [1], "additional_information": {"text_token_ids": [[]]}},
-    )
-
-
-class TestVoxCPM2Serving:
-    def test_voxcpm2_model_type_detection(self, voxcpm2_server):
-        assert voxcpm2_server._tts_model_type == "voxcpm2"
-        assert voxcpm2_server._is_tts is True
-
-    def test_voxcpm2_accepts_any_text_input(self, voxcpm2_server):
-        """VoxCPM2 skips the strict voxcpm validator (see `_validate_tts_request`)."""
-        request = OpenAICreateSpeechRequest(input="مرحباً", instructions="warm tone")
-        assert voxcpm2_server._validate_tts_request(request) is None
-
-    def test_build_prompt_text_only(self, voxcpm2_server, mock_build_voxcpm2_prompt):
-        """No instructions, no cfg_value: text flows through unchanged, no cfg stash."""
-        request = OpenAICreateSpeechRequest(input="hello voxcpm2")
-        prompt = asyncio.run(voxcpm2_server._build_voxcpm2_prompt(request))
-
-        assert mock_build_voxcpm2_prompt.call_args.kwargs["text"] == "hello voxcpm2"
-        assert "cfg_value" not in prompt.get("additional_information", {})
-
-    def test_build_prompt_prepends_instructions(self, voxcpm2_server, mock_build_voxcpm2_prompt):
-        """`instructions` wraps in parens and prepends to text (native VoxCPM2 convention)."""
-        request = OpenAICreateSpeechRequest(
-            input="hello voxcpm2",
-            instructions="A warm young woman",
-        )
-        asyncio.run(voxcpm2_server._build_voxcpm2_prompt(request))
-
-        assert mock_build_voxcpm2_prompt.call_args.kwargs["text"] == "(A warm young woman)hello voxcpm2"
-
-    def test_build_prompt_strips_instructions_whitespace(self, voxcpm2_server, mock_build_voxcpm2_prompt):
-        """Leading/trailing whitespace in `instructions` is stripped before prepending."""
-        request = OpenAICreateSpeechRequest(input="hello", instructions="  calm radio  ")
-        asyncio.run(voxcpm2_server._build_voxcpm2_prompt(request))
-
-        assert mock_build_voxcpm2_prompt.call_args.kwargs["text"] == "(calm radio)hello"
-
-    def test_build_prompt_stashes_cfg_value(self, voxcpm2_server, mock_build_voxcpm2_prompt):
-        """`extra_params['cfg_value']` lands in `additional_information` for the talker to lift."""
-        request = OpenAICreateSpeechRequest(input="hello", extra_params={"cfg_value": 2.7})
-        prompt = asyncio.run(voxcpm2_server._build_voxcpm2_prompt(request))
-
-        assert prompt["additional_information"]["cfg_value"] == 2.7
-
-    def test_build_prompt_omits_cfg_value_when_extra_params_missing(self, voxcpm2_server, mock_build_voxcpm2_prompt):
-        """Omitting `cfg_value` must not add the key (talker falls back to its default)."""
-        request = OpenAICreateSpeechRequest(input="hello")
-        prompt = asyncio.run(voxcpm2_server._build_voxcpm2_prompt(request))
-
-        assert "cfg_value" not in prompt["additional_information"]
-
-    def test_build_prompt_omits_cfg_value_when_extra_params_has_other_keys(
-        self, voxcpm2_server, mock_build_voxcpm2_prompt
-    ):
-        """`extra_params` set but without `cfg_value` must not affect the prompt."""
-        request = OpenAICreateSpeechRequest(input="hello", extra_params={"some_other_knob": 1})
-        prompt = asyncio.run(voxcpm2_server._build_voxcpm2_prompt(request))
-
-        assert "cfg_value" not in prompt["additional_information"]
-
-    def test_build_prompt_instructions_and_cfg_together(self, voxcpm2_server, mock_build_voxcpm2_prompt):
-        """Both features compose: text prefixed AND cfg_value stashed."""
-        request = OpenAICreateSpeechRequest(
-            input="hello",
-            instructions="excited",
-            extra_params={"cfg_value": 2.5},
-        )
-        prompt = asyncio.run(voxcpm2_server._build_voxcpm2_prompt(request))
-
-        assert mock_build_voxcpm2_prompt.call_args.kwargs["text"] == "(excited)hello"
-        assert prompt["additional_information"]["cfg_value"] == 2.5
-
-    def test_build_prompt_hifi_cloning_ref_audio_ref_text_cfg(
-        self, voxcpm2_server, mock_build_voxcpm2_prompt, mocker: MockerFixture
-    ):
-        """Hi-Fi Cloning (`ref_audio` + `ref_text`) resolves audio and threads cfg_value."""
-        voxcpm2_server._resolve_ref_audio = AsyncMock(return_value=([0.1, -0.1, 0.2], 16000))
-
-        request = OpenAICreateSpeechRequest(
-            input="clone me",
-            ref_audio="data:audio/wav;base64,QUJD",
-            ref_text="reference transcript",
-            extra_params={"cfg_value": 2.7},
-        )
-        prompt = asyncio.run(voxcpm2_server._build_voxcpm2_prompt(request))
-
-        kwargs = mock_build_voxcpm2_prompt.call_args.kwargs
-        assert kwargs["text"] == "clone me"
-        assert kwargs["ref_audio"] == [0.1, -0.1, 0.2]
-        assert kwargs["ref_sr"] == 16000
-        assert kwargs["ref_text"] == "reference transcript"
-        assert prompt["additional_information"]["cfg_value"] == 2.7
-        voxcpm2_server._resolve_ref_audio.assert_awaited_once_with("data:audio/wav;base64,QUJD")
-
-    def test_validate_rejects_instructions_with_ref_text(self, voxcpm2_server):
-        """Hi-Fi Cloning (ref_audio + ref_text) + instructions returns a 400-level
-        error string from the validator instead of silently dropping `instructions`.
-
-        Matches upstream voxcpm CLI's hard error on the same combo
-        (src/voxcpm/cli.py:128-131: --control + --prompt-text).
-        """
-        request = OpenAICreateSpeechRequest(
-            input="hello",
-            ref_audio="data:audio/wav;base64,QUJD",
-            ref_text="reference transcript",
-            instructions="this is incompatible with Hi-Fi mode",
-        )
-        error = voxcpm2_server._validate_tts_request(request)
-        assert error is not None
-        assert "Hi-Fi" in error
-        assert "ref_text" in error
-
-    def test_validate_accepts_instructions_without_ref_text(self, voxcpm2_server):
-        """`instructions` is fine in Voice Design and Controllable Cloning modes
-        (i.e. whenever `ref_text` is not also set)."""
-        request = OpenAICreateSpeechRequest(input="hello", instructions="warm tone")
-        assert voxcpm2_server._validate_tts_request(request) is None
-
-    def test_validate_accepts_ref_text_without_instructions(self, voxcpm2_server):
-        """Hi-Fi Cloning (ref_audio + ref_text, no instructions) is fine."""
-        request = OpenAICreateSpeechRequest(
-            input="hello",
-            ref_audio="data:audio/wav;base64,QUJD",
-            ref_text="reference transcript",
-        )
-        assert voxcpm2_server._validate_tts_request(request) is None
-
-    def test_validate_rejects_overlong_instructions(self, voxcpm2_server):
-        """Instructions longer than `_max_instructions_length` (500 default) are rejected
-        with an error message that includes the actual length, the cap, and an
-        upstream-context hint so users coming from the voxcpm CLI understand why
-        the limit exists.
-        """
-        oversize = "x" * (voxcpm2_server._max_instructions_length + 1)
-        request = OpenAICreateSpeechRequest(input="hello", instructions=oversize)
-        error = voxcpm2_server._validate_tts_request(request)
-        assert error is not None
-        assert str(len(oversize)) in error
-        assert str(voxcpm2_server._max_instructions_length) in error
-        assert "Upstream voxcpm has no cap" in error
-
-    def test_validate_accepts_at_limit_instructions(self, voxcpm2_server):
-        """Instructions exactly at `_max_instructions_length` are accepted."""
-        at_limit = "x" * voxcpm2_server._max_instructions_length
-        request = OpenAICreateSpeechRequest(input="hello", instructions=at_limit)
-        assert voxcpm2_server._validate_tts_request(request) is None
-
-    def test_prepare_speech_generation_runs_validator_for_voxcpm2_length(
-        self, voxcpm2_server, mock_build_voxcpm2_prompt, mocker: MockerFixture
-    ):
-        """Single-request `/v1/audio/speech` path must invoke `_validate_tts_request`
-        before building the prompt, so the instructions-length cap actually fires
-        for normal (non-batch) requests too. Regression guard for #3118 P2.
-        """
-        oversize = "x" * (voxcpm2_server._max_instructions_length + 1)
-        request = OpenAICreateSpeechRequest(input="hello", instructions=oversize)
-        with pytest.raises(ValueError, match=str(voxcpm2_server._max_instructions_length)):
-            asyncio.run(voxcpm2_server._prepare_speech_generation(request))
-        # The prompt builder must NOT be reached when validation fails.
-        mock_build_voxcpm2_prompt.assert_not_called()
-
-    def test_prepare_speech_generation_runs_validator_for_voxcpm2_hifi(
-        self, voxcpm2_server, mock_build_voxcpm2_prompt, mocker: MockerFixture
-    ):
-        """Single-request path must also surface the Hi-Fi mode rejection
-        (instructions + ref_text), not silently strip `instructions`.
-        """
-        request = OpenAICreateSpeechRequest(
-            input="hello",
-            ref_audio="data:audio/wav;base64,QUJD",
-            ref_text="reference transcript",
-            instructions="incompatible with Hi-Fi mode",
-        )
-        with pytest.raises(ValueError, match="Hi-Fi"):
-            asyncio.run(voxcpm2_server._prepare_speech_generation(request))
-        mock_build_voxcpm2_prompt.assert_not_called()
-
-    @pytest.mark.parametrize("cfg", [0.1, 0.5, 1.5, 2.0, 2.7, 5.0, 10.0])
-    def test_cfg_value_accepts_range(self, voxcpm2_server, mock_build_voxcpm2_prompt, cfg):
-        """`_build_voxcpm2_prompt` accepts cfg_value within 0.1-10.0 inclusive."""
-        request = OpenAICreateSpeechRequest(input="hello", extra_params={"cfg_value": cfg})
-        prompt = asyncio.run(voxcpm2_server._build_voxcpm2_prompt(request))
-        assert prompt["additional_information"]["cfg_value"] == cfg
-
-    @pytest.mark.parametrize("cfg", [0.0, -1.0, 10.5, 100.0])
-    def test_cfg_value_rejects_out_of_range(self, voxcpm2_server, mock_build_voxcpm2_prompt, cfg):
-        """Out-of-range `cfg_value` (validated in `_build_voxcpm2_prompt`) raises ValueError."""
-        request = OpenAICreateSpeechRequest(input="hello", extra_params={"cfg_value": cfg})
-        with pytest.raises(ValueError, match="out of range"):
-            asyncio.run(voxcpm2_server._build_voxcpm2_prompt(request))
-
-    @pytest.mark.parametrize("bad", ["abc", None, [1, 2], {"x": 1}])
-    def test_cfg_value_rejects_non_numeric(self, voxcpm2_server, mock_build_voxcpm2_prompt, bad):
-        """Non-numeric `extra_params['cfg_value']` raises ValueError instead of crashing the talker."""
-        request = OpenAICreateSpeechRequest(input="hello", extra_params={"cfg_value": bad})
-        with pytest.raises(ValueError, match="must be a number|out of range"):
-            asyncio.run(voxcpm2_server._build_voxcpm2_prompt(request))
-
-
-class TestVoxCPM2DecodeStepCfgPropagation:
-    """Static AST regression guard for #3118 P1: ensure both `_finish_prefill`
-    and `_finish_decode` invoke `_run_cfm` with the per-request `cfg_value`
-    kwarg. Originally only `_finish_prefill` threaded the kwarg, so requests
-    longer than the prefill patch fell back to `self._cfg_value = 2.0` for
-    every decode step. This test parses the talker source and asserts the
-    kwarg is present at every `_run_cfm` call site inside the two methods,
-    so a future refactor that drops the kwarg from either method fails CI
-    even before the behavioral cfg-sweep tests run.
-    """
-
-    @pytest.fixture(scope="class")
-    def talker_class_ast(self):
-        import ast
-        import inspect
-
-        from vllm_omni.model_executor.models.voxcpm2 import voxcpm2_talker
-
-        src = inspect.getsource(voxcpm2_talker)
-        tree = ast.parse(src)
-        cls = next(
-            n
-            for n in ast.walk(tree)
-            if isinstance(n, ast.ClassDef) and n.name == "VoxCPM2TalkerForConditionalGeneration"
-        )
-        return cls
-
-    def _run_cfm_calls_in(self, cls_ast, method_name):
-        import ast
-
-        method = next(n for n in cls_ast.body if isinstance(n, ast.FunctionDef) and n.name == method_name)
-        return [
-            n
-            for n in ast.walk(method)
-            if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) and n.func.attr == "_run_cfm"
-        ]
-
-    @pytest.mark.parametrize("method_name", ["_finish_prefill", "_finish_decode"])
-    def test_run_cfm_called_with_cfg_value_kwarg(self, talker_class_ast, method_name):
-        calls = self._run_cfm_calls_in(talker_class_ast, method_name)
-        assert calls, f"{method_name} must call self._run_cfm at least once"
-        for call in calls:
-            kwarg_names = {kw.arg for kw in call.keywords if kw.arg}
-            assert "cfg_value" in kwarg_names, (
-                f"{method_name}: every self._run_cfm call must thread the "
-                f"per-request `cfg_value` kwarg (the bug fixed in #3118 P1 was "
-                f"that `_finish_decode` dropped this kwarg, so all decode steps "
-                f"silently fell back to the talker default `self._cfg_value`)."
-            )

tests/entrypoints/openai_api/test_video_stream_handler.py

@@ -278,12 +278,14 @@ class TestHandleQueryErrors:
     @pytest.mark.asyncio
     async def test_query_returns_error_response(self):
         """create_chat_completion returns ErrorResponse → error forwarded."""
-        from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+        from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
 
         err = ErrorResponse(
-            message="model overloaded",
-            type="server_error",
-            code=503,
+            error=ErrorInfo(
+                message="model overloaded",
+                type="server_error",
+                code=503,
+            ),
         )
 
         chat = AsyncMock()

tests/entrypoints/test_omni_entrypoints.py

@@ -426,8 +426,9 @@ def test_openai_serving_models_can_consume_async_omni_compat_attrs():
 
     assert serving_models.model_config is model_config
     assert serving_models.renderer is renderer
-    assert serving_models.io_processor is io_processor
     assert serving_models.input_processor is input_processor
+    # vLLM 0.20 keeps io_processor on the engine client instead of copying it.
+    assert serving_models.engine_client.io_processor is io_processor
 
 
 def test_get_diffusion_od_config_returns_diffusion_stage_config():

tests/entrypoints/test_serve.py

@@ -108,7 +108,7 @@ def test_run_headless_registers_stage_once_and_launches_all_local_engines(mocker
     assert manager_kwargs["local_client"] is False
     assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001"
     assert manager_kwargs["log_stats"] is False
-    engine_manager.join_first.assert_called_once_with()
+    engine_manager.monitor_engine_liveness.assert_called_once_with()
     engine_manager.shutdown.assert_called_once_with()
 
 

tests/examples/helpers.py

@@ -351,3 +351,33 @@ def strip_trailing_audio_saved_line(text: str) -> str:
     while lines and lines[-1].strip().startswith("Audio saved to"):
         lines.pop()
     return "\n".join(lines).strip()
+
+
+def strip_audio_saved_to_lines(text: str) -> str:
+    """Remove every line starting with ``Audio saved to`` (streaming prints one per chunk).
+
+    Without this, :func:`extract_content_after_keyword` with ``content:`` and ``DOTALL``
+    keeps those lines inside the captured text segment.
+    """
+    lines = [ln for ln in text.splitlines() if not ln.strip().startswith("Audio saved to")]
+    return "\n".join(lines).strip()
+
+
+def extract_last_audio_saved_path(text: str) -> str:
+    """Return the filesystem path from the last ``Audio saved to`` line.
+
+    Non-streaming output has a single line; streaming prints one path per chunk.
+    Do not use :func:`extract_content_after_keyword` with ``Audio saved to`` for
+    streaming: greedy ``.+`` under ``DOTALL`` concatenates every path and body into one
+    invalid string (Linux ``File name too long`` when opening it as a path).
+    """
+    last_path: str | None = None
+    for line in text.splitlines():
+        stripped = line.strip()
+        if stripped.startswith("Audio saved to"):
+            rest = stripped[len("Audio saved to") :].strip()
+            if rest:
+                last_path = rest
+    if last_path is None:
+        raise AssertionError("'Audio saved to' line with a path not found in command output")
+    return last_path

tests/examples/offline_inference/test_qwen3_tts_estimator.py

@@ -0,0 +1,121 @@
+"""
+Offline inference regression tests: Qwen3-TTS _estimate_prompt_len.
+
+Covers the four ref_audio input forms handled by the estimator in
+examples/offline_inference/qwen3_tts/end2end.py:
+  - bare local path
+  - file:// URI
+  - http:// URL
+  - data: URI
+"""
+
+import base64
+import io
+import sys
+import wave
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+_examples_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "offline_inference" / "qwen3_tts")
+sys.path.insert(0, _examples_dir)
+from end2end import _estimate_prompt_len  # noqa: E402
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+TEST_SR = 16000
+TEST_DURATION_S = 3
+CODEC_FRAME_RATE = 12
+EXPECTED_REF_CODE_LEN = TEST_DURATION_S * CODEC_FRAME_RATE  # 36, distinct from fallback 2048
+
+
+def _wav_bytes() -> bytes:
+    buf = io.BytesIO()
+    samples = np.zeros(TEST_SR * TEST_DURATION_S, dtype=np.int16)
+    with wave.open(buf, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(TEST_SR)
+        wf.writeframes(samples.tobytes())
+    return buf.getvalue()
+
+
+@pytest.fixture
+def wav_file(tmp_path: Path) -> Path:
+    path = tmp_path / "ref.wav"
+    path.write_bytes(_wav_bytes())
+    return path
+
+
+@pytest.fixture
+def estimator_cache(monkeypatch):
+    """
+    Bypass model loading by pre-populating the estimator cache, and shortcut
+    Qwen3TTSTalker.estimate_prompt_len_from_additional_information so that
+    the test directly observes the real _estimate_ref_code_len closure.
+    """
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+        Qwen3TTSTalkerForConditionalGeneration,
+    )
+
+    def _shortcut(**kwargs):
+        ref_audio = kwargs["additional_information"]["ref_audio"][0]
+        return kwargs["estimate_ref_code_len"](ref_audio)
+
+    monkeypatch.setattr(
+        Qwen3TTSTalkerForConditionalGeneration,
+        "estimate_prompt_len_from_additional_information",
+        staticmethod(_shortcut),
+    )
+
+    mock_tcfg = SimpleNamespace(codec_frame_rate=CODEC_FRAME_RATE, codec_language_id=None, spk_is_dialect=None)
+    return {"test_model": (MagicMock(), mock_tcfg, None)}
+
+
+def _run(cache: dict, ref_audio: str) -> int | None:
+    return _estimate_prompt_len(
+        additional_information={"task_type": ["Base"], "ref_audio": [ref_audio]},
+        model_name="test_model",
+        _cache=cache,
+    )
+
+
+def test_bare_local_path(estimator_cache, monkeypatch, wav_file):
+    """Bare local path must be loaded via load_audio, not MediaConnector."""
+    from vllm.multimodal.media import MediaConnector
+
+    def _must_not_be_called(self, path):
+        pytest.fail(f"MediaConnector.fetch_audio should not be called for bare local path: {path}")
+
+    monkeypatch.setattr(MediaConnector, "fetch_audio", _must_not_be_called)
+
+    assert _run(estimator_cache, str(wav_file)) == EXPECTED_REF_CODE_LEN
+
+
+def test_file_uri(estimator_cache, wav_file):
+    """file:// URI is routed through MediaConnector, which handles it natively."""
+    assert _run(estimator_cache, f"file://{wav_file}") == EXPECTED_REF_CODE_LEN
+
+
+def test_http_url(estimator_cache, monkeypatch):
+    """http:// URL is routed through MediaConnector.fetch_audio (mocked to avoid network)."""
+    from vllm.multimodal.media import MediaConnector
+
+    audio = np.zeros(TEST_SR * TEST_DURATION_S, dtype=np.float32)
+
+    def _fake_fetch_audio(self, path):
+        return audio, TEST_SR
+
+    monkeypatch.setattr(MediaConnector, "fetch_audio", _fake_fetch_audio)
+
+    assert _run(estimator_cache, "http://example.com/ref.wav") == EXPECTED_REF_CODE_LEN
+
+
+def test_data_uri(estimator_cache):
+    """data: URI is routed through MediaConnector, which decodes base64 natively."""
+    data_uri = "data:audio/wav;base64," + base64.b64encode(_wav_bytes()).decode("ascii")
+    assert _run(estimator_cache, data_uri) == EXPECTED_REF_CODE_LEN

tests/examples/offline_inference/test_text_to_image.py

@@ -13,7 +13,6 @@ from tests.helpers.mark import hardware_marks
 
 pytestmark = [pytest.mark.full_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})]
 
-
 T2I_SCRIPT = EXAMPLES / "offline_inference" / "text_to_image" / "text_to_image.py"
 README_PATH = T2I_SCRIPT.with_name("README.md")
 EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_t2i"

tests/examples/online_serving/test_qwen2_5_omni.py

@@ -4,13 +4,18 @@ See examples/online_serving/qwen2_5_omni/README.md
 """
 
 import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
 from pathlib import Path
 
 import pytest
 
 from tests.examples.helpers import (
     extract_content_after_keyword,
+    extract_last_audio_saved_path,
     run_cmd,
+    strip_audio_saved_to_lines,
     strip_trailing_audio_saved_line,
 )
 from tests.helpers.mark import hardware_test
@@ -20,8 +25,6 @@ from tests.helpers.stage_config import get_deploy_config_path
 
 pytestmark = [pytest.mark.full_model, pytest.mark.example, pytest.mark.omni]
 
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
 models = ["Qwen/Qwen2.5-Omni-7B"]
 
 # Single CI deploy YAML; rocm/xpu deltas are picked automatically via the
@@ -55,8 +58,8 @@ def test_send_multimodal_request_001(omni_server) -> None:
     text_content = strip_trailing_audio_saved_line(text_content_tmp)
 
     # Verify text output same as audio output
-    wav_path = extract_content_after_keyword("Audio saved to", result)
-    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path.strip()}")
+    wav_path = extract_last_audio_saved_path(result)
+    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path}")
     print(f"text content is: {text_content}")
     print(f"audio content is: {audio_content}")
 
@@ -91,8 +94,8 @@ def test_send_multimodal_request_002(omni_server) -> None:
     text_content = strip_trailing_audio_saved_line(text_content_tmp)
 
     # Verify text output same as audio output
-    wav_path = extract_content_after_keyword("Audio saved to", result)
-    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path.strip()}")
+    wav_path = extract_last_audio_saved_path(result)
+    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path}")
     print(f"text content is: {text_content}")
     assert all(keyword in text_content for keyword in ["baby", "book"]), (
         "The output does not contain any of the keywords in video description."
@@ -176,8 +179,8 @@ def test_modality_control_002(omni_server) -> None:
 
     result = run_cmd(command)
     # Verify text output same as audio output
-    wav_path = extract_content_after_keyword("Audio saved to", result)
-    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path.strip()}")
+    wav_path = extract_last_audio_saved_path(result)
+    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path}")
     print(f"audio content is: {audio_content}")
     assert all(keyword in audio_content for keyword in ["baby", "book"]), (
         "The output does not contain any of the keywords in video description."
@@ -207,8 +210,8 @@ def test_modality_control_003(omni_server) -> None:
     text_content = strip_trailing_audio_saved_line(text_content_tmp)
 
     # Verify text output same as audio output
-    wav_path = extract_content_after_keyword("Audio saved to", result)
-    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path.strip()}")
+    wav_path = extract_last_audio_saved_path(result)
+    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path}")
     print(f"text content is: {text_content}")
     assert all(keyword in text_content for keyword in ["baby", "book"]), (
         "The output does not contain any of the keywords in video description."
@@ -239,11 +242,11 @@ def test_stream_001(omni_server) -> None:
     result = run_cmd(command)
 
     text_content_tmp = extract_content_after_keyword("content:", result)
-    text_content = strip_trailing_audio_saved_line(text_content_tmp)
+    text_content = strip_audio_saved_to_lines(text_content_tmp)
 
     # Verify text output same as audio output
-    wav_path = extract_content_after_keyword("Audio saved to", result)
-    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path.strip()}")
+    wav_path = extract_last_audio_saved_path(result)
+    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path}")
     print(f"text content is: {text_content}")
     assert all(keyword in text_content for keyword in ["baby", "book"]), (
         "The output does not contain any of the keywords in video description."

tests/examples/online_serving/test_qwen3_omni.py

@@ -4,13 +4,18 @@ See examples/online_serving/qwen3_omni/README.md
 """
 
 import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
 from pathlib import Path
 
 import pytest
 
 from tests.examples.helpers import (
     extract_content_after_keyword,
+    extract_last_audio_saved_path,
     run_cmd,
+    strip_audio_saved_to_lines,
     strip_trailing_audio_saved_line,
 )
 from tests.helpers.mark import hardware_test
@@ -20,8 +25,6 @@ from tests.helpers.stage_config import get_deploy_config_path
 
 pytestmark = [pytest.mark.full_model, pytest.mark.example, pytest.mark.omni]
 
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 
@@ -51,9 +54,9 @@ def test_send_multimodal_request_001(omni_server) -> None:
     result = run_cmd(command)
     text_content_tmp = extract_content_after_keyword("Chat completion output from text:", result)
     text_content = strip_trailing_audio_saved_line(text_content_tmp)
-    wav_path = extract_content_after_keyword("Audio saved to", result)
+    wav_path = extract_last_audio_saved_path(result)
     # Verify text output same as audio output
-    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path.strip()}")
+    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path}")
     print(f"text content is: {text_content}")
     print(f"audio content is: {audio_content}")
 
@@ -83,8 +86,8 @@ def test_send_multimodal_request_002(omni_server) -> None:
     text_content = strip_trailing_audio_saved_line(text_content_tmp)
 
     # Verify text output same as audio output
-    wav_path = extract_content_after_keyword("Audio saved to", result)
-    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path.strip()}")
+    wav_path = extract_last_audio_saved_path(result)
+    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path}")
     print(f"text content is: {text_content}")
     print(f"audio content is: {audio_content}")
     assert all(keyword in text_content for keyword in ["baby", "book"]), (
@@ -148,8 +151,8 @@ def test_modality_control_002(omni_server) -> None:
 
     result = run_cmd(command)
     # Verify text output same as audio output
-    wav_path = extract_content_after_keyword("Audio saved to", result)
-    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path.strip()}")
+    wav_path = extract_last_audio_saved_path(result)
+    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path}")
     print(f"audio content is: {audio_content}")
     assert "cherry blossom" in audio_content, "The output does not contain any of the keywords."
 
@@ -174,8 +177,8 @@ def test_modality_control_003(omni_server) -> None:
     text_content = strip_trailing_audio_saved_line(text_content_tmp)
 
     # Verify text output same as audio output
-    wav_path = extract_content_after_keyword("Audio saved to", result)
-    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path.strip()}")
+    wav_path = extract_last_audio_saved_path(result)
+    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path}")
     print(f"text content is: {text_content}")
     assert "cherry blossom" in audio_content, "The output does not contain any of the keywords."
     print(f"audio content is: {audio_content}")
@@ -200,11 +203,11 @@ def test_stream_001(omni_server) -> None:
     result = run_cmd(command)
 
     text_content_tmp = extract_content_after_keyword("content:", result)
-    text_content = strip_trailing_audio_saved_line(text_content_tmp)
+    text_content = strip_audio_saved_to_lines(text_content_tmp)
 
     # Verify text output same as audio output
-    wav_path = extract_content_after_keyword("Audio saved to", result)
-    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path.strip()}")
+    wav_path = extract_last_audio_saved_path(result)
+    audio_content = convert_audio_file_to_text(output_path=f"./{wav_path}")
     print(f"text content is: {text_content}")
     assert "cherry blossom" in audio_content, "The output does not contain any of the keywords."
     print(f"audio content is: {audio_content}")

tests/helpers/mark.py

@@ -3,7 +3,14 @@
 """Pytest marks and decorators for hardware / resource selection (CUDA, ROCm, …)."""
 
 import pytest
-from vllm.utils.torch_utils import cuda_device_count_stateless
+
+try:
+    from vllm.utils.torch_utils import cuda_device_count_stateless
+except ImportError:
+    import torch
+
+    def cuda_device_count_stateless() -> int:
+        return torch.cuda.device_count()
 
 # Re-exported from tests.helpers.env (GPU wait + DeviceMemoryMonitor).
 

tests/helpers/media.py

@@ -496,6 +496,33 @@ def generate_synthetic_image(
     }
 
 
+_TEST_ASSETS_ROOT = Path(__file__).resolve().parents[1] / "assets"
+
+_AUDIO_MIME_BY_SUFFIX = {
+    ".wav": "audio/wav",
+    ".mp3": "audio/mpeg",
+    ".flac": "audio/flac",
+    ".ogg": "audio/ogg",
+}
+
+
+def test_asset_path(relative_path: str | os.PathLike) -> Path:
+    """Resolve a path under ``tests/assets/`` to its absolute on-disk location."""
+    return _TEST_ASSETS_ROOT / Path(relative_path)
+
+
+def load_test_audio_data_url(relative_path: str | os.PathLike) -> str:
+    """Load a vendored test audio file under ``tests/assets/`` as a base64 data URL.
+
+    Used by tests that need real reference audio (e.g. voice cloning) without
+    relying on the server's ability to fetch external URLs at request time.
+    """
+    path = test_asset_path(relative_path)
+    mime = _AUDIO_MIME_BY_SUFFIX.get(path.suffix.lower(), "application/octet-stream")
+    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
+    return f"data:{mime};base64,{encoded}"
+
+
 def decode_b64_image(b64: str):
     img = Image.open(io.BytesIO(base64.b64decode(b64)))
     img.load()
@@ -653,5 +680,7 @@ __all__ = [
     "generate_synthetic_audio",
     "generate_synthetic_image",
     "generate_synthetic_video",
+    "load_test_audio_data_url",
     "preprocess_text",
+    "test_asset_path",
 ]

tests/helpers/runtime.py

@@ -2,6 +2,7 @@
 
 import base64
 import concurrent.futures
+import errno
 import io
 import json
 import os
@@ -56,10 +57,37 @@ except Exception:  # pragma: no cover
         return None
 
 
-def get_open_port() -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        return int(s.getsockname()[1])
+def get_open_port(host: str = "127.0.0.1", *, max_attempts: int = 128) -> int:
+    """Return a local TCP port that is suitable for binding a new listener.
+
+    A single ``bind(host, 0)`` / close cycle leaves a race where another process can
+    take the same port number before PyTorch/vLLM bind it, yielding
+    ``EADDRINUSE`` / ``DistNetworkError``. We therefore:
+
+    #. Allocate an ephemeral port on *host*.
+    #. Immediately attempt ``bind(host, port)`` again. If that fails with
+       ``errno.EADDRINUSE``, retry from step 1.
+
+    Raises ``RuntimeError`` if no free port is found after *max_attempts* (e.g. port
+    exhaustion under heavy parallel tests).
+    """
+    last_exc: OSError | None = None
+    for _ in range(max_attempts):
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind((host, 0))
+            port = int(s.getsockname()[1])
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
+                probe.bind((host, port))
+        except OSError as exc:
+            last_exc = exc
+            if exc.errno == errno.EADDRINUSE:
+                continue
+            raise
+        return port
+    raise RuntimeError(
+        f"Could not obtain a free TCP port on {host!r} after {max_attempts} attempts (last error: {last_exc!r})"
+    ) from last_exc
 
 
 def dummy_messages_from_mix_data(
@@ -1084,7 +1112,14 @@ class OmniRunner:
         seed: int = 42,
         stage_init_timeout: int = 600,
         batch_timeout: int = 10,
-        init_timeout: int = 900,
+        # Bumped from 900s -> 1800s to give CI cold-cache loads of large
+        # diffusion models enough headroom (Buildkite #8418 hit a 6-second
+        # overrun loading Tongyi-MAI/Z-Image-Turbo: weights alone took 690s,
+        # the full stage was ready at ~896s, but the orchestrator wrapper
+        # finished at ~906s, just past the previous 900s ceiling). Engine
+        # production default in AsyncOmniEngine remains 600s; this only
+        # affects the test runner wrapper.
+        init_timeout: int = 1800,
         shm_threshold_bytes: int = 65536,
         log_stats: bool = False,
         stage_configs_path: str | None = None,

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -411,10 +411,10 @@ def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler():
     runner.input_batch = _DummyInputBatch()
     runner.model = SimpleNamespace(
         prefer_model_sampler=True,
-        sample=lambda logits, sampling_metadata: seen_histories.append(
-            [list(x) for x in sampling_metadata.output_token_ids]
-        )
-        or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None),
+        sample=lambda logits, sampling_metadata: (
+            seen_histories.append([list(x) for x in sampling_metadata.output_token_ids])
+            or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None)
+        ),
     )
     runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
 
@@ -450,10 +450,10 @@ def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler():
     runner.input_batch = _DummyInputBatch()
     runner.model = SimpleNamespace(
         prefer_model_sampler=True,
-        sample=lambda logits, sampling_metadata: seen_histories.append(
-            [list(x) for x in sampling_metadata.output_token_ids]
-        )
-        or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None),
+        sample=lambda logits, sampling_metadata: (
+            seen_histories.append([list(x) for x in sampling_metadata.output_token_ids])
+            or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None)
+        ),
     )
     runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
 

tests/model_executor/models/dynin_omni/test_dynin_omni_token2audio.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from transformers import PreTrainedModel
+
+from vllm_omni.model_executor.models.dynin_omni.dynin_omni_token2audio import (
+    _ensure_transformers_tied_weights_compat,
+)
+
+
+def test_transformers_tied_weights_compat(monkeypatch):
+    monkeypatch.delattr(PreTrainedModel, "all_tied_weights_keys", raising=False)
+
+    _ensure_transformers_tied_weights_compat()
+
+    model = PreTrainedModel.__new__(PreTrainedModel)
+    model._tied_weights_keys = ["decoder.weight"]
+    model._dynamic_tied_weights_keys = {"lm_head.weight": "decoder.weight"}
+
+    assert model.all_tied_weights_keys == {
+        "decoder.weight": None,
+        "lm_head.weight": "decoder.weight",
+    }

vllm_omni/benchmarks/data_modules/seed_tts_eval.py

@@ -351,6 +351,7 @@ def _ensure_en_asr() -> None:
     with _lock:
         if _en_processor is not None:
             return
+        import torch
         from transformers import WhisperForConditionalGeneration, WhisperProcessor
 
         _device = _get_eval_device()
@@ -361,7 +362,13 @@ def _ensure_en_asr() -> None:
             _device,
         )
         _en_processor = WhisperProcessor.from_pretrained(mid)
-        _en_model = WhisperForConditionalGeneration.from_pretrained(mid).to(_device)
+        # Force float32 weights/bias to match `_transcribe_en_f32_16k`'s fp32 protocol
+        # and `WhisperProcessor`'s default fp32 ``input_features``. transformers >=5.x
+        # honors ``model.config.torch_dtype`` in ``from_pretrained``, and
+        # ``openai/whisper-large-v3`` ships ``torch_dtype: float16`` in its config —
+        # without this override conv1 raises
+        # ``RuntimeError: Input type (float) and bias type (c10::Half) should be the same``.
+        _en_model = WhisperForConditionalGeneration.from_pretrained(mid, torch_dtype=torch.float32).to(_device)
         _en_model.eval()
 
 
@@ -392,17 +399,32 @@ def _transcribe_en_f32_16k(wav_f32: np.ndarray) -> str:
         return ""
     with _lock:
         assert _en_processor is not None and _en_model is not None and _device is not None
-        inputs = _en_processor(wav_f32, sampling_rate=16000, return_tensors="pt")
+        try:
+            inputs = _en_processor(
+                wav_f32,
+                sampling_rate=16000,
+                return_tensors="pt",
+                return_attention_mask=True,
+            )
+        except TypeError:
+            inputs = _en_processor(wav_f32, sampling_rate=16000, return_tensors="pt")
         input_features = inputs.input_features.to(_device)
+        attention_mask = getattr(inputs, "attention_mask", None)
+        if attention_mask is None and isinstance(inputs, dict):
+            attention_mask = inputs.get("attention_mask")
+        generate_kwargs: dict[str, Any] = {}
+        if attention_mask is not None:
+            generate_kwargs["attention_mask"] = attention_mask.to(_device)
         with torch.no_grad():
             try:
                 forced = _en_processor.get_decoder_prompt_ids(language="english", task="transcribe")
-                predicted_ids = _en_model.generate(input_features, forced_decoder_ids=forced)
+                predicted_ids = _en_model.generate(input_features, forced_decoder_ids=forced, **generate_kwargs)
             except Exception:
                 predicted_ids = _en_model.generate(
                     input_features,
                     language="english",
                     task="transcribe",
+                    **generate_kwargs,
                 )
         text = _en_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
     return (text or "").strip()

vllm_omni/benchmarks/patch/patch.py

@@ -1,5 +1,4 @@
 import asyncio
-import base64
 import contextlib
 import io
 import json
@@ -15,6 +14,7 @@ from datetime import datetime
 from typing import Literal
 
 import aiohttp
+import pybase64 as base64
 from pydub import AudioSegment
 from tqdm.asyncio import tqdm
 from vllm.benchmarks import datasets
@@ -620,7 +620,8 @@ if "daily-omni" not in OPENAI_COMPATIBLE_BACKENDS:
 # ruff: noqa: E402
 # Prevent import order from causing patch failures
 from vllm.benchmarks import serve
-from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, get_request, wait_for_endpoint
+from vllm.benchmarks.lib.ready_checker import wait_for_endpoint
+from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, get_request
 
 from vllm_omni.benchmarks.metrics.metrics import MultiModalsBenchmarkMetrics, calculate_metrics
 

vllm_omni/config/model.py

@@ -225,6 +225,12 @@ class OmniModelConfig(ModelConfig):
         new_hf_text_config = self.draw_hf_text_config()
         if new_hf_text_config is not self.hf_text_config:
             self.hf_text_config = new_hf_text_config
+            # Recalculate model_arch_config since it derives head counts,
+            # hidden size, etc. from hf_text_config.  Without this the
+            # FlashAttentionMetadataBuilder uses the wrong num_heads_q /
+            # num_heads_kv (from the thinker) for talker stages, causing
+            # FA3 scheduler_metadata shape mismatches at runtime.
+            self.model_arch_config = self.get_model_arch_config()
             # Recalculate dependent attributes
             self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None)
             # Recalculate max_model_len since it depends on hf_text_config

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -340,6 +340,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     request_id=req_id,
                 )
 
+            # Free encoder inputs only after the step has actually executed.
+            if request.has_encoder_inputs:
+                self._free_encoder_inputs(request)
+
             stopped = False
             is_segment_finished = False
             new_logprobs = None
@@ -364,6 +368,20 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if not stopped and self._process_kv_transfer_trigger(request, new_token_ids):
                 stopped = True
 
+            if new_token_ids and self.structured_output_manager.should_advance(request):
+                struct_output_request = request.structured_output_request
+                assert struct_output_request is not None
+                assert struct_output_request.grammar is not None
+                if not struct_output_request.grammar.accept_tokens(req_id, new_token_ids):
+                    logger.error(
+                        "Unexpected: grammar rejected tokens %s for request %s. Terminating request.",
+                        new_token_ids,
+                        req_id,
+                    )
+                    request.status = RequestStatus.FINISHED_ERROR
+                    request.resumable = False
+                    stopped = True
+
             if stopped:
                 routed_experts = self._get_routed_experts(request)
 
@@ -388,18 +406,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
                 new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
 
-            if new_token_ids and self.structured_output_manager.should_advance(request):
-                struct_output_request = request.structured_output_request
-                assert struct_output_request is not None
-                assert struct_output_request.grammar is not None
-                ok = struct_output_request.grammar.accept_tokens(req_id, new_token_ids)
-                if not ok:
-                    logger.warning(
-                        "Unexpected: grammar rejected tokens %s for request %s.",
-                        new_token_ids,
-                        req_id,
-                    )
-
             if num_nans_in_logits is not None and req_id in num_nans_in_logits:
                 request.num_nans_in_logits = num_nans_in_logits[req_id]
 
@@ -417,10 +423,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                         pooling_output=pooler_output,
                         stop_reason=request.stop_reason,
                         events=request.take_events(),
+                        prefill_stats=request.take_prefill_stats(),
                         kv_transfer_params=kv_transfer_params,
                         trace_headers=request.trace_headers,
-                        num_cached_tokens=request.num_cached_tokens,
-                        num_external_computed_tokens=request.num_external_computed_tokens,
                         routed_experts=routed_experts,
                         num_nans_in_logits=request.num_nans_in_logits,
                         is_segment_finished=is_segment_finished,
@@ -453,7 +458,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                         finish_reason=request.get_finished_reason(),
                         events=request.take_events(),
                         trace_headers=request.trace_headers,
-                        num_cached_tokens=request.num_cached_tokens,
                     )
                 )
                 if self.chunk_transfer_adapter is not None:

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -120,8 +120,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             req_to_new_blocks[request.request_id] = new_blocks
             num_scheduled_tokens[request.request_id] = num_new_tokens
             cached_prompt_token_ids[request.request_id] = request.prompt_token_ids
-            if request.num_cached_tokens < 0:
-                request.num_cached_tokens = num_computed_tokens
             cached_additional_information[request.request_id] = getattr(request, "additional_information", None)
             token_budget -= num_new_tokens
             scheduled_running_reqs.append(request)
@@ -147,9 +145,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Pop the finished request from waiting queue and don't schedule it
                 self.waiting.pop_request()
                 continue
-            # Count the number of prefix cached tokens.
-            if request.num_cached_tokens < 0:
-                request.num_cached_tokens = request.num_computed_tokens
 
             # async_chunk: wait for the first upstream chunk (don't start with placeholders).
             if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0:
@@ -191,8 +186,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             req_to_new_blocks[request.request_id] = new_blocks
             num_scheduled_tokens[request.request_id] = num_new_tokens
-            if request.num_cached_tokens < 0:
-                request.num_cached_tokens = 0
             token_budget -= num_new_tokens
             scheduled_new_reqs.append(request)
 
@@ -441,6 +434,10 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     request_id=req_id,
                 )
 
+            # Free encoder inputs only after the step has actually executed.
+            if request.has_encoder_inputs:
+                self._free_encoder_inputs(request)
+
             stopped = False
             new_logprobs = None
             new_token_ids = generated_token_ids
@@ -502,11 +499,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
             if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
-                # Add EngineCoreOutput for this Request.
-                num_cached = request.num_cached_tokens
-                if num_cached < 0:
-                    logger.warning("Negative num_cached_tokens (%d) for request %s, clamping to 0", num_cached, req_id)
-                    num_cached = 0
                 outputs[request.client_index].append(
                     EngineCoreOutput(
                         request_id=req_id,
@@ -517,10 +509,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                         pooling_output=pooler_output,
                         stop_reason=request.stop_reason,
                         events=request.take_events(),
+                        prefill_stats=request.take_prefill_stats(),
                         kv_transfer_params=kv_transfer_params,
                         trace_headers=request.trace_headers,
-                        num_cached_tokens=num_cached,
-                        num_external_computed_tokens=request.num_external_computed_tokens,
                         routed_experts=routed_experts,
                         num_nans_in_logits=request.num_nans_in_logits,
                     )
@@ -542,12 +533,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
             self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
             for request in requests:
-                num_cached = request.num_cached_tokens
-                if num_cached < 0:
-                    logger.warning(
-                        "Negative num_cached_tokens (%d) for request %s, clamping to 0", num_cached, request.request_id
-                    )
-                    num_cached = 0
                 outputs[request.client_index].append(
                     EngineCoreOutput(
                         request_id=request.request_id,
@@ -555,7 +540,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                         finish_reason=request.get_finished_reason(),
                         events=request.take_events(),
                         trace_headers=request.trace_headers,
-                        num_cached_tokens=num_cached,
                     )
                 )
                 if self.chunk_transfer_adapter is not None:

vllm_omni/diffusion/data.py

@@ -776,6 +776,18 @@ class OmniDiffusionConfig:
         if "diffusers_call_kwargs" in kwargs and kwargs["diffusers_call_kwargs"] is None:
             kwargs["diffusers_call_kwargs"] = {}
 
+        # Forward top-level parallel knobs (e.g. --tensor-parallel-size from CLI)
+        # into parallel_config so the diffusion engine sees them.
+        par = kwargs.get("parallel_config", {})
+        if isinstance(par, Mapping):
+            par = dict(par)
+            if par.get("tensor_parallel_size") is None:
+                par.pop("tensor_parallel_size", None)
+            tensor_parallel_size = kwargs.get("tensor_parallel_size")
+            if tensor_parallel_size is not None and "tensor_parallel_size" not in par:
+                par["tensor_parallel_size"] = tensor_parallel_size
+            kwargs["parallel_config"] = par
+
         # Filter kwargs to only include valid fields
         valid_fields = {f.name for f in fields(cls)}
         filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields}

vllm_omni/diffusion/model_loader/hub_prefetch.py

@@ -0,0 +1,314 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Best-effort HuggingFace Hub prefetch for multi-subfolder pipelines.
+
+This module exists to defend diffusion pipelines against a race condition we
+hit after the `transformers` v5 rebase (see Buildkite vllm-omni-rebase
+#1043 Qwen-Image-Edit-2509 failure):
+
+When several diffusion worker processes start in parallel and each calls
+``SomeModel.from_pretrained(model_id, subfolder="text_encoder", ...)`` with a
+cold HuggingFace cache, transformers v5's cache-resolution (``cached_files``)
+can observe a partially-written shard set written by a peer worker and raise
+``OSError: <model_id> does not appear to have a file named
+text_encoder/model-00002-of-00002.safetensors`` even though the peer will
+eventually finish writing it.
+
+Why ``origin/main`` does not need this helper
+---------------------------------------------
+The exact same ``__init__`` code lives on ``origin/main`` (e.g. the Qwen-Image
+``pipeline_qwen_image_edit_plus.py`` in build vllm-omni#7412 passes without
+any prefetch), so the race is NOT a behavioural change in vLLM-Omni itself.
+Two environmental factors mask the race on main:
+
+* ``origin/main`` is pinned (transitively, via vLLM main) to
+  ``transformers`` 4.x. In 4.x the per-file ``cached_file`` path resolves
+  shards **lazily**, one at a time, so each ``hf_hub_download`` blocks on its
+  own single-file ``.lock`` and the second worker naturally waits for the
+  first worker's atomic rename. ``transformers>=5.0`` rewrote this into
+  ``cached_files`` (plural) which batch-resolves every shard listed in the
+  index up-front via ``os.path.isfile`` and raises immediately if any shard
+  is still sitting under its ``*.incomplete`` name. Same wave of v5 changes
+  that introduced ``tie_weights(missing_keys=..., recompute_mapping=...)``
+  (see the Dynin shim in ``dynin_omni_token2text.py``).
+* CI shares ``HF_HOME=/fsx/hf_cache`` across pipelines (both the
+  ``vllm-omni`` and ``vllm-omni-rebase`` pipelines mount the same FS). That
+  cache is normally warm for long-lived repos like ``Qwen-Image-Edit-2509``,
+  so most builds never go through the download path at all. Build 1043
+  happened to hit a partially-evicted cache AND transformers v5's stricter
+  resolver simultaneously, which is why the failure looks 'rebase-specific'
+  but is really a latent race that main was absorbing via (1).
+
+``huggingface_hub.snapshot_download`` does take per-blob ``.lock`` files,
+but those locks are only acquired once each blob is mid-write - they do not
+cover the surrounding ``cached_files`` shard-list resolution that
+transformers v5 performs eagerly. So we additionally wrap the
+``snapshot_download`` call in our own node-wide ``fcntl.flock`` keyed on the
+repo id (see ``_repo_prefetch_lock``). The first concurrent worker / process
+to reach the helper fully materialises the snapshot; subsequent entrants
+block on the flock and then find a warm cache, so their ``from_pretrained``
+calls never observe a half-written shard set. For a warm cache the snapshot
+call is a near-noop (it only stat()s the files), so this is also cheap on
+``origin/main`` should we ever backport it there.
+
+The helper is intentionally best-effort: prefetch failures (offline, gated
+repos, transient 5xx, missing flock support) are logged and swallowed so the
+subsequent ``from_pretrained`` call can surface the real, specific error to
+the user rather than being masked here.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import errno
+import logging
+import os
+from collections.abc import Iterable, Iterator
+
+logger = logging.getLogger(__name__)
+
+
+def _node_lock_dir() -> str:
+    """Return a node-local directory suitable for the prefetch lock files.
+
+    We co-locate the lock file with the HF cache it is serialising
+    (``$HF_HOME/locks/vllm-omni-prefetch/``) so they share a filesystem -
+    ``fcntl.flock`` semantics across NFS / Lustre are only well-defined
+    when the underlying FS supports POSIX advisory locking. CI's
+    ``/fsx/hf_cache`` (FSx for Lustre) does. If ``HF_HOME`` is unwritable
+    (read-only baked image, etc.) we fall back to ``XDG_CACHE_HOME``,
+    ``~/.cache/huggingface``, then ``/tmp`` so unit tests outside CI keep
+    working.
+    """
+    candidates: list[str] = []
+    hf_home = os.environ.get("HF_HOME")
+    if hf_home:
+        candidates.append(os.path.join(hf_home, "locks", "vllm-omni-prefetch"))
+    xdg_cache = os.environ.get("XDG_CACHE_HOME")
+    if xdg_cache:
+        candidates.append(os.path.join(xdg_cache, "huggingface", "locks", "vllm-omni-prefetch"))
+    candidates.append(os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "locks", "vllm-omni-prefetch"))
+    candidates.append(os.path.join("/tmp", "vllm-omni-prefetch-locks"))
+
+    for cand in candidates:
+        try:
+            os.makedirs(cand, exist_ok=True)
+            probe = os.path.join(cand, ".write_check")
+            with open(probe, "a"):
+                pass
+            with contextlib.suppress(OSError):
+                os.remove(probe)
+            return cand
+        except OSError:
+            continue
+    fallback = os.path.join("/tmp", "vllm-omni-prefetch-locks")
+    os.makedirs(fallback, exist_ok=True)
+    return fallback
+
+
+def _safe_repo_filename(model: str) -> str:
+    """Translate an HF repo id into a filesystem-safe lock filename."""
+    return model.replace("/", "__").replace(os.sep, "__") + ".lock"
+
+
+@contextlib.contextmanager
+def _repo_prefetch_lock(model: str) -> Iterator[None]:
+    """Hold an exclusive ``fcntl.flock`` keyed on the HF repo id.
+
+    Why we still need a node-wide lock on top of ``snapshot_download``'s
+    own per-blob ``.lock`` files:
+
+    * ``huggingface_hub.snapshot_download`` does serialise concurrent
+      writers within the same call, but its lock lives in the per-blob
+      cache directory and is only acquired once each blob is being
+      written. With ``transformers`` v5 the multi-shard ``cached_files``
+      resolver inspecting the cache between blob acquires is exactly
+      the window that races - a peer worker's ``model-00001-of-00002.
+      safetensors.incomplete`` makes ``cached_files`` raise ``OSError:
+      <repo> does not appear to have a file named ...`` even though
+      the peer is mid-download (Buildkite #8418).
+    * Multiple unrelated processes on the same node (concurrent CI jobs,
+      multiple ``OmniServer`` instances, every ``DiffusionWorker``
+      subprocess from the multiproc executor) can enter this code path
+      in parallel. ``flock`` on a per-repo file makes the whole
+      ``snapshot_download(...)`` call a true node-wide critical
+      section, so the *first* entrant fully populates the cache and
+      every subsequent ``from_pretrained`` sees a warm, complete tree.
+
+    Best-effort: if we cannot create / acquire the lock (no ``fcntl``,
+    read-only FS, NFS without flock support) we log and run the download
+    anyway. Worst case behaviour reverts to "snapshot_download +
+    transformers v5 race" which is exactly the pre-fix state.
+    """
+    try:
+        import fcntl  # type: ignore[import-not-found]
+    except ImportError:  # pragma: no cover - non-POSIX (Windows)
+        logger.debug("fcntl unavailable on this platform; running prefetch unlocked")
+        yield
+        return
+
+    try:
+        lock_dir = _node_lock_dir()
+    except OSError as exc:
+        logger.warning("Could not allocate lock dir for prefetch of %s (%s); skipping flock", model, exc)
+        yield
+        return
+
+    lock_path = os.path.join(lock_dir, _safe_repo_filename(model))
+
+    try:
+        fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644)
+    except OSError as exc:
+        logger.warning("Could not open prefetch lock %s (%s); proceeding without flock", lock_path, exc)
+        yield
+        return
+
+    locked = False
+    try:
+        try:
+            fcntl.flock(fd, fcntl.LOCK_EX)
+            locked = True
+            logger.debug("Acquired prefetch flock for %s at %s", model, lock_path)
+        except OSError as exc:
+            # FS without flock support (ENOLCK / EOPNOTSUPP / EINVAL).
+            # EACCES on some Lustre configs without ``flock`` mount opt.
+            if exc.errno in (errno.ENOLCK, errno.EOPNOTSUPP, errno.EACCES, errno.EINVAL):
+                logger.warning(
+                    "fcntl.flock not supported on '%s' (%s: %s); running prefetch unlocked",
+                    lock_path,
+                    type(exc).__name__,
+                    exc,
+                )
+            else:
+                raise
+        yield
+    finally:
+        if locked:
+            with contextlib.suppress(OSError):
+                fcntl.flock(fd, fcntl.LOCK_UN)
+        with contextlib.suppress(OSError):
+            os.close(fd)
+
+
+def prefetch_subfolders(
+    model: str,
+    subfolders: Iterable[str],
+    *,
+    local_files_only: bool = False,
+    include_root_metadata: bool = True,
+) -> None:
+    """Materialise ``model``'s ``subfolders`` in the HF cache before loading.
+
+    Args:
+        model: A HuggingFace Hub repo id (e.g. ``"Qwen/Qwen-Image-Edit-2509"``)
+            or a local directory path. Local paths are a no-op.
+        subfolders: Iterable of subfolder names (e.g. ``["text_encoder",
+            "vae"]``) whose contents need to be fully present before any
+            worker calls ``from_pretrained(subfolder=...)``.
+        local_files_only: When True, skip the prefetch entirely. The caller
+            has explicitly promised the cache is already populated (as happens
+            for local model checkouts), so hitting the network would defeat
+            the intent and may fail in air-gapped environments.
+        include_root_metadata: When True, also pull ``*.json`` at the repo
+            root so ``model_index.json`` / ``config.json`` resolution during
+            ``from_pretrained`` also hits a warm cache.
+    """
+    if local_files_only or not model or os.path.isdir(model):
+        return
+
+    try:
+        from huggingface_hub import snapshot_download
+    except ImportError:  # pragma: no cover - huggingface_hub is a hard dep
+        logger.debug("huggingface_hub unavailable; skipping prefetch of %s", model)
+        return
+
+    allow_patterns: list[str] = []
+    for sub in subfolders:
+        sub = (sub or "").strip("/")
+        if not sub:
+            continue
+        # hf_hub globbing is shell-style: `text_encoder/*` catches the index +
+        # any flat files, `text_encoder/**` catches nested safetensors shards
+        # that some repos place under e.g. ``text_encoder/pytorch_model/``.
+        allow_patterns.extend([f"{sub}/*", f"{sub}/**"])
+
+    if include_root_metadata:
+        allow_patterns.extend(["*.json", "*.txt"])
+
+    if not allow_patterns:
+        return
+
+    # Wrap the download in a node-wide ``fcntl`` flock keyed on the repo
+    # id. The first concurrent entrant fully populates the cache; every
+    # subsequent worker / process blocks here, then no-ops because the
+    # snapshot is now warm. This is what makes the prefetch race-free
+    # even when many ``DiffusionWorker`` subprocesses (or multiple
+    # OmniServer instances on the same node) hit this code in parallel.
+    try:
+        with _repo_prefetch_lock(model):
+            snapshot_download(
+                repo_id=model,
+                allow_patterns=allow_patterns,
+            )
+    except Exception as exc:
+        # Best-effort: propagate only via logging. The subsequent
+        # ``from_pretrained`` call will raise a clearer, call-site-specific
+        # error (auth, 404, disk full, ...) that we'd rather surface - EXCEPT
+        # for auth/gating, which we escalate here with an explicit hint so
+        # readers of CI logs don't have to correlate the generic "OSError:
+        # <repo> does not appear to have a file named ..." that
+        # ``from_pretrained`` would otherwise emit much later with an
+        # unrelated-looking message.
+        if _looks_like_auth_error(exc):
+            logger.error(
+                "Hub prefetch for '%s' failed with an authentication / gated "
+                "repository error (%s: %s). The CI HF_TOKEN must (1) be set "
+                "in the step env, (2) be valid, and (3) belong to an account "
+                "that has accepted the model license on huggingface.co. See "
+                "docs/contributing/ci/hf_credentials.md.",
+                model,
+                type(exc).__name__,
+                exc,
+            )
+        else:
+            logger.warning(
+                "Hub prefetch for repo '%s' subfolders %s failed (%s: %s); "
+                "falling back to on-demand download in from_pretrained",
+                model,
+                list(subfolders),
+                type(exc).__name__,
+                exc,
+            )
+
+
+def _looks_like_auth_error(exc: BaseException) -> bool:
+    """Classify prefetch exceptions as auth/gating failures.
+
+    ``huggingface_hub`` raises bespoke ``GatedRepoError`` /
+    ``RepositoryNotFoundError`` subclasses when the token is missing or
+    lacks license acceptance; but older releases (and third-party
+    transport layers) sometimes only surface this as a generic
+    ``HfHubHTTPError`` / ``requests.HTTPError`` with status 401/403. We
+    check both code paths so the branch above is stable across
+    ``huggingface_hub`` versions.
+    """
+    try:
+        from huggingface_hub.errors import (  # type: ignore[import-not-found]
+            GatedRepoError,
+            RepositoryNotFoundError,
+        )
+
+        if isinstance(exc, GatedRepoError | RepositoryNotFoundError):
+            return True
+    except ImportError:  # pragma: no cover - very old huggingface_hub
+        pass
+
+    status = getattr(getattr(exc, "response", None), "status_code", None)
+    if status in (401, 403):
+        return True
+    # Last-resort string heuristic - ``snapshot_download`` on some
+    # transports wraps the 401 as a plain OSError whose message is the
+    # only load-bearing signal.
+    msg = str(exc).lower()
+    return "401 client error" in msg or "403 client error" in msg or "gatedrepo" in msg

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -114,17 +114,32 @@ class BagelRotaryEmbedding(nn.Module):
     def __init__(self, config):
         super().__init__()
 
-        if config.rope_scaling is not None:
-            # Delegate to HF's rope-scaling helpers for non-default types.
+        # transformers>=5.0 stores rope params under ``rope_parameters`` and
+        # always populates ``rope_scaling`` (even for the default type), so we
+        # infer ``rope_type`` from either and fall back to "default".
+        rope_scaling = getattr(config, "rope_scaling", None) or {}
+        rope_parameters = getattr(config, "rope_parameters", None) or {}
+        rope_type = (
+            rope_scaling.get("rope_type") or rope_scaling.get("type") or rope_parameters.get("rope_type") or "default"
+        )
+
+        if rope_type == "default":
+            # transformers>=5.0 removed the 'default' entry from
+            # ROPE_INIT_FUNCTIONS; use the plain sinusoidal formula.
+            rope_theta = (
+                getattr(config, "rope_theta", None)
+                or rope_parameters.get("rope_theta")
+                or rope_scaling.get("rope_theta")
+                or 10000.0
+            )
+            dim = config.hidden_size // config.num_attention_heads
+            inv_freq = 1.0 / (rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+            self.attention_scaling = 1.0
+        else:
             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)
 

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -39,6 +39,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
@@ -209,6 +210,15 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
         model = od_config.model
         local_files_only = os.path.exists(model)
 
+        # Avoid the transformers v5 multi-worker subfolder race (see
+        # ``vllm_omni/diffusion/model_loader/hub_prefetch.py`` for the full
+        # analysis; L4 build #1043 hit this on FLUX.2-klein-4B's text_encoder).
+        prefetch_subfolders(
+            model,
+            ["scheduler", "text_encoder", "tokenizer", "vae"],
+            local_files_only=local_files_only,
+        )
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model,
             subfolder="scheduler",

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py

@@ -4,7 +4,7 @@
 from typing import Any
 
 import vllm.forward_context as _vllm_fc
-from vllm.model_executor.layers.fused_moe import SharedFusedMoE
+from vllm.model_executor.layers.fused_moe import FusedMoE
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
 from vllm_omni.platforms import current_omni_platform
@@ -13,7 +13,7 @@ from vllm_omni.platforms import current_omni_platform
 def _set_forward_context_num_tokens(num_tokens: int) -> None:
     """Set num_tokens on the vLLM ForwardContext for MoE routing.
 
-    After the rebase to vLLM 0.18.0, SharedFusedMoE expects
+    After the rebase to vLLM 0.18.0, FusedMoE expects
     ForwardContext.num_tokens to be set. Without it, MoE expert
     routing may produce incorrect results (silent correctness bug).
     """
@@ -25,7 +25,7 @@ def _set_forward_context_num_tokens(num_tokens: int) -> None:
         forward_context.in_profile_run = False
 
 
-class HunyuanFusedMoEDefault(SharedFusedMoE):
+class HunyuanFusedMoEDefault(FusedMoE):
     def __init__(self, *, prefix: str = "", **kwargs: Any) -> None:
         super().__init__(prefix=prefix, **kwargs)
         self._prefix = prefix

vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py

@@ -48,7 +48,6 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import (
 )
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.t5_encoder.t5_gemma_encoder import T5GemmaEncoderModelTP
-from vllm_omni.diffusion.models.utils import _load_json
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
     DiffusionPipelineProfilerMixin,
 )
@@ -1641,6 +1640,20 @@ def get_magi_human_post_process_func(*args, **kwargs):
 # ===========================================================================
 
 
+def _load_json(model_path: str, filename: str, local_files_only: bool = True) -> dict:
+    """Load a JSON config file from a local path or HuggingFace Hub repo."""
+    if local_files_only:
+        path = os.path.join(model_path, *filename.split("/"))
+        with open(path) as f:
+            return json.load(f)
+    else:
+        from huggingface_hub import hf_hub_download
+
+        cached = hf_hub_download(repo_id=model_path, filename=filename)
+        with open(cached) as f:
+            return json.load(f)
+
+
 def _resolve_subdir(
     model_path: str,
     subfolder: str,

vllm_omni/diffusion/models/nextstep_1_1/modeling_nextstep_llama.py

@@ -15,6 +15,36 @@ from vllm.model_executor.layers.linear import (
     RowParallelLinear,
 )
 
+
+def _default_rope_init(config, device=None, seq_len=None, layer_type=None):
+    """Vanilla sinusoidal RoPE (no scaling).
+
+    transformers>=5.0 removed the 'default' entry from
+    ``ROPE_INIT_FUNCTIONS``; reimplement the original behaviour so that
+    ``rope_type='default'`` configs keep working.
+
+    transformers>=5.0 also no longer promotes ``rope_theta`` to a top-level
+    attribute on configs that only declare it inside ``rope_parameters`` /
+    ``rope_scaling`` (e.g. NextStep-1.1). Resolve it defensively before
+    computing the inverse frequency so the init does not explode on those
+    configs.
+    """
+    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    rope_theta = getattr(config, "rope_theta", None)
+    if rope_theta is None:
+        for attr in ("rope_parameters", "rope_scaling"):
+            params = getattr(config, attr, None)
+            if isinstance(params, dict) and "rope_theta" in params:
+                rope_theta = params["rope_theta"]
+                break
+    if rope_theta is None:
+        rope_theta = 10000.0
+    inv_freq = 1.0 / (
+        float(rope_theta) ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim)
+    )
+    return inv_freq, 1.0
+
+
 # ---------------------------------------------------------------------------
 # Utilities (inlined from remote utils/model_utils.py)
 # ---------------------------------------------------------------------------
@@ -77,7 +107,7 @@ class LlamaRotaryEmbedding(nn.Module):
         super().__init__()
         self.rope_type = "default"
         self.config = config
-        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS.get(self.rope_type, _default_rope_init)
         inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
 

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -26,6 +26,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
 )
@@ -274,6 +275,14 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         # Check if model is a local path
         local_files_only = os.path.exists(model)
 
+        # See pipeline_qwen_image_edit_plus: guard against transformers v5
+        # multi-worker race on partial subfolder shard sets (Buildkite #1043).
+        prefetch_subfolders(
+            model,
+            ["scheduler", "text_encoder", "vae", "tokenizer"],
+            local_files_only=local_files_only,
+        )
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -27,6 +27,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
@@ -244,6 +245,14 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         # Check if model is a local path
         local_files_only = os.path.exists(model)
 
+        # See pipeline_qwen_image_edit_plus: guard against transformers v5
+        # multi-worker race on partial subfolder shard sets (Buildkite #1043).
+        prefetch_subfolders(
+            model,
+            ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
+            local_files_only=local_files_only,
+        )
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -25,6 +25,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
 from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
@@ -209,6 +210,18 @@ class QwenImageEditPlusPipeline(
         # Check if model is a local path
         local_files_only = os.path.exists(model)
 
+        # Defend against a transformers v5 multi-worker race where a peer
+        # worker's partially-written shard set makes our from_pretrained
+        # subfolder resolution fail with OSError (observed in Buildkite
+        # #1043 for Qwen/Qwen-Image-Edit-2509). snapshot_download takes a
+        # per-repo file lock so the first worker downloads and the rest
+        # wait on a warm cache before loading.
+        prefetch_subfolders(
+            model,
+            ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
+            local_files_only=local_files_only,
+        )
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -24,6 +24,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.autoencoder_kl_qwenimage import (
     AutoencoderKLQwenImage,
@@ -217,6 +218,14 @@ class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParalle
         # Check if model is a local path
         local_files_only = os.path.exists(model)
 
+        # See pipeline_qwen_image_edit_plus: guard against transformers v5
+        # multi-worker race on partial subfolder shard sets (Buildkite #1043).
+        prefetch_subfolders(
+            model,
+            ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
+            local_files_only=local_files_only,
+        )
+
         # modules keep same as transformers & diffusers
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only

vllm_omni/diffusion/models/sd3/pipeline_sd3.py

@@ -19,6 +19,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import Distribu
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
 from vllm_omni.diffusion.models.sd3.sd3_transformer import (
     SD3Transformer2DModel,
 )
@@ -152,6 +153,41 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         # Check if model is a local path
         local_files_only = os.path.exists(model)
 
+        # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
+        # SD3.5 loads six subfolders in a row, each with multi-shard
+        # safetensors - it is the worst-case fan-out for the race window.
+        prefetch_subfolders(
+            model,
+            [
+                "scheduler",
+                "tokenizer",
+                "tokenizer_2",
+                "tokenizer_3",
+                "text_encoder",
+                "text_encoder_2",
+                "text_encoder_3",
+                "vae",
+            ],
+            local_files_only=local_files_only,
+        )
+
+        # ``transformers`` v5 changed ``from_pretrained``'s default
+        # ``torch_dtype`` from fp32 to whatever ``config.torch_dtype`` is
+        # baked into the checkpoint (fp16 for ``stabilityai/stable-
+        # diffusion-3.5-medium``). Without an explicit override the text
+        # encoders + VAE end up in fp16 while the rest of the pipeline
+        # (transformer, ``CombinedTimestepTextProjEmbeddings``,
+        # ``PatchEmbed``) is materialised in ``od_config.dtype`` (bf16
+        # by default). The first time latents flow through the
+        # transformer this surfaces as ``RuntimeError: Input type
+        # (c10::Half) and bias type (c10::BFloat16) should be the same``
+        # inside ``SD3Transformer2DModel.pos_embed`` (Buildkite #8418
+        # ``test_sd3_medium``). Pin every weight-loading
+        # ``from_pretrained`` to ``od_config.dtype`` so the whole
+        # pipeline is dtype-consistent. Tokenizers and the scheduler
+        # don't carry weights and therefore don't take ``torch_dtype``.
+        dtype = od_config.dtype
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
@@ -163,20 +199,30 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
             model, subfolder="tokenizer_3", local_files_only=local_files_only
         )
         self.text_encoder = CLIPTextModelWithProjection.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
+            model,
+            subfolder="text_encoder",
+            torch_dtype=dtype,
+            local_files_only=local_files_only,
         )
         self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
-            model, subfolder="text_encoder_2", local_files_only=local_files_only
+            model,
+            subfolder="text_encoder_2",
+            torch_dtype=dtype,
+            local_files_only=local_files_only,
         )
         self.text_encoder_3 = T5EncoderModel.from_pretrained(
             model,
             subfolder="text_encoder_3",
+            torch_dtype=dtype,
             local_files_only=local_files_only,
         )
         self.transformer = SD3Transformer2DModel(od_config=od_config)
 
         self.vae = DistributedAutoencoderKL.from_pretrained(
-            model, subfolder="vae", local_files_only=local_files_only
+            model,
+            subfolder="vae",
+            torch_dtype=dtype,
+            local_files_only=local_files_only,
         ).to(self.device)
 
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8

vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py

@@ -360,7 +360,6 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil
         negative_prompt: str | list[str] | None = None,
         audio_end_in_s: float | None = None,
         audio_start_in_s: float = 0.0,
-        num_inference_steps: int = 100,
         guidance_scale: float = 7.0,
         num_waveforms_per_prompt: int = 1,
         generator: torch.Generator | list[torch.Generator] | None = None,
@@ -378,7 +377,6 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil
             negative_prompt: Negative prompt for CFG
             audio_end_in_s: Audio end time in seconds (max ~47s for stable-audio-open-1.0)
             audio_start_in_s: Audio start time in seconds
-            num_inference_steps: Number of denoising steps
             guidance_scale: CFG scale
             num_waveforms_per_prompt: Number of audio outputs per prompt
             generator: Random generator for reproducibility
@@ -398,8 +396,12 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
-
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        num_inference_steps = req.sampling_params.num_inference_steps
+        if num_inference_steps is None:
+            num_inference_steps = 100  # Default steps
+        num_inference_steps = req.sampling_params.num_inference_steps
+        if num_inference_steps is None:
+            num_inference_steps = 50  # Default steps
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
 

vllm_omni/diffusion/models/wan2_2/patch_diffusers.py

@@ -9,6 +9,10 @@ from vllm_omni.diffusion.layers.norm import RMSNormVAE
 def patch_wan_rms_norm():
     """Patch diffusers Wan RMSNorm implementation with RMSNormVAE."""
 
-    for module_name, module in sys.modules.items():
+    # NOTE: iterate over a snapshot of sys.modules. `hasattr` can trigger lazy
+    # submodule imports (e.g. transformers' `_LazyModule.__getattr__`), which
+    # mutate sys.modules during iteration and raise
+    # `RuntimeError: dictionary changed size during iteration`.
+    for module_name, module in list(sys.modules.items()):
         if hasattr(module, "WanRMS_norm"):
             setattr(module, "WanRMS_norm", RMSNormVAE)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -22,6 +22,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import Dist
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
@@ -316,6 +317,13 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
                 )
             )
 
+        # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
+        prefetch_subfolders(
+            model,
+            ["tokenizer", "text_encoder", "vae"],
+            local_files_only=local_files_only,
+        )
+
         self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
         self.text_encoder = UMT5EncoderModel.from_pretrained(
             model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -23,6 +23,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import Dist
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
@@ -196,6 +197,12 @@ class Wan22I2VPipeline(
         # Image encoder (CLIP) - optional, for Wan2.1-style I2V
         self.has_image_encoder = "image_encoder" in model_index and model_index["image_encoder"][0] is not None
 
+        # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
+        subfolders = ["tokenizer", "text_encoder", "vae"]
+        if self.has_image_encoder:
+            subfolders.extend(["image_processor", "image_encoder"])
+        prefetch_subfolders(model, subfolders, local_files_only=local_files_only)
+
         if self.has_image_encoder:
             self.image_processor = CLIPImageProcessor.from_pretrained(
                 model, subfolder="image_processor", local_files_only=local_files_only

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py

@@ -34,6 +34,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import Omni
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
@@ -176,6 +177,13 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress
             ),
         ]
 
+        # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
+        prefetch_subfolders(
+            model,
+            ["tokenizer", "text_encoder", "vae"],
+            local_files_only=local_files_only,
+        )
+
         # Text encoder
         self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
         self.text_encoder = UMT5EncoderModel.from_pretrained(

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -35,6 +35,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
 from vllm_omni.diffusion.models.z_image.z_image_transformer import (
     ZImageTransformer2DModel,
 )
@@ -179,6 +180,14 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
         self._execution_device = get_local_device()
         model = od_config.model
         local_files_only = os.path.exists(model)
+
+        # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
+        prefetch_subfolders(
+            model,
+            ["scheduler", "text_encoder", "vae", "tokenizer"],
+            local_files_only=local_files_only,
+        )
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )

vllm_omni/distributed/kv_transfer/monkey_patch.py

@@ -182,7 +182,9 @@ def apply_mooncake_connector_patch() -> bool:
     PatchedClass = _create_patched_mooncake_connector()
 
     _mc_module.MooncakeConnector = PatchedClass
-    for _, module in sys.modules.items():
+    # Snapshot sys.modules: hasattr() may trigger lazy submodule imports that
+    # mutate sys.modules during iteration.
+    for _, module in list(sys.modules.items()):
         if hasattr(module, "MooncakeConnector") and module.MooncakeConnector is _OriginalClass:
             module.MooncakeConnector = PatchedClass
 

vllm_omni/engine/async_omni_engine.py

@@ -13,6 +13,7 @@ import dataclasses
 import json
 import os
 import queue
+import sys
 import threading
 import time
 import uuid
@@ -65,6 +66,7 @@ from vllm_omni.engine.stage_init_utils import (
     StartedLlmStage,
     _inject_inferred_kv_tp_topology,
     acquire_device_locks,
+    acquire_diffusion_device_locks,
     build_diffusion_config,
     build_engine_args_dict,
     build_vllm_config,
@@ -110,6 +112,13 @@ _PARENT_ARGS_KEEP: frozenset[str] = frozenset(
         "worker_extension_cls",
         "allowed_local_media_path",
         "allowed_media_domains",
+        # Legacy stage-config YAMLs may intentionally leave parallel or
+        # distributed knobs unspecified at the stage level and rely on
+        # top-level CLI values to fill them in during the per-stage merge.
+        # Keep these fields so stages that omit them can inherit CLI values,
+        # while stages with explicit YAML values still win because the legacy
+        # stage-config loader prefers stage-local engine args.
+        "tensor_parallel_size",
     }
 )
 
@@ -352,6 +361,7 @@ class AsyncOmniEngine:
             name="orchestrator",
         )
         self.orchestrator_thread.start()
+
         self._wait_for_orchestrator_init(startup_future, startup_timeout)
 
         # Stage runtime fields are assigned directly on self by the bootstrap thread.
@@ -518,7 +528,8 @@ class AsyncOmniEngine:
                 stage_id=metadata.stage_id,
             )
             logger.info("[AsyncOmniEngine] Stage %s remote engine handshake started", metadata.stage_id)
-            with launch_cm as (engine_manager, coordinator, addresses):
+            (engine_manager, coordinator, addresses, _tensor_queue) = launch_cm.__enter__()
+            try:
                 started_stage = StartedLlmStage(
                     stage_id=metadata.stage_id,
                     metadata=metadata,
@@ -528,6 +539,11 @@ class AsyncOmniEngine:
                     coordinator=coordinator,
                     addresses=addresses,
                 )
+            except BaseException:
+                if not launch_cm.__exit__(*sys.exc_info()):
+                    raise
+            else:
+                launch_cm.__exit__(None, None, None)
             logger.info("[AsyncOmniEngine] Stage %s remote engine startup completed", metadata.stage_id)
             assert started_stage is not None
             return started_stage
@@ -541,11 +557,14 @@ class AsyncOmniEngine:
         stage_cfg: Any,
         metadata: Any,
         omni_master_server: OmniMasterServer,
+        stage_init_timeout: int,
     ) -> StageDiffusionClient:
         """Launch a local diffusion stage on OmniMasterServer-allocated sockets."""
         proc = None
+        lock_fds: list[int] = []
         try:
             od_config = build_diffusion_config(self.model, stage_cfg, metadata)
+            lock_fds = acquire_diffusion_device_locks(metadata.stage_id, od_config, stage_init_timeout)
             handshake_address, request_address, response_address = register_stage_with_omni_master(
                 omni_master_address=omni_master_server.address,
                 omni_master_port=omni_master_server.port,
@@ -564,7 +583,7 @@ class AsyncOmniEngine:
                 request_address=request_address,
                 response_address=response_address,
             )
-            complete_diffusion_handshake(proc, handshake_address)
+            complete_diffusion_handshake(proc, handshake_address, stage_init_timeout)
             logger.info(
                 "[AsyncOmniEngine] Stage %s diffusion startup completed",
                 metadata.stage_id,
@@ -580,6 +599,9 @@ class AsyncOmniEngine:
             if proc is not None:
                 terminate_alive_proc(proc)
             raise
+        finally:
+            if lock_fds:
+                release_device_locks(lock_fds)
 
     def _create_remote_diffusion_stage(
         self,
@@ -785,6 +807,7 @@ class AsyncOmniEngine:
                                         stage_cfg,
                                         metadata,
                                         self._omni_master_server,
+                                        stage_init_timeout,
                                     )
                                 else:
                                     use_inline = True if self.num_stages == 1 else False
@@ -1445,8 +1468,6 @@ class AsyncOmniEngine:
                         cfg.engine_args.enable_sleep_mode = global_sleep_mode
                 if getattr(cfg, "stage_type", None) != "diffusion":
                     continue
-                if not hasattr(cfg, "engine_args") or cfg.engine_args is None:
-                    cfg.engine_args = OmegaConf.create({})
                 if kwargs.get("lora_path") is not None:
                     if not hasattr(cfg.engine_args, "lora_path") or cfg.engine_args.lora_path is None:
                         cfg.engine_args.lora_path = kwargs["lora_path"]

vllm_omni/engine/output_processor.py

@@ -128,9 +128,12 @@ class OmniRequestState(RequestState):
                 if isinstance(v, list) and v and isinstance(v[0], torch.Tensor):
                     try:
                         if k == "audio":
-                            # Audio chunks are usually 2D (i.e., 1, N); concatenate
-                            # on the last axis to preserve the channel dimension.
-                            self.mm_accumulated[k] = torch.cat(v, dim=-1)
+                            # Preserve channel dimension when chunks are compatible;
+                            # fall back to 1-D waveform concatenation for uneven chunks.
+                            try:
+                                self.mm_accumulated[k] = torch.cat(v, dim=-1)
+                            except RuntimeError:
+                                self.mm_accumulated[k] = torch.cat([t.reshape(-1) for t in v], dim=0)
                         elif k == "sr":
                             # Sample rate is a constant scalar, keep last value.
                             self.mm_accumulated[k] = v[-1]

vllm_omni/engine/stage_engine_startup.py

@@ -445,13 +445,22 @@ def _wait_for_omni_engine_startup(
             logger.debug("[omni] HELLO from engine %d", eng_index)
 
         elif status == "READY" and engine.state == CoreEngineState.CONNECTED:
-            num_gpu_blocks = (cache_config.num_gpu_blocks or 0) + msg["num_gpu_blocks"]
-            cache_config.num_gpu_blocks = num_gpu_blocks
+            # Upstream vllm >=0.19 dropped `num_gpu_blocks` from the READY
+            # handshake payload (the field is now communicated out-of-band via
+            # stats/health topics); tolerate both legacy and new message
+            # shapes so the omni handshake keeps working across rebases.
+            reported_blocks = msg.get("num_gpu_blocks")
+            if reported_blocks is not None:
+                cache_config.num_gpu_blocks = (cache_config.num_gpu_blocks or 0) + int(reported_blocks)
             if engine_addresses.frontend_stats_publish_address is None:
                 engine_addresses.frontend_stats_publish_address = msg.get("dp_stats_address")
             start_pending -= 1
             engine.state = CoreEngineState.READY
-            logger.debug("[omni] READY from engine %d (num_gpu_blocks=%d)", eng_index, msg["num_gpu_blocks"])
+            logger.debug(
+                "[omni] READY from engine %d (num_gpu_blocks=%s)",
+                eng_index,
+                "unknown" if reported_blocks is None else reported_blocks,
+            )
 
         else:
             raise RuntimeError(f"[omni] Unexpected status '{status}' from engine {eng_index} in state {engine.state}.")
@@ -462,7 +471,7 @@ def connect_remote_engine_cores(
     vllm_config: VllmConfig,
     omni_master_server: OmniMasterServer,
     stage_id: int,
-) -> Iterator[tuple[None, DPCoordinator | None, EngineZmqAddresses]]:
+) -> Iterator[tuple[None, DPCoordinator | None, EngineZmqAddresses, None]]:
     """Wait for remote engine cores to connect through the omni handshake."""
     addresses = omni_master_server.get_zmq_addresses(stage_id)
     parallel_config = vllm_config.parallel_config
@@ -491,7 +500,7 @@ def connect_remote_engine_cores(
     handshake_bind_address = omni_master_server.get_allocation(stage_id).handshake_bind_address
 
     with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket:
-        yield None, coordinator, addresses
+        yield None, coordinator, addresses, None
 
         _wait_for_omni_engine_startup(
             handshake_socket,

vllm_omni/engine/stage_init_utils.py

@@ -600,6 +600,32 @@ def release_device_locks(lock_fds: list[int]) -> None:
             pass
 
 
+def acquire_diffusion_device_locks(
+    stage_id: int,
+    od_config: Any,
+    stage_init_timeout: int,
+) -> list[int]:
+    """Acquire init locks for the GPU set used by a diffusion stage.
+
+    Diffusion stages express their device count via ``OmniDiffusionConfig``'s
+    ``parallel_config.world_size`` rather than the LLM-style
+    ``tensor_parallel_size`` knob, so adapt to the shape that
+    ``acquire_device_locks`` understands.
+    """
+    parallel_config = getattr(od_config, "parallel_config", None)
+    world_size = getattr(parallel_config, "world_size", 1)
+    try:
+        world_size = max(1, int(world_size))
+    except (TypeError, ValueError):
+        world_size = 1
+
+    return acquire_device_locks(
+        stage_id,
+        {"tensor_parallel_size": world_size},
+        stage_init_timeout,
+    )
+
+
 def load_omni_transfer_config_for_model(model: str, config_path: str | None) -> Any:
     """Load omni transfer config from an explicit path or resolved model config.
 

vllm_omni/entrypoints/cli/serve.py

@@ -718,7 +718,9 @@ def run_headless(args: argparse.Namespace) -> None:
             executor_class=executor_class,
             log_stats=log_stats,
         )
-        engine_manager.join_first()
+        # vllm>=0.19 renamed CoreEngineProcManager.join_first() to
+        # monitor_engine_liveness() (see upstream PR #35862).
+        engine_manager.monitor_engine_liveness()
     finally:
         logger.info("[Headless] Shutting down stage %d.", stage_id)
         if engine_manager is not None:

vllm_omni/entrypoints/openai/api_server.py

@@ -65,8 +65,8 @@ from vllm.entrypoints.openai.speech_to_text.serving import (
 from vllm.entrypoints.openai.utils import validate_json_request
 from vllm.entrypoints.pooling.classify.serving import ServingClassification
 from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding
-from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling
-from vllm.entrypoints.pooling.score.serving import ServingScores
+from vllm.entrypoints.pooling.pooling.serving import ServingPooling as OpenAIServingPooling
+from vllm.entrypoints.pooling.scoring.serving import ServingScores
 from vllm.entrypoints.serve.disagg.serving import ServingTokens
 
 # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic.
@@ -95,7 +95,11 @@ from vllm_omni.entrypoints.openai.image_api_utils import (
     parse_size,
     validate_layered_layers,
 )
-from vllm_omni.entrypoints.openai.protocol.audio import BatchSpeechRequest, OpenAICreateSpeechRequest
+from vllm_omni.entrypoints.openai.protocol.audio import (
+    BatchSpeechRequest,
+    OpenAICreateAudioGenerateRequest,
+    OpenAICreateSpeechRequest,
+)
 from vllm_omni.entrypoints.openai.protocol.images import (
     ImageData,
     ImageGenerationRequest,
@@ -112,6 +116,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoResponse,
 )
 from vllm_omni.entrypoints.openai.realtime_connection import RealtimeConnection
+from vllm_omni.entrypoints.openai.serving_audio_generate import OmniOpenAIServingAudioGenerate
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
@@ -313,6 +318,9 @@ class _DiffusionServingModels:
     def base_model_paths(self) -> list[BaseModelPath]:
         return self._base_model_paths
 
+    def is_base_model(self, model_name: str) -> bool:
+        return any(p.name == model_name for p in self._base_model_paths)
+
     def __getattr__(self, name):
         return self._Unsupported(name)
 
@@ -581,10 +589,10 @@ async def omni_init_app_state(
 
     # For omni models
     state.stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None
+    model_name = served_model_names[0] if served_model_names else args.model
 
     # Pure Diffusion mode: use simplified initialization logic
     if is_pure_diffusion:
-        model_name = served_model_names[0] if served_model_names else args.model
         state.vllm_config = None
         state.diffusion_engine = engine_client
         state.openai_serving_models = _DiffusionServingModels(base_model_paths)
@@ -597,6 +605,16 @@ async def omni_init_app_state(
             model_name=model_name,
         )
 
+        # audio related
+        state.openai_serving_speech = None
+        state.openai_serving_audio_generate = OmniOpenAIServingAudioGenerate.for_diffusion(
+            engine_client,
+            state.openai_serving_models,
+            request_logger=request_logger,
+            model_name=model_name,
+        )
+
+        # video related
         diffusion_stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None
         state.openai_serving_video = OmniOpenAIServingVideo.for_diffusion(
             diffusion_engine=engine_client,  # type: ignore
@@ -652,51 +670,35 @@ async def omni_init_app_state(
     )
     lora_modules = process_lora_modules(args.lora_modules, default_mm_loras)
 
-    # Ensure input_processor, io_processor, and model_config exist for OpenAIServingModels compatibility
+    # Ensure `input_processor` and `model_config` exist on the engine client
+    # for OpenAIServingModels compatibility.
+    #
+    # vLLM 0.20 dropped the `io_processor` kwarg from OpenAIServingRender and
+    # neither `vllm.entrypoints.openai.*` nor `vllm.entrypoints.serve.*` reads
+    # `engine_client.io_processor` anymore, so we no longer need to back-fill
+    # it here. AsyncOmni still sets `self.io_processor` in its own __init__
+    # for any vllm-omni internal callers that rely on it.
     if (
         not hasattr(engine_client, "input_processor")
         or engine_client.input_processor is None
-        or not hasattr(engine_client, "io_processor")
-        or engine_client.io_processor is None
         or not hasattr(engine_client, "model_config")
         or engine_client.model_config is None
     ):
         if vllm_config is not None:
-            # Try to initialize processors if vllm_config is available
             try:
-                from vllm.plugins.io_processors import get_io_processor
                 from vllm.v1.engine.input_processor import InputProcessor
 
                 tokenizer = await engine_client.get_tokenizer()
                 if tokenizer is not None:
-                    # Initialize input_processor
                     if not hasattr(engine_client, "input_processor") or engine_client.input_processor is None:
                         engine_client.input_processor = InputProcessor(
                             vllm_config=vllm_config,
                         )
                         logger.info("Initialized input_processor for AsyncOmni")
 
-                    # Initialize model_config
                     if not hasattr(engine_client, "model_config") or engine_client.model_config is None:
                         engine_client.model_config = vllm_config.model_config
                         logger.info("Initialized model_config for AsyncOmni")
-
-                    # Initialize io_processor
-                    if not hasattr(engine_client, "io_processor") or engine_client.io_processor is None:
-                        model_config = (
-                            engine_client.model_config
-                            if hasattr(engine_client, "model_config")
-                            else vllm_config.model_config
-                        )
-                        io_processor_plugin = model_config.io_processor_plugin
-                        renderer = getattr(engine_client, "renderer", None)
-                        if renderer is None:
-                            from vllm.renderers import renderer_from_config
-
-                            renderer = renderer_from_config(vllm_config)
-                            engine_client.renderer = renderer
-                        engine_client.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin)
-                        logger.info("Initialized io_processor for AsyncOmni")
                 else:
                     logger.warning("Cannot initialize processors: tokenizer is None. OpenAIServingModels may fail.")
             except Exception as e:
@@ -714,10 +716,15 @@ async def omni_init_app_state(
     )
     await state.openai_serving_models.init_static_loras()
 
+    # NOTE: kept aligned with vllm 0.20 `init_app_state`:
+    # - dropped the `io_processor` kwarg (no longer accepted by 0.20);
+    #   io_processor stays on `engine_client` and downstream serving classes
+    #   read it from there.
+    # - pass `reasoning_parser` so render-time `adjust_request` runs for
+    #   reasoning models (matches `vllm.entrypoints.openai.api_server`).
     state.openai_serving_render = OpenAIServingRender(
         model_config=engine_client.model_config,
         renderer=engine_client.renderer,
-        io_processor=engine_client.io_processor,
         model_registry=state.openai_serving_models.registry,
         request_logger=request_logger,
         chat_template=resolved_chat_template,
@@ -726,6 +733,7 @@ async def omni_init_app_state(
         enable_auto_tools=args.enable_auto_tool_choice,
         exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none,
         tool_parser=args.tool_call_parser,
+        reasoning_parser=args.structured_outputs_config.reasoning_parser,
         default_chat_template_kwargs=args.default_chat_template_kwargs,
         log_error_stack=args.log_error_stack,
     )
@@ -906,7 +914,11 @@ async def omni_init_app_state(
     )
 
     state.openai_serving_speech = OmniOpenAIServingSpeech(
-        engine_client, state.openai_serving_models, request_logger=request_logger
+        engine_client, state.openai_serving_models, request_logger=request_logger, model_name=model_name
+    )
+
+    state.openai_serving_audio_generate = OmniOpenAIServingAudioGenerate(
+        engine_client, state.openai_serving_models, request_logger=request_logger, model_name=model_name
     )
 
     state.openai_streaming_speech = OmniStreamingSpeechHandler(
@@ -949,6 +961,10 @@ def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None:
     return request.app.state.openai_serving_speech
 
 
+def OmniAudioGenerate(request: Request) -> OmniOpenAIServingAudioGenerate | None:
+    return getattr(request.app.state, "openai_serving_audio_generate", None)
+
+
 @router.post(
     "/v1/chat/completions",
     dependencies=[Depends(validate_json_request)],
@@ -1115,6 +1131,40 @@ async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request)
         raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
 
 
+@router.post(
+    "/v1/audio/generate",
+    dependencies=[Depends(validate_json_request)],
+    responses={
+        HTTPStatus.OK.value: {"content": {"audio/*": {}}},
+        HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
+        HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
+        HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
+    },
+)
+@with_cancellation
+@load_aware_call
+async def create_audio_generate(request: OpenAICreateAudioGenerateRequest, raw_request: Request):
+    handler = OmniAudioGenerate(raw_request)
+    if handler is None:
+        base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None)
+        if base_server is None:
+            raise HTTPException(
+                status_code=HTTPStatus.NOT_FOUND.value,
+                detail="The model does not support Audio Generate API",
+            )
+        return base_server.create_error_response(message="The model does not support Audio Generate API")
+    try:
+        result = await handler.create_audio_generate(request, raw_request)
+        if isinstance(result, ErrorResponse):
+            return JSONResponse(
+                content=result.model_dump(),
+                status_code=result.error.code if result.error else 400,
+            )
+        return result
+    except Exception as e:
+        raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
+
+
 @router.get(
     "/v1/audio/voices",
     responses={

vllm_omni/entrypoints/openai/audio_utils_mixin.py

@@ -42,6 +42,10 @@ class AudioMixin:
                 "Only mono (1D) and stereo (2D) are supported."
             )
 
+        if audio_tensor.ndim == 2 and audio_tensor.shape[0] == 2:
+            # Convert from [channels, samples] to [samples, channels]
+            audio_tensor = audio_tensor.T
+
         audio_tensor, sample_rate = self._apply_speed_adjustment(audio_tensor, speed, sample_rate)
 
         supported_formats = {

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -83,12 +83,7 @@ class OpenAICreateSpeechRequest(BaseModel):
     )
     extra_params: dict[str, Any] | None = Field(
         default=None,
-        description=(
-            "Optional model-specific parameters passed directly to the model's extra_args. "
-            "VoxCPM2 reads `cfg_value` (float in 0.1-10.0; default 2.0) from this dict to set "
-            "the classifier-free guidance scale per request; Voxtral TTS reads `cfg_alpha` here "
-            "the same way."
-        ),
+        description=("Optional model-specific parameters passed directly to the model's extra_args."),
     )
 
     @field_validator("stream_format")
@@ -129,6 +124,53 @@ class OpenAICreateSpeechRequest(BaseModel):
         return self
 
 
+class OpenAICreateAudioGenerateRequest(BaseModel):
+    """Request model for audio generation via diffusion models (e.g. Stable Audio)."""
+
+    input: str = Field(
+        description="Text prompt describing the audio to generate",
+    )
+    model: str | None = None
+    response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav"
+    speed: float | None = Field(
+        default=1.0,
+        ge=0.25,
+        le=4.0,
+    )
+    stream_format: Literal["sse", "audio"] | None = "audio"
+    audio_length: float | None = Field(
+        default=None,
+        description="Audio length in seconds",
+    )
+    audio_start: float | None = Field(
+        default=0.0,
+        description="Audio start time in seconds",
+    )
+    negative_prompt: str | None = Field(
+        default=None,
+        description="Negative prompt for classifier-free guidance",
+    )
+    guidance_scale: float | None = Field(
+        default=None,
+        description="Guidance scale for diffusion models",
+    )
+    num_inference_steps: int | None = Field(
+        default=None,
+        description="Number of inference steps",
+    )
+    seed: int | None = Field(
+        default=None,
+        description="Random seed for reproducibility",
+    )
+
+    @field_validator("stream_format")
+    @classmethod
+    def validate_stream_format(cls, v: str) -> str:
+        if v == "sse":
+            raise ValueError("'sse' is not a supported stream_format yet. Please use 'audio'.")
+        return v
+
+
 class CreateAudio(BaseModel):
     audio_tensor: np.ndarray
     sample_rate: int = 24000

vllm_omni/entrypoints/openai/serving_audio_generate.py

@@ -0,0 +1,168 @@
+import asyncio
+
+import torch
+from fastapi import Request
+from fastapi.responses import Response
+from vllm.entrypoints.openai.engine.serving import OpenAIServing
+from vllm.logger import init_logger
+from vllm.utils import random_uuid
+
+from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
+from vllm_omni.entrypoints.openai.protocol.audio import (
+    AudioResponse,
+    CreateAudio,
+    OpenAICreateAudioGenerateRequest,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+logger = init_logger(__name__)
+
+
+class OmniOpenAIServingAudioGenerate(OpenAIServing, AudioMixin):
+    """Serving class for audio generation via diffusion models (e.g. Stable Audio)."""
+
+    def __init__(self, *args, **kwargs):
+        self.model_name = kwargs.pop("model_name", None)
+        super().__init__(*args, **kwargs)
+        self.diffusion_mode = False
+
+    @classmethod
+    def for_diffusion(cls, *args, **kwargs) -> "OmniOpenAIServingAudioGenerate":
+        """Create an instance configured to run in diffusion mode."""
+        instance = cls(*args, **kwargs)
+        instance.diffusion_mode = True
+        return instance
+
+    def _is_stable_audio_model(self) -> bool:
+        return self.engine_client.model_type == "StableAudioPipeline"
+
+    async def create_audio_generate(
+        self,
+        request: OpenAICreateAudioGenerateRequest,
+        raw_request: Request | None = None,
+    ):
+        """
+        Generate audio using diffusion-based models (e.g. Stable Audio).
+
+        This endpoint is designed for audio generation models as
+        opposed to TTS models that specifically generate speech.
+        """
+
+        error_check_ret = await self._check_model(request)
+        if error_check_ret is not None:
+            logger.error("Error with model %s", error_check_ret)
+            return error_check_ret
+
+        if self.engine_client.errored:
+            raise self.engine_client.dead_error
+
+        request_id = f"audiogen-{random_uuid()}"
+
+        try:
+            default_sr = 44100  # Default sample rate for Stable Audio
+
+            # Build prompt for diffusion audio generation
+            prompt = {
+                "prompt": request.input,
+            }
+            if request.negative_prompt:
+                prompt["negative_prompt"] = request.negative_prompt
+
+            # Build sampling params for diffusion
+            sampling_params_list = [OmniDiffusionSamplingParams(num_outputs_per_prompt=1)]
+
+            # Create generator if seed provided
+            if request.seed is not None:
+                from vllm_omni.platforms import current_omni_platform
+
+                rng = torch.Generator(device=current_omni_platform.device_type).manual_seed(request.seed)
+                sampling_params_list[0].generator = rng
+
+            if request.guidance_scale is not None:
+                sampling_params_list[0].guidance_scale = request.guidance_scale
+
+            if request.num_inference_steps is not None:
+                sampling_params_list[0].num_inference_steps = request.num_inference_steps
+
+            # Set up audio duration parameters
+            if request.audio_length is not None:
+                audio_length = request.audio_length
+                audio_start = request.audio_start if request.audio_start is not None else 0.0
+                audio_end_in_s = audio_start + audio_length
+                sampling_params_list[0].extra_args = {
+                    "audio_start_in_s": audio_start,
+                    "audio_end_in_s": audio_end_in_s,
+                }
+
+            logger.info(
+                "Audio generation request %s: prompt=%r",
+                request_id,
+                request.input[:50] + "..." if len(request.input) > 50 else request.input,
+            )
+
+            generator = self.engine_client.generate(
+                prompt=prompt,
+                request_id=request_id,
+                sampling_params_list=sampling_params_list,
+                output_modalities=["audio"],
+            )
+
+            final_output: OmniRequestOutput | None = None
+            async for res in generator:
+                final_output = res
+
+            if final_output is None:
+                return self.create_error_response("No output generated from the model.")
+
+            # Extract audio from output
+            audio_output = None
+            if hasattr(final_output, "multimodal_output") and final_output.multimodal_output:
+                audio_output = final_output.multimodal_output
+            if not audio_output and hasattr(final_output, "request_output"):
+                if final_output.request_output and hasattr(final_output.request_output, "multimodal_output"):
+                    audio_output = final_output.request_output.multimodal_output
+
+            # Check for audio data using either "audio" or "model_outputs" key
+            audio_key = None
+            if audio_output:
+                if "audio" in audio_output:
+                    audio_key = "audio"
+                elif "model_outputs" in audio_output:
+                    audio_key = "model_outputs"
+
+            if not audio_output or audio_key is None:
+                return self.create_error_response("Audio generation model did not produce audio output.")
+
+            audio_tensor = audio_output[audio_key]
+            sample_rate = audio_output.get("sr", default_sr)
+            if hasattr(sample_rate, "item"):
+                sample_rate = sample_rate.item()
+
+            # Convert tensor to numpy
+            if hasattr(audio_tensor, "float"):
+                audio_tensor = audio_tensor.float().detach().cpu().numpy()
+
+            # Squeeze batch dimension if present, but preserve channel dimension for stereo
+            if audio_tensor.ndim > 1:
+                audio_tensor = audio_tensor.squeeze()
+
+            audio_obj = CreateAudio(
+                audio_tensor=audio_tensor,
+                sample_rate=int(sample_rate),
+                response_format=request.response_format or "wav",
+                speed=request.speed or 1.0,
+                stream_format=request.stream_format,
+                base64_encode=False,
+            )
+
+            audio_response: AudioResponse = self.create_audio(audio_obj)
+            return Response(content=audio_response.audio_data, media_type=audio_response.media_type)
+
+        except asyncio.CancelledError:
+            return self.create_error_response("Client disconnected")
+        except ValueError as e:
+            return self.create_error_response(e)
+        except Exception as e:
+            logger.exception("Audio generation failed: %s", e)
+            return self.create_error_response(f"Audio generation failed: {e}")

vllm_omni/entrypoints/openai/serving_speech.py

@@ -197,6 +197,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         return instance
 
     def __init__(self, *args, **kwargs):
+        self.model_name = kwargs.pop("model_name", None)
         super().__init__(*args, **kwargs)
         # Initialize uploaded speakers storage (ephemeral — cleared on restart)
         speech_voice_samples_dir = os.environ.get("SPEECH_VOICE_SAMPLES", "/tmp/voice_samples")
@@ -492,39 +493,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         ref_sr = None
         if request.ref_audio is not None:
             ref_audio, ref_sr = await self._resolve_ref_audio(request.ref_audio)
-        # VoxCPM2 uses a text-prefix convention for style/voice instructions:
-        # the native CLI wraps the instruction in parens and prepends it.
-        # Mirror that here so `instructions` works through the OpenAI endpoint.
-        # The Hi-Fi Cloning mode (ref_audio + ref_text) + instructions combo
-        # is rejected upstream in `_validate_tts_request`, so by the time we
-        # reach this builder we can assume the combination is valid.
-        text = request.input
-        if request.instructions:
-            text = f"({request.instructions.strip()}){text}"
-        prompt = build_voxcpm2_prompt(
+        return build_voxcpm2_prompt(
             hf_config=self.engine_client.model_config.hf_config,
             tokenizer=self._voxcpm2_tokenizer,
             split_map=self._voxcpm2_split_map,
-            text=text,
+            text=request.input,
             ref_audio=ref_audio,
             ref_sr=ref_sr,
             ref_text=request.ref_text,
         )
-        # Per-request cfg_value threads via additional_information; the talker
-        # lifts it onto _RequestState in preprocess() and uses it in _run_cfm.
-        # Following the convention from #2338 (Voxtral TTS / cfg_alpha), the
-        # value is provided via the shared `extra_params` dict rather than as
-        # a top-level protocol field, keeping OpenAICreateSpeechRequest free
-        # of model-specific knobs.
-        if request.extra_params and "cfg_value" in request.extra_params:
-            try:
-                cfg_value = float(request.extra_params["cfg_value"])
-            except (TypeError, ValueError) as e:
-                raise ValueError(f"extra_params['cfg_value'] must be a number: {e}") from e
-            if not 0.1 <= cfg_value <= 10.0:
-                raise ValueError(f"extra_params['cfg_value']={cfg_value} out of range (0.1-10.0)")
-            prompt.setdefault("additional_information", {})["cfg_value"] = cfg_value
-        return prompt
 
     def _get_uploaded_audio_data(self, voice_name: str) -> str | None:
         """Get base64 encoded audio data for uploaded voice."""
@@ -880,27 +857,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         if self._tts_model_type == "voxcpm":
             return self._validate_voxcpm_request(request)
         if self._tts_model_type == "voxcpm2":
-            # Hi-Fi Cloning mode (`ref_audio` + `ref_text`) does not use the
-            # control instruction. Upstream voxcpm CLI errors out on the same
-            # combination (src/voxcpm/cli.py:128-131: --control cannot be used
-            # together with --prompt-text); we surface a 400 to match instead
-            # of silently dropping user input.
-            if request.instructions and request.ref_text:
-                return (
-                    "`instructions` is incompatible with `ref_text` in Hi-Fi Cloning "
-                    "mode (the control instruction is ignored when ref_text is set)."
-                )
-            # Cap `instructions` length for parity with qwen3_tts / voxtral /
-            # fish_tts (prevents blowing the prefill budget via an unbounded
-            # parenthetical prefix). Upstream voxcpm has no cap; this is a
-            # serving-side safety net, see PR #3118 discussion.
-            if request.instructions and len(request.instructions) > self._max_instructions_length:
-                return (
-                    f"Instructions too long ({len(request.instructions)}/"
-                    f"{self._max_instructions_length} chars). Upstream voxcpm has no cap; "
-                    f"this limit prevents exceeding the prefill budget."
-                )
-            return None
+            return None  # VoxCPM2 accepts any text input
         if self._tts_model_type == "ming_flash_omni_tts":
             return self._validate_ming_tts_request(request)
         if self._tts_model_type == "moss_tts_nano":
@@ -1761,12 +1718,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             if request.instructions:
                 prompt["instruct"] = request.instructions
         elif self._tts_model_type == "voxcpm2":
-            # Run the same validator the batch endpoint uses so single-request
-            # /v1/audio/speech enforces the instructions-length cap and any
-            # other voxcpm2-specific checks (e.g. Hi-Fi mode + instructions).
-            validation_error = self._validate_tts_request(request)
-            if validation_error:
-                raise ValueError(validation_error)
             prompt = await self._build_voxcpm2_prompt(request)
             tts_params = {}
         elif self._is_tts:
@@ -2000,6 +1951,14 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
             if not request.input or not request.input.strip():
                 raise ValueError("Input text cannot be empty")
 
+            # Validate ref_audio format up-front so that bogus inputs return a
+            # 4xx instead of falling through to MediaConnector and surfacing as
+            # a 500 Internal Server Error (e.g. test_voice_clone_invalid_ref_audio_format).
+            if request.ref_audio is not None:
+                fmt_err = self._validate_ref_audio_format(request.ref_audio)
+                if fmt_err:
+                    return self._diffusion_error_response(fmt_err, status_code=400)
+
             request_id = f"speech-{random_uuid()}"
             prompt: dict[str, Any] = {"input": request.input}
             if request.ref_audio:
@@ -2066,16 +2025,28 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         except (EngineGenerateError, EngineDeadError):
             raise  # Propagate to the global Omni exception handler
         except ValueError as e:
-            return self._diffusion_error_response(str(e))
+            # ValueError represents invalid client input (bad ref_audio URI,
+            # empty text, model-side validation failures forwarded as ValueError, ...).
+            # Return 400 to match `create_speech`'s `self.create_error_response(e)`
+            # default (HTTPStatus.BAD_REQUEST) rather than masking it as a 500.
+            return self._diffusion_error_response(str(e), status_code=400)
         except Exception as e:
             logger.exception("Diffusion speech generation failed: %s", e)
             return self._diffusion_error_response(f"Speech generation failed: {e}")
 
     @staticmethod
-    def _diffusion_error_response(message: str) -> Response:
-        """Create a JSON error response without depending on OpenAIServing."""
-        error_body = json.dumps({"error": {"message": message, "type": "server_error", "param": None, "code": 500}})
-        return Response(content=error_body, media_type="application/json", status_code=500)
+    def _diffusion_error_response(message: str, status_code: int = 500) -> Response:
+        """Create a JSON error response without depending on OpenAIServing.
+
+        Args:
+            message: Error message to surface to the client.
+            status_code: HTTP status code; defaults to 500. Pass a 4xx code for
+                client-input validation failures so the response semantics match
+                the OpenAI-compatible behavior used by ``create_speech``.
+        """
+        err_type = "BadRequestError" if 400 <= status_code < 500 else "server_error"
+        error_body = json.dumps({"error": {"message": message, "type": err_type, "param": None, "code": status_code}})
+        return Response(content=error_body, media_type="application/json", status_code=status_code)
 
     async def create_speech(
         self,

vllm_omni/inputs/data.py

@@ -16,7 +16,8 @@ except ImportError:
 
 
 import torch
-from vllm.inputs import EmbedsPrompt, TextPrompt, TokensInput, TokensPrompt
+from vllm.inputs import EmbedsPrompt, TextPrompt, TokensPrompt
+from vllm.inputs.engine import TokensInput
 
 
 class OmniTextPrompt(TextPrompt):

vllm_omni/model_executor/models/bagel/bagel.py

@@ -53,19 +53,34 @@ from vllm_omni.diffusion.models.bagel.pipeline_bagel import default_ae_params
 
 
 class OmniBagelProcessor(BagelProcessor):
+    # transformers>=5.0 ProcessorMixin.get_attributes() only scans the leaf
+    # class's __dict__ for ``<attribute>_class`` hints; redeclare them here
+    # so from_pretrained() correctly sets ``self.image_processor`` and
+    # ``self.tokenizer`` on the OmniBagelProcessor instance.
+    image_processor_class = "SiglipImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
     def __call__(self, text=None, images=None, **kwargs):
         is_img2img = kwargs.pop("is_img2img", False)
 
         if is_img2img and images is not None:
-            image_kwargs = kwargs.copy()
+            # transformers>=5.0 enforces strict kwarg typing on image
+            # processors, so split generic kwargs into text/image buckets
+            # via the standard ProcessorMixin helper before dispatch.
+            from vllm.transformers_utils.processors.bagel import BagelProcessorKwargs
+
+            output_kwargs = self._merge_kwargs(
+                BagelProcessorKwargs,
+                tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+                **kwargs,
+            )
+            image_kwargs = dict(output_kwargs["images_kwargs"])
             image_kwargs["do_resize"] = False
             image_kwargs["do_rescale"] = True
-            if "return_tensors" not in image_kwargs:
-                image_kwargs["return_tensors"] = "pt"
-
+            image_kwargs.setdefault("return_tensors", "pt")
             pixel_values = self.image_processor(images, **image_kwargs)
 
-            text_inputs = self.tokenizer(text, **kwargs) if text is not None else None
+            text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) if text is not None else None
 
             if pixel_values is not None and text_inputs is not None:
                 combined = dict(text_inputs)