[Perf] Reduce IPC overhead for single-stage diffusion serving for Wan2.2 (#1715)

Signed-off-by: samithuang <285365963@qq.com>
Signed-off-by: Samit <285365963@qq.com>

.buildkite/release-pipeline.yaml

@@ -0,0 +1,38 @@
+steps:
+  - label: ":docker: Build image"
+    key: image-build
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L2 Test
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-ready.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L3 Test
+  - label: "Upload Merge Pipeline"
+    depends_on: image-build
+    key: upload-merge-pipeline
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-merge.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L4 Test
+  - label: "Upload Nightly Pipeline"
+    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/test-nightly.yml

@@ -1,5 +1,5 @@
 steps:
-  - label: "Omni Model Test with H100"
+  - label: ":full_moon: Omni Model Test with H100"
     timeout_in_minutes: 90
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1"
@@ -41,7 +41,7 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: "Omni Model Test"
+  - label: ":full_moon: Omni Model Test"
     timeout_in_minutes: 60
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1"
@@ -56,18 +56,44 @@ steps:
           always-pull: true
           shm-size: "8gb"
           propagate-environment: true
+          shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Omni Model Perf Test"
-    timeout_in_minutes: 120
+  - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
+    timeout_in_minutes: 30
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - |
+        huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
+    agents:
+      queue: "gpu_4_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: ":full_moon: Omni Model Perf Test with H100"
+    key: nightly-performance
+    timeout_in_minutes: 180
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export BENCHMARK_DIR=tests
       - pytest -s -v tests/perf/scripts/run_benchmark.py
+      - buildkite-agent artifact upload "tests/*.json"
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -96,3 +122,18 @@ steps:
                 hostPath:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
+
+  - label: ":email: Nightly Perf Collection & Email"
+    key: nightly-perf-distribution
+    depends_on: nightly-performance
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - pip install openpyxl
+      - export DEFAULT_INPUT_DIR=tests
+      - export DEFAULT_OUTPUT_DIR=tests
+      - buildkite-agent artifact download "tests/*.json" . --step nightly-performance
+      - python tools/nightly/generate_nightly_perf_excel.py
+      - python tools/nightly/send_nightly_perf_email.py
+      - buildkite-agent artifact upload "tests/*.xlsx"
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/test-ready.yml

@@ -214,9 +214,10 @@ steps:
     commands:
       - |
         timeout 20m bash -c '
+          huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
           export VLLM_LOGGING_LEVEL=DEBUG
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
         '
     agents:
       queue: "gpu_4_queue"

README.md

@@ -1,6 +1,6 @@
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
+    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
     <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
   </picture>
 </p>
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://drive.google.com/file/d/12k_DOo2xoEROvd4xQxYNuEi6C2GECoRx/view?usp=drive_link"><b>Slides</b></a> |
 </p>
 
 
@@ -17,6 +17,8 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 *Latest News* 🔥
 
+- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
+- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
 - [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.

benchmarks/qwen3-tts/vllm_omni/bench_async_chunk.py

@@ -0,0 +1,301 @@
+"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
+
+Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
+concurrency levels for both async_chunk modes. Saves results as JSON.
+
+Usage:
+    # Run against a server already serving with a given config:
+    python bench_async_chunk.py \
+        --host 127.0.0.1 --port 8000 \
+        --config-name async_chunk_on \
+        --num-prompts 50 \
+        --max-concurrency 1 10 \
+        --result-dir results/
+"""
+
+import argparse
+import asyncio
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+from tqdm.asyncio import tqdm
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
+    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
+]
+
+
+@dataclass
+class RequestResult:
+    success: bool = False
+    ttfp: float = 0.0
+    e2e: float = 0.0
+    audio_bytes: int = 0
+    audio_duration: float = 0.0
+    rtf: float = 0.0
+    prompt: str = ""
+    error: str = ""
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 0
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0
+    request_throughput: float = 0.0
+    per_request: list = field(default_factory=list)
+
+
+def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
+    return num_bytes / sample_width / sample_rate
+
+
+async def send_tts_request(
+    session: aiohttp.ClientSession,
+    api_url: str,
+    prompt: str,
+    voice: str = "vivian",
+    language: str = "English",
+    stream: bool = True,
+    pbar: tqdm | None = None,
+) -> RequestResult:
+    payload = {
+        "input": prompt,
+        "voice": voice,
+        "language": language,
+        "stream": stream,
+        "response_format": "pcm",
+    }
+
+    result = RequestResult(prompt=prompt)
+    st = time.perf_counter()
+
+    try:
+        async with session.post(api_url, json=payload) as response:
+            if response.status != 200:
+                result.error = f"HTTP {response.status}: {await response.text()}"
+                return result
+
+            first_chunk = True
+            total_bytes = 0
+
+            async for chunk in response.content.iter_any():
+                if first_chunk and len(chunk) > 0:
+                    result.ttfp = time.perf_counter() - st
+                    first_chunk = False
+                total_bytes += len(chunk)
+
+            result.e2e = time.perf_counter() - st
+            result.audio_bytes = total_bytes
+            result.audio_duration = pcm_bytes_to_duration(total_bytes)
+            if result.audio_duration > 0:
+                result.rtf = result.e2e / result.audio_duration
+            result.success = True
+
+    except Exception as e:
+        result.error = str(e)
+        result.e2e = time.perf_counter() - st
+
+    if pbar:
+        pbar.update(1)
+    return result
+
+
+async def run_benchmark(
+    host: str,
+    port: int,
+    num_prompts: int,
+    max_concurrency: int,
+    num_warmups: int = 3,
+    voice: str = "vivian",
+    language: str = "English",
+    stream: bool = True,
+) -> BenchmarkResult:
+    api_url = f"http://{host}:{port}/v1/audio/speech"
+
+    connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
+    session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
+
+    if num_warmups > 0:
+        print(f"  Warming up with {num_warmups} requests...")
+        warmup_tasks = [
+            send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
+            for i in range(num_warmups)
+        ]
+        await asyncio.gather(*warmup_tasks)
+        print("  Warmup done.")
+
+    request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
+
+    print(f"  Running {num_prompts} requests with concurrency={max_concurrency}...")
+    semaphore = asyncio.Semaphore(max_concurrency)
+    pbar = tqdm(total=num_prompts, desc=f"  concurrency={max_concurrency}")
+
+    async def limited_request(prompt):
+        async with semaphore:
+            return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
+
+    start_time = time.perf_counter()
+    tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
+    results: list[RequestResult] = await asyncio.gather(*tasks)
+    duration = time.perf_counter() - start_time
+    pbar.close()
+
+    await session.close()
+
+    successful = [r for r in results if r.success]
+    failed = [r for r in results if not r.success]
+
+    bench = BenchmarkResult(
+        concurrency=max_concurrency,
+        num_prompts=num_prompts,
+        completed=len(successful),
+        failed=len(failed),
+        duration_s=duration,
+    )
+
+    if successful:
+        ttfps = [r.ttfp * 1000 for r in successful]
+        e2es = [r.e2e * 1000 for r in successful]
+        rtfs = [r.rtf for r in successful]
+        audio_durs = [r.audio_duration for r in successful]
+
+        bench.mean_ttfp_ms = float(np.mean(ttfps))
+        bench.median_ttfp_ms = float(np.median(ttfps))
+        bench.std_ttfp_ms = float(np.std(ttfps))
+        bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
+        bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
+        bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
+
+        bench.mean_e2e_ms = float(np.mean(e2es))
+        bench.median_e2e_ms = float(np.median(e2es))
+        bench.std_e2e_ms = float(np.std(e2es))
+        bench.p90_e2e_ms = float(np.percentile(e2es, 90))
+        bench.p95_e2e_ms = float(np.percentile(e2es, 95))
+        bench.p99_e2e_ms = float(np.percentile(e2es, 99))
+
+        bench.mean_rtf = float(np.mean(rtfs))
+        bench.median_rtf = float(np.median(rtfs))
+        bench.std_rtf = float(np.std(rtfs))
+
+        bench.mean_audio_duration_s = float(np.mean(audio_durs))
+        bench.total_audio_duration_s = float(np.sum(audio_durs))
+        bench.audio_throughput = bench.total_audio_duration_s / duration
+        bench.request_throughput = len(successful) / duration
+
+        bench.per_request = [
+            {
+                "ttfp_ms": r.ttfp * 1000,
+                "e2e_ms": r.e2e * 1000,
+                "rtf": r.rtf,
+                "audio_duration_s": r.audio_duration,
+                "prompt": r.prompt,
+            }
+            for r in successful
+        ]
+
+    print(f"\n{'=' * 60}")
+    print(f"  Concurrency: {max_concurrency}  |  Completed: {bench.completed}  |  Failed: {bench.failed}")
+    print(f"  Duration: {duration:.2f}s  |  Throughput: {bench.request_throughput:.2f} req/s")
+    print(
+        f"  TTFP (ms):  mean={bench.mean_ttfp_ms:.1f}  median={bench.median_ttfp_ms:.1f}"
+        f"  p90={bench.p90_ttfp_ms:.1f}  p99={bench.p99_ttfp_ms:.1f}"
+    )
+    print(
+        f"  E2E (ms):   mean={bench.mean_e2e_ms:.1f}  median={bench.median_e2e_ms:.1f}"
+        f"  p90={bench.p90_e2e_ms:.1f}  p99={bench.p99_e2e_ms:.1f}"
+    )
+    print(f"  RTF:        mean={bench.mean_rtf:.3f}  median={bench.median_rtf:.3f}")
+    print(f"  Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
+    print(f"{'=' * 60}\n")
+
+    if failed:
+        for r in failed[:3]:
+            print(f"  [ERROR] {r.error[:200]}")
+
+    return bench
+
+
+async def main(args):
+    all_results = []
+
+    for concurrency in args.max_concurrency:
+        result = await run_benchmark(
+            host=args.host,
+            port=args.port,
+            num_prompts=args.num_prompts,
+            max_concurrency=concurrency,
+            num_warmups=args.num_warmups,
+            voice=args.voice,
+            language=args.language,
+            stream=args.stream,
+        )
+        result.config_name = args.config_name
+        all_results.append(asdict(result))
+
+    result_dir = Path(args.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump(all_results, f, indent=2)
+    print(f"Results saved to {result_file}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--num-prompts", type=int, default=50)
+    parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
+    parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--voice", type=str, default="vivian")
+    parser.add_argument("--language", type=str, default="English")
+    parser.add_argument("--stream", action="store_true", default=True)
+    parser.add_argument("--no-stream", dest="stream", action="store_false")
+    parser.add_argument("--config-name", type=str, default="async_chunk_on")
+    parser.add_argument("--result-dir", type=str, default="results")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    asyncio.run(main(args))

benchmarks/qwen3-tts/vllm_omni/plot_async_chunk.py

@@ -0,0 +1,249 @@
+"""Plot TTFP comparison: async_chunk off vs on.
+
+Generates a bar chart with improvement arrows, matching the Qwen3-Omni
+async_chunk benchmark figure style.
+
+Usage:
+    python plot_async_chunk.py \
+        --off results/bench_async_chunk_off_*.json \
+        --on results/bench_async_chunk_on_*.json \
+        --output results/qwen3_tts_async_chunk_ttfp.png
+
+    # Also supports E2E and RTF metrics:
+    python plot_async_chunk.py \
+        --off results/bench_async_chunk_off_*.json \
+        --on results/bench_async_chunk_on_*.json \
+        --metric e2e \
+        --output results/qwen3_tts_async_chunk_e2e.png
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+METRIC_CONFIG = {
+    "ttfp": {
+        "key": "mean_ttfp_ms",
+        "ylabel": "TTFP (s)",
+        "title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
+        "to_seconds": True,
+    },
+    "e2e": {
+        "key": "mean_e2e_ms",
+        "ylabel": "E2E (s)",
+        "title": "E2E Latency - Qwen3-TTS, by concurrency",
+        "to_seconds": True,
+    },
+    "rtf": {
+        "key": "mean_rtf",
+        "ylabel": "RTF",
+        "title": "Real-Time Factor - Qwen3-TTS, by concurrency",
+        "to_seconds": False,
+    },
+}
+
+
+def load_results(path: str) -> list[dict]:
+    with open(path) as f:
+        return json.load(f)
+
+
+def plot_ttfp_comparison(
+    off_results: list[dict],
+    on_results: list[dict],
+    metric: str,
+    output_path: str,
+    title_override: str | None = None,
+):
+    cfg = METRIC_CONFIG[metric]
+    key = cfg["key"]
+    to_seconds = cfg["to_seconds"]
+
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    off_vals = []
+    on_vals = []
+    for c in concurrencies:
+        v_off = off_map[c][key]
+        v_on = on_map[c][key]
+        if to_seconds:
+            v_off /= 1000.0
+            v_on /= 1000.0
+        off_vals.append(v_off)
+        on_vals.append(v_on)
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+
+    x = np.arange(len(concurrencies))
+    width = 0.3
+
+    ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
+    ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
+
+    # Draw improvement arrows and labels
+    for i in range(len(concurrencies)):
+        v_off = off_vals[i]
+        v_on = on_vals[i]
+        if v_on > 0:
+            improvement = v_off / v_on
+        else:
+            improvement = float("inf")
+
+        # Arrow from top of off-bar to top of on-bar
+        arrow_start_x = x[i] - width / 2
+        arrow_start_y = v_off * 0.95
+        arrow_end_x = x[i] + width / 2
+        arrow_end_y = v_on * 1.05
+
+        ax.annotate(
+            "",
+            xy=(arrow_end_x, arrow_end_y),
+            xytext=(arrow_start_x, arrow_start_y),
+            arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
+        )
+
+        # Improvement label
+        label_x = (arrow_start_x + arrow_end_x) / 2
+        label_y = arrow_start_y + (v_off - v_on) * 0.15
+        ax.text(
+            label_x,
+            label_y,
+            f"{improvement:.1f}x improvement",
+            ha="center",
+            va="bottom",
+            fontsize=10,
+            color="red",
+            fontweight="bold",
+        )
+
+    title = title_override or cfg["title"]
+    ax.set_title(title, fontsize=13, fontweight="bold")
+    ax.set_ylabel(cfg["ylabel"], fontsize=12)
+    ax.set_xlabel("Max concurrency", fontsize=12)
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_yscale("log")
+    ax.legend(loc="upper left", fontsize=11)
+    ax.grid(axis="y", alpha=0.3, linestyle="--")
+    ax.set_axisbelow(True)
+
+    plt.tight_layout()
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
+    """Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+    fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
+
+    for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
+        cfg = METRIC_CONFIG[metric]
+        key = cfg["key"]
+        to_seconds = cfg["to_seconds"]
+
+        off_vals = []
+        on_vals = []
+        for c in concurrencies:
+            v_off = off_map[c][key]
+            v_on = on_map[c][key]
+            if to_seconds:
+                v_off /= 1000.0
+                v_on /= 1000.0
+            off_vals.append(v_off)
+            on_vals.append(v_on)
+
+        x = np.arange(len(concurrencies))
+        width = 0.3
+        ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
+        ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
+
+        for i in range(len(concurrencies)):
+            if on_vals[i] > 0:
+                improvement = off_vals[i] / on_vals[i]
+                ax.annotate(
+                    "",
+                    xy=(x[i] + width / 2, on_vals[i] * 1.05),
+                    xytext=(x[i] - width / 2, off_vals[i] * 0.95),
+                    arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
+                )
+                label_y = off_vals[i] * 0.85
+                ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
+
+        ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
+        ax.set_ylabel(cfg["ylabel"], fontsize=11)
+        ax.set_xlabel("Max concurrency", fontsize=11)
+        ax.set_xticks(x)
+        ax.set_xticklabels([str(c) for c in concurrencies])
+        if metric != "rtf":
+            ax.set_yscale("log")
+        ax.legend(fontsize=9)
+        ax.grid(axis="y", alpha=0.3, linestyle="--")
+        ax.set_axisbelow(True)
+
+    plt.tight_layout()
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def print_table(off_results: list[dict], on_results: list[dict]):
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    print("\n## Benchmark Results: async_chunk off vs on\n")
+    print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
+    print("| --- | --- | --- | --- | --- |")
+
+    for name, key, fmt in [
+        ("TTFP (ms)", "mean_ttfp_ms", ".1f"),
+        ("E2E (ms)", "mean_e2e_ms", ".1f"),
+        ("RTF", "mean_rtf", ".3f"),
+        ("Throughput", "audio_throughput", ".2f"),
+    ]:
+        for c in concurrencies:
+            v_off = off_map[c].get(key, 0)
+            v_on = on_map[c].get(key, 0)
+            if v_on > 0 and key != "audio_throughput":
+                ratio = f"{v_off / v_on:.1f}x"
+            elif v_off > 0 and key == "audio_throughput":
+                ratio = f"{v_on / v_off:.1f}x"
+            else:
+                ratio = "N/A"
+            print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
+    parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
+    parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
+    parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
+    parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
+    parser.add_argument("--title", type=str, default=None, help="Custom title override")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    off_results = load_results(args.off)
+    on_results = load_results(args.on)
+
+    print_table(off_results, on_results)
+
+    if args.metric == "all":
+        plot_all_metrics(off_results, on_results, args.output)
+    else:
+        plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)

benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh

@@ -0,0 +1,167 @@
+#!/bin/bash
+# Qwen3-TTS async_chunk on vs off Benchmark
+#
+# Starts two servers (async_chunk on and off), benchmarks both,
+# and generates comparison plots.
+#
+# Usage:
+#   bash run_async_chunk_benchmark.sh
+#
+# Environment variables:
+#   GPU_DEVICE       - GPU index (default: 0)
+#   MODEL            - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
+#   NUM_PROMPTS      - Prompts per concurrency level (default: 50)
+#   CONCURRENCY      - Space-separated concurrency levels (default: "1 10")
+#   PORT_ON          - Port for async_chunk on server (default: 8000)
+#   PORT_OFF         - Port for async_chunk off server (default: 8001)
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+cd "$PROJECT_ROOT"
+
+GPU_DEVICE="${GPU_DEVICE:-0}"
+MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
+NUM_PROMPTS="${NUM_PROMPTS:-50}"
+CONCURRENCY="${CONCURRENCY:-1 10}"
+NUM_WARMUPS="${NUM_WARMUPS:-3}"
+PORT_ON="${PORT_ON:-8000}"
+PORT_OFF="${PORT_OFF:-8001}"
+RESULT_DIR="${SCRIPT_DIR}/results"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+
+STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
+STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
+
+mkdir -p "${RESULT_DIR}"
+
+echo "============================================================"
+echo " Qwen3-TTS async_chunk Benchmark"
+echo "============================================================"
+echo " GPU:            ${GPU_DEVICE}"
+echo " Model:          ${MODEL}"
+echo " Prompts:        ${NUM_PROMPTS}"
+echo " Concurrency:    ${CONCURRENCY}"
+echo " Port (on/off):  ${PORT_ON} / ${PORT_OFF}"
+echo " Results:        ${RESULT_DIR}"
+echo "============================================================"
+
+cleanup() {
+    echo "Cleaning up servers..."
+    kill "$PID_ON" 2>/dev/null || true
+    kill "$PID_OFF" 2>/dev/null || true
+    wait "$PID_ON" 2>/dev/null || true
+    wait "$PID_OFF" 2>/dev/null || true
+}
+trap cleanup EXIT
+
+wait_for_server() {
+    local port=$1
+    local name=$2
+    local max_wait=300
+    local elapsed=0
+    echo "Waiting for ${name} server on port ${port}..."
+    while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
+        sleep 5
+        elapsed=$((elapsed + 5))
+        if [ $elapsed -ge $max_wait ]; then
+            echo "ERROR: ${name} server failed to start within ${max_wait}s"
+            exit 1
+        fi
+    done
+    echo "${name} server ready (${elapsed}s)"
+}
+
+# ---- Phase 1: Start async_chunk ON server ----
+echo ""
+echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
+CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
+    --stage-configs-path "${STAGE_CONFIG_ON}" \
+    --host 0.0.0.0 --port "${PORT_ON}" \
+    --trust-remote-code --enforce-eager --omni \
+    > "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
+PID_ON=$!
+
+wait_for_server "${PORT_ON}" "async_chunk_on"
+
+echo "[Phase 1] Benchmarking async_chunk ON..."
+# shellcheck disable=SC2086
+python "${SCRIPT_DIR}/bench_async_chunk.py" \
+    --host 127.0.0.1 --port "${PORT_ON}" \
+    --config-name "async_chunk_on" \
+    --num-prompts "${NUM_PROMPTS}" \
+    --max-concurrency ${CONCURRENCY} \
+    --num-warmups "${NUM_WARMUPS}" \
+    --result-dir "${RESULT_DIR}"
+
+echo "[Phase 1] Stopping async_chunk ON server..."
+kill "$PID_ON" 2>/dev/null || true
+wait "$PID_ON" 2>/dev/null || true
+sleep 5
+
+# ---- Phase 2: Start async_chunk OFF server ----
+echo ""
+echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
+CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
+    --stage-configs-path "${STAGE_CONFIG_OFF}" \
+    --host 0.0.0.0 --port "${PORT_OFF}" \
+    --trust-remote-code --enforce-eager --omni \
+    > "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
+PID_OFF=$!
+
+wait_for_server "${PORT_OFF}" "async_chunk_off"
+
+echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
+# shellcheck disable=SC2086
+python "${SCRIPT_DIR}/bench_async_chunk.py" \
+    --host 127.0.0.1 --port "${PORT_OFF}" \
+    --config-name "async_chunk_off" \
+    --num-prompts "${NUM_PROMPTS}" \
+    --max-concurrency ${CONCURRENCY} \
+    --num-warmups "${NUM_WARMUPS}" \
+    --no-stream \
+    --result-dir "${RESULT_DIR}"
+
+echo "[Phase 2] Stopping async_chunk OFF server..."
+kill "$PID_OFF" 2>/dev/null || true
+wait "$PID_OFF" 2>/dev/null || true
+
+# ---- Phase 3: Plot results ----
+echo ""
+echo "[Phase 3] Generating plots..."
+
+# Find the latest result files
+RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
+RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
+
+if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
+    echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
+    exit 1
+fi
+
+echo "  ON results:  ${RESULT_ON}"
+echo "  OFF results: ${RESULT_OFF}"
+
+# TTFP comparison (main figure)
+python "${SCRIPT_DIR}/plot_async_chunk.py" \
+    --off "${RESULT_OFF}" \
+    --on "${RESULT_ON}" \
+    --metric ttfp \
+    --output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
+
+# All metrics comparison
+python "${SCRIPT_DIR}/plot_async_chunk.py" \
+    --off "${RESULT_OFF}" \
+    --on "${RESULT_ON}" \
+    --metric all \
+    --output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
+
+echo ""
+echo "============================================================"
+echo " Benchmark complete!"
+echo " Results: ${RESULT_DIR}/"
+echo " Plots:"
+echo "   - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
+echo "   - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
+echo "============================================================"

docker/Dockerfile.ci

@@ -1,9 +1,8 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.16.0
+ARG VLLM_BASE_TAG=v0.17.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
-
 COPY . .
 
 # Install system dependencies

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.16.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.

docker/Dockerfile.xpu

@@ -9,8 +9,7 @@ ARG PYTHON_VERSION=3.12
 ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
 
 RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
-    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
-    add-apt-repository -y ppa:kobuk-team/intel-graphics
+    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
 
 RUN apt clean && apt-get update -y && \
     apt-get install -y --no-install-recommends --fix-missing \
@@ -31,9 +30,22 @@ RUN apt clean && apt-get update -y && \
     python3-pip
 
 RUN apt update && apt upgrade -y && \
-    apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \
     apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
 
+# Install UMD with fixed version
+RUN mkdir neo && \
+    cd neo && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
+    dpkg -i *.deb && \
+    cd .. && \
+    rm -rf neo
+
 ENV PATH="/root/.local/bin:$PATH"
 ENV VIRTUAL_ENV="/opt/venv"
 ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python

docs/.nav.yml

@@ -66,6 +66,7 @@ nav:
   - Model Implementation:
     - contributing/model/README.md
     - contributing/model/adding_omni_model.md
+    - contributing/model/adding_tts_model.md
     - contributing/model/adding_diffusion_model.md
   - CI: contributing/ci
   - Design Documents:

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)
 
 Currently, the main options are maintained by stage configs for each model.
 
@@ -19,3 +19,4 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 - **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
 - **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
 - **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
+- **[CPU Offloading](../user_guide/diffusion/cpu_offload_diffusion.md)** - Enable CPU offloading (model-level and layerwise) for for DiT models

docs/configuration/stage_configs.md

@@ -135,6 +135,14 @@ Each stage in the `stage_args` list contains the following configuration options
 
 A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`).
 
+### `prompt_expand_func` (Optional)
+
+A custom Python function hook for the LLM stage (Stage 0) that expands a single incoming prompt object into multiple prompts. This is primarily used for multi-modal Classifier-Free Guidance (CFG), where it generates the necessary companion requests (like a negative text prompt) and tags them with internal roles (e.g., `cfg_text`). This ensures the upstream LLM generates the needed contextual hidden states for both the conditional and unconditional generations simultaneously.
+
+### `cfg_kv_collect_func` (Optional)
+
+A custom Python function hook for downstream diffusion stages (Stage 1+) to collect, map, and process the KV caches transferred from the companion requests fired by `prompt_expand_func`. It aggregates the hidden condition states cleanly (e.g., binding them as `cfg_text_past_key_values` and `cfg_text_kv_metadata`), allowing the diffusion runtime to perform CFG smoothly without redundantly evaluating text paths on the DiT workers.
+
 ### `runtime`
 
 Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed.

docs/contributing/README.md

@@ -71,14 +71,7 @@ For additional features and advanced configurations, refer to the:
 ### Testing
 
 vLLM-Omni uses `pytest` to test the codebase.
-
-```bash
-# Run all tests
-pytest tests/
-
-# Run tests for a single test file with detailed output
-pytest -s -v tests/test_omni_llm.py
-```
+Please refer to the [test instructions](./ci/test_guide.md) for detailed testing information.
 
 !!! warning
     Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
@@ -123,6 +116,9 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 !!! note
     If the PR spans more than one category, please include all relevant prefixes.
 
+### Local Test
+Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
+
 ### Code Quality
 
 The PR needs to meet the following code quality standards:

docs/contributing/ci/test_guide.md

@@ -0,0 +1,75 @@
+# Test Guide
+## Setting Up the Test Environment
+### Creating a Container
+vLLM-Omni provides an official Docker image for deployment. These images are built upon vLLM Docker images and are available on [Docker Hub](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which vLLM release it is based on.
+For a local test environment, you can follow the steps below to create a container:
+## Installing Dependencies
+### vLLM & vLLM-Omni
+vLLM-Omni is built based on vLLM. You can follow [install guide](../../getting_started/installation/README.md) to build your local environment.
+
+### Test Case Dependencies
+When running test cases, you may need to install the following dependencies:
+
+```bash
+uv pip install ".[dev]"
+apt-get install -y ffmpeg
+```
+
+## Running Tests
+Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases.
+
+=== "L1 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "core_model and cpu"
+    ```
+    The latest test command is available in the "Simple Unit Test" step of this [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
+
+=== "L2 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "core_model and not cpu" --run-level=core_model
+    ```
+    If you only want to run a specific test case, you can use:
+    ```bash
+    pytest -s -v test_xxxx.py --run-level=core_model
+    ```
+    If you only want to run specific test cases on a particular platform, you can use:
+    ```bash
+    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
+    ```
+    The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
+
+=== "L3 level & L4 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "advanced_model" --run-level=advanced_model
+    ```
+    If you only want to run L3 test case, you can use:
+    ```bash
+    pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model
+    ```
+    If you only want to run a specific test case, you can use:
+    ```bash
+    pytest -s -v test_xxxx.py --run-level=advanced_model
+    ```
+    If you only want to run specific test cases on a particular platform, you can use:
+    ```bash
+    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
+    ```
+    Note: To run performance tests, use:
+    ```bash
+    pytest -s -v perf/scripts/run_benchmark.py
+    ```
+
+    The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
+
+    The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml).
+
+You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
+
+## Adding New Test Cases
+Please refer to the [L5 Layering Specification document](./CI_5levels.md).

docs/contributing/model/adding_omni_model.md

@@ -460,7 +460,7 @@ def forward(self, ...):
 These keys are then accessible in your stage transition function:
 ```python
 # In stage_input_processors/qwen3_omni.py
-thinker_embeddings = output.multimodal_output["0"]  # Access by key
+thinker_prefill_embeddings = output.multimodal_output["0"]  # Access by key
 thinker_hidden_states = output.multimodal_output["24"]
 ```
 
@@ -513,11 +513,11 @@ def thinker2talker(
     for thinker_output in thinker_outputs:
         output = thinker_output.outputs[0]
         # Extract thinker embeddings and hidden states
-        thinker_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
+        thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
         thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
 
         info = {
-            "thinker_embeddings": thinker_embeddings,
+            "thinker_prefill_embeddings": thinker_prefill_embeddings,
             "thinker_hidden_states": thinker_hidden_states,
             "thinker_sequences": thinker_output.prompt_token_ids + output.token_ids,
             "thinker_input_ids": thinker_output.prompt_token_ids,

docs/contributing/model/adding_tts_model.md

@@ -0,0 +1,584 @@
+# Adding a TTS Model
+
+This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS**
+as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the
+key optimizations all TTS models in this repo should follow.
+
+## Table of Contents
+
+1. [Overview](#overview)
+2. [Directory Structure](#directory-structure)
+3. [Step-by-Step Implementation](#step-by-step-implementation)
+4. [Key Components](#key-components)
+5. [Model Registration](#model-registration)
+6. [Stage Configuration](#stage-configuration)
+7. [Stage Input Processors](#stage-input-processors)
+8. [Testing](#testing)
+9. [Summary](#summary)
+
+## Overview
+
+vLLM-Omni supports TTS models as multi-stage pipelines where each stage runs independently
+and can be placed on different devices. Qwen3-TTS has two stages:
+
+| Stage | Name | Input | Output |
+|-------|------|-------|--------|
+| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes |
+| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform |
+
+Each stage is a separate model class configured independently via YAML. The two stages
+are connected by the `async_chunk` framework, which enables inter-stage streaming for
+low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk_design.md)).
+
+### Without async_chunk (batch mode)
+
+Stage 0 runs to completion before Stage 1 starts, resulting in long first-packet latency:
+
+```mermaid
+flowchart TB
+    subgraph stage0["Stage 0: AR Code Predictor"]
+        direction LR
+        P[Prefill] --> D1[Decode 1]
+        D1 --> D2[Decode 2]
+        D2 --> Dots1["..."]
+        Dots1 --> DN[Decode N]
+    end
+
+    subgraph stage1["Stage 1: Code2Wav"]
+        direction LR
+        DEC[Decode all codes at once]
+    end
+
+    stage0 -- "all N codes" --> stage1
+    stage1 --> FPL["First Packet Latency = Stage 0 + Stage 1"]
+
+    style stage0 fill:#dae8fc,stroke:#6c8ebf
+    style stage1 fill:#f8d7c8,stroke:#d4856a
+    style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
+```
+
+### With async_chunk (streaming mode)
+
+Stage 0 sends codec codes to Stage 1 every `chunk_size=25` tokens. Stage 1 begins decoding
+immediately, reducing first-packet latency from the full AR time to just the first chunk:
+
+```mermaid
+flowchart TB
+    subgraph stage0["Stage 0: Code Predictor (AR)"]
+        direction LR
+        P[Prefill] --> D1["Decode 1-25"]
+        D1 --> D2["Decode 26-50"]
+        D2 --> Dots1["..."]
+        Dots1 --> DN["Decode N"]
+    end
+
+    subgraph stage1["Stage 1: Code2Wav"]
+        direction LR
+        C1["Chunk 1\n(25 frames)"] --> C2["Chunk 2\n(context + 25)"]
+        C2 --> Dots2["..."]
+        Dots2 --> CN["Final chunk"]
+    end
+
+    D1 -. "chunk 1 (25 codes)" .-> C1
+    D2 -. "chunk 2 (context + 25)" .-> C2
+    DN -. "final" .-> CN
+
+    stage0 --> FPL["⏱ First Packet Latency = Prefill + 25 decode steps only"]
+
+    style stage0 fill:#dae8fc,stroke:#6c8ebf
+    style stage1 fill:#e8d4f8,stroke:#8a6cad
+    style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
+```
+
+Key parameters: `chunk_size=25`, `left_context_size=25` (validated defaults from Qwen3-TTS
+and Qwen3-Omni).
+
+## Directory Structure
+
+When adding a new TTS model, create the following structure:
+
+```
+vllm_omni/model_executor/models/
+  your_model_name/
+    __init__.py
+    your_model.py                    # Unified class (stage dispatch)
+    your_model_ar_stage.py           # Stage 0: AR stage
+    your_model_decoder.py            # Stage 1: audio decoder
+
+vllm_omni/model_executor/stage_input_processors/
+  your_model_name.py                 # Stage 0 -> Stage 1 transition
+
+vllm_omni/model_executor/stage_configs/
+  your_model_name.yaml               # Batch mode config
+  your_model_name_async_chunk.yaml   # Streaming mode config
+```
+
+**Qwen3-TTS reference files:**
+
+| File | Purpose |
+|------|---------|
+| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
+| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR |
+| `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder |
+| `stage_configs/qwen3_tts.yaml` | Stage config (async_chunk enabled) |
+| `stage_configs/qwen3_tts_batch.yaml` | Batch mode config |
+| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
+
+## Step-by-Step Implementation
+
+### Step 1: Implement Stage 0 - AR Stage
+
+Stage 0 is the autoregressive stage that generates intermediate audio representations.
+**It must use vLLM's native decoder layers with fused ops and PagedAttention** for the LLM
+backbone - this is the primary source of speedup over HuggingFace inference.
+
+#### 1.1 Use vLLM Decoder Layers Directly
+
+Build your transformer layers from the corresponding vLLM decoder layer class (e.g.
+`Qwen3DecoderLayer` for Qwen3-based backbones, or the equivalent for LLaMA, Qwen2, etc.).
+Do not wrap the HuggingFace model directly - that bypasses PagedAttention and fused kernels.
+
+```python
+# your_model_ar_stage.py
+
+from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer
+
+class YourTTSARStage(nn.Module):
+
+    def __init__(self, config, vllm_config, prefix):
+        self.layers = nn.ModuleList([
+            Qwen3DecoderLayer(
+                config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}"
+            )
+            for i in range(config.num_hidden_layers)
+        ])
+        self.lm_head = ParallelLMHead(config.codec_size, config.hidden_size)
+```
+
+See `qwen3_tts_code_predictor_vllm.py` for the full implementation.
+
+#### 1.2 Forward Pass
+
+Implement `forward()` to return an `OmniOutput` with intermediate data for Stage 1:
+
+```python
+def forward(self, input_ids, positions, intermediate_tensors=None,
+            inputs_embeds=None, **kwargs) -> OmniOutput:
+    hidden_states = self.run_layers(input_ids, positions, intermediate_tensors, inputs_embeds)
+    logits = self.lm_head(hidden_states)
+
+    return OmniOutput(
+        text_hidden_states=hidden_states,
+        multimodal_outputs={
+            "audio_codes": self.extract_codes(logits),
+        },
+    )
+```
+
+The keys in `multimodal_outputs` are what your stage input processor will read to build
+Stage 1 inputs.
+
+#### 1.3 Weight Loading with Fused QKV
+
+When using vLLM's fused `QKVParallelLinear`, pack the HF `q_proj`/`k_proj`/`v_proj` weights
+into `qkv_proj` using `stacked_params_mapping`. See the `load_weights()` method in
+`qwen3_tts_code_predictor_vllm.py` for the standard pattern - it can be reused as-is
+for any Qwen-family backbone.
+
+#### 1.4 Custom Stop Condition (if needed)
+
+Some TTS models use a learned stop head rather than an EOS token. If your model does this,
+implement it inside `sample()`:
+
+```python
+def sample(self, logits, sampling_metadata) -> SamplerOutput | None:
+    output = self.sampler(logits, sampling_metadata)
+    if self._stop_head_fired():
+        output = mark_as_finished(output)
+    return output
+```
+
+### Step 2: Implement Stage 1 - Decoder
+
+Stage 1 decodes Stage 0 output into audio. It runs outside the scheduler (no PagedAttention
+needed). Implement `chunked_decode_streaming()` to support async_chunk streaming:
+
+```python
+# your_model_decoder.py
+
+class YourTTSDecoder(nn.Module):
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        # Initialize your audio decoder (SpeechTokenizer, HiFiGAN, etc.)
+
+    def forward(self, codes: torch.Tensor, **kwargs) -> torch.Tensor:
+        return self.decoder(codes)
+
+    def chunked_decode_streaming(self, codes, chunk_size=25,
+                                  left_context_size=25) -> torch.Tensor:
+        """Decode with a sliding context window for smooth chunk boundaries."""
+        end_index = codes.shape[-1]
+        context_size = 0 if end_index <= chunk_size else left_context_size
+        wav_chunk = self(codes)
+        # Trim left context to avoid duplicate audio
+        return wav_chunk[..., context_size * self.total_upsample:]
+```
+
+### Step 3: Implement the Unified Model Class
+
+The unified class dispatches to the correct stage based on `model_stage` in the config:
+
+```python
+# your_model.py
+
+class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP):
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        self.model_stage = vllm_config.model_config.model_stage
+
+        if self.model_stage == "ar_stage":
+            ar_vllm_config = vllm_config.with_hf_config(
+                vllm_config.model_config.hf_config.ar_config,
+                architectures=["YourTTSARStageForConditionalGeneration"],
+            )
+            self.ar_stage = init_vllm_registered_model(
+                vllm_config=ar_vllm_config,
+                prefix=maybe_prefix(prefix, "ar"),
+                hf_config=ar_vllm_config.model_config.hf_config,
+                architectures=["YourTTSARStageForConditionalGeneration"],
+            )
+            self.model = self.ar_stage
+
+        elif self.model_stage == "decoder":
+            self.decoder = YourTTSDecoder(vllm_config=vllm_config, prefix=prefix)
+            self.model = self.decoder
+```
+
+### Step 4: Create `__init__.py`
+
+```python
+# vllm_omni/model_executor/models/your_model_name/__init__.py
+from .your_model import YourTTSModelForConditionalGeneration
+
+__all__ = ["YourTTSModelForConditionalGeneration"]
+```
+
+## Key Components
+
+### Model Interfaces
+
+Your unified model class should implement the appropriate interfaces:
+
+- **`SupportsPP`**: Required for pipeline parallelism support (all models should implement this)
+- **`SupportsMultiModal`**: Only if your model accepts multimodal inputs (e.g. reference audio for voice cloning)
+
+### Output Format
+
+Use `OmniOutput` so the orchestrator can route intermediate data between stages:
+
+```python
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+return OmniOutput(
+    text_hidden_states=hidden_states,
+    multimodal_outputs={
+        "audio_codes": codec_codes,
+    },
+)
+```
+
+### Weight Loading from a Single Checkpoint
+
+If both stages load from one checkpoint, separate them by prefix in the unified class:
+
+```python
+def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+    ar_weights, decoder_weights = [], []
+    for name, tensor in weights:
+        if name.startswith("decoder."):
+            decoder_weights.append((name, tensor))
+        else:
+            ar_weights.append((name, tensor))
+
+    if self.model_stage == "ar_stage":
+        return self.ar_stage.load_weights(ar_weights)
+    elif self.model_stage == "decoder":
+        return self.decoder.load_weights(decoder_weights)
+```
+
+## Model Registration
+
+Register all stage classes in `vllm_omni/model_executor/models/registry.py`:
+
+```python
+_OMNI_MODELS = {
+    # (package_name, module_name, class_name)
+    "YourTTSModelForConditionalGeneration": (
+        "your_model_name", "your_model",
+        "YourTTSModelForConditionalGeneration",
+    ),
+    "YourTTSARStageForConditionalGeneration": (
+        "your_model_name", "your_model_ar_stage",
+        "YourTTSARStageForConditionalGeneration",
+    ),
+    "YourTTSDecoder": (
+        "your_model_name", "your_model_decoder",
+        "YourTTSDecoder",
+    ),
+}
+```
+
+The registry uses lazy loading - model classes are only imported when needed.
+
+## Stage Configuration
+
+Each stage has a `worker_type` that determines how it is scheduled:
+
+- `worker_type: ar` - autoregressive stage, uses `OmniARScheduler` with PagedAttention
+- `worker_type: generation` - non-AR stage (e.g. decoder), uses `OmniGenerationScheduler`
+
+Key configuration fields:
+
+| Field | Description |
+|-------|-------------|
+| `model_stage` | Which stage to initialize (`ar_stage`, `decoder`, etc.) |
+| `model_arch` | Architecture name, must match `registry.py` |
+| `engine_input_source` | List of upstream stage IDs that provide input (e.g. `[0]`) |
+| `engine_output_type` | Output type: `latent` for intermediate, `audio` for final |
+| `custom_process_next_stage_input_func` | Async chunk processor function path (streaming only) |
+| `final_output` | Whether this stage produces the final user-facing output |
+| `final_output_type` | Type of final output (`audio`, `text`, etc.) |
+
+### Batch mode
+
+```yaml
+# stage_configs/your_model_name.yaml
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 64
+    engine_args:
+      model_stage: ar_stage
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      engine_output_type: latent
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 2048
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: decoder
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: audio
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+```
+
+### Streaming mode (async_chunk)
+
+Add `async_chunk: true` at the top level and specify `custom_process_next_stage_input_func`
+on Stage 0 to define how intermediate outputs are chunked and forwarded:
+
+```yaml
+# stage_configs/your_model_name_async_chunk.yaml
+
+async_chunk: true
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 64
+    engine_args:
+      model_stage: ar_stage
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      engine_output_type: latent
+      custom_process_next_stage_input_func: >
+        vllm_omni.model_executor.stage_input_processors.your_model_name.ar2decoder_async_chunk
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 2048
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+    engine_args:
+      model_stage: decoder
+      model_arch: YourTTSModelForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: audio
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+```
+
+## Stage Input Processors
+
+Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours in
+`vllm_omni/model_executor/stage_input_processors/your_model_name.py`.
+
+See `stage_input_processors/qwen3_tts.py` for the full reference implementation.
+
+### Data structures
+
+Understanding what's available in stage outputs:
+
+- `stage_list[source_id].engine_outputs` - list of `EngineCoreOutput` objects
+- Each `EngineCoreOutput` has `outputs` - list of `RequestOutput` objects
+- Each `RequestOutput` has:
+  - `token_ids` - generated token IDs
+  - `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs`
+  - `prompt_token_ids` - original prompt token IDs
+
+### Batch mode (non-streaming)
+
+Collects all Stage 0 outputs and forwards them to Stage 1 in one shot:
+
+```python
+def ar2decoder(
+    stage_list: list[Any],
+    engine_input_source: list[int],
+    prompt: OmniTokensPrompt | TextPrompt | None = None,
+    requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    source_id = engine_input_source[0]
+    decoder_inputs = []
+
+    for output in stage_list[source_id].engine_outputs:
+        result = output.outputs[0]
+        codes = result.multimodal_output["audio_codes"].cpu()
+        decoder_inputs.append(
+            OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist())
+        )
+
+    return decoder_inputs
+```
+
+### Streaming mode (async_chunk)
+
+Buffers Stage 0 outputs and forwards a chunk to Stage 1 once `chunk_size` frames
+have accumulated. The function signature follows the `OmniChunkTransferAdapter` protocol:
+
+```python
+def ar2decoder_async_chunk(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any] | None,
+    request: Any,
+    is_finished: bool = False,
+) -> dict[str, Any] | None:
+    """Forward chunks of AR output to the decoder stage."""
+    request_id = request.external_req_id
+    finished = bool(is_finished or request.is_finished())
+
+    # Extract and buffer the latest frame
+    if isinstance(pooling_output, dict):
+        frame = extract_frame(pooling_output)
+        if frame is not None:
+            transfer_manager.code_prompt_token_ids[request_id].append(
+                frame.cpu().tolist()
+            )
+    elif not finished:
+        return None
+
+    # Read chunk config from connector
+    chunk_size = 25
+    left_context_size = 25
+
+    length = len(transfer_manager.code_prompt_token_ids[request_id])
+    if length <= 0:
+        if finished:
+            return {"codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
+        return None
+
+    # Wait until a full chunk is ready (or request is finished)
+    chunk_length = length % chunk_size
+    if chunk_length != 0 and not finished:
+        return None
+
+    # Build context window: left_context + chunk
+    context_length = chunk_length if chunk_length != 0 else chunk_size
+    end_index = min(length, left_context_size + context_length)
+    window = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
+
+    return {
+        "codes": torch.tensor(window).transpose(0, 1).reshape(-1).tolist(),
+        "left_context_size": max(0, int(end_index - context_length)),
+        "finished": torch.tensor(finished, dtype=torch.bool),
+    }
+```
+
+Key points:
+- `transfer_manager` is the `OmniChunkTransferAdapter` that owns the chunk lifecycle
+- Each call appends one AR decode step's output; a chunk is emitted every `chunk_size` steps
+- The final (possibly partial) chunk is flushed when `is_finished` is true
+- `left_context_size` frames of overlap are included for smooth audio boundaries
+
+## Testing
+
+For general testing conventions, see [tests_style.md](../ci/tests_style.md).
+
+Recommended test cases for a new TTS model:
+
+1. **Single request** - verify waveform output shape and sample rate
+2. **Batched requests** - verify each request in the batch finishes independently
+3. **async_chunk streaming** - verify audio chunks arrive incrementally and decode correctly
+4. **Speaker conditioning** (if applicable) - verify different speaker inputs produce different outputs
+
+Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py`
+
+## Adding a Model Recipe
+
+After implementing and testing your model, add a model recipe to the
+[vllm-project/recipes](https://github.com/vllm-project/recipes) repository so users can
+get started quickly. See [Adding an Omni-Modality Model](./adding_omni_model.md#adding-a-model-recipe)
+for the expected format.
+
+## Summary
+
+Adding a TTS model to vLLM-Omni involves:
+
+1. **Create model directory** with AR stage, decoder stage, and unified class
+2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly
+3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()`
+4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration`
+5. **Register** all stage classes in `registry.py`
+6. **YAML configs** - provide both batch and `async_chunk` variants
+7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25
+8. **Tests** - cover single request, batching, and async_chunk streaming
+9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
+
+### Qwen3-TTS Reference Files
+
+| File | Purpose |
+|------|---------|
+| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
+| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops |
+| `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` |
+| `stage_configs/qwen3_tts.yaml` | Stage configuration |
+| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
+
+For more information, see:
+
+- [Architecture Overview](../../design/architecture_overview.md)
+- [Async Chunk Design](../../design/feature/async_chunk_design.md)
+- [Stage Configuration Guide](../../configuration/stage_configs.md)

docs/design/architecture_overview.md

@@ -92,6 +92,13 @@ The framework achieves high performance through several optimization techniques:
     * **Quantization:** Supports various quantization implementations including FP8 and AWQ.
     * **FusedOps:** Allows for custom and third-party integration.
 
+### Classifier-Free Guidance (CFG) Companion Flow
+
+vLLM-Omni natively models Classifier-Free Guidance (CFG) across disaggregated multi-stage setups via a "companion request" paradigm, eliminating redundant textual/multimodal context computation boundaries:
+1. **Prompt Expansion:** In the initial autoregressive (AR) stage, a customized `prompt_expand_func` hook intercepts incoming generation prompts and pairs them directly with negative companion prompts (e.g., a default negative prompt) on the fly, tagging the secondary prompt with a specific internal role (`cfg_text`).
+2. **Synchronized KV Cache Transfer:** The AR stage evaluates both the primary and companion sequence batches concurrently. The `OmniConnector` captures these specific structural dependencies and reliably passes the positive and negative outcome KV caches seamlessly across stage boundaries via shared memory or network protocols.
+3. **KV Cache Collection & Injection:** Upon reaching the downstream Diffusion (DiT) Engine, an assigned `cfg_kv_collect_func` automatically intercepts the mapped companion caches (`cfg_text_past_key_values`). These auxiliary dependencies are natively gathered and seamlessly bound to the primary generation sequence variables, enabling the DiT Engine to cleanly implement cross-attention CFG guidance over accurate conditioning and unconditioning structures in parallel.
+
 ### Flexibility and Usability
 
 vLLM-Omni is designed to be flexible and straightforward for users:

docs/design/feature/async_chunk_design.md

@@ -19,7 +19,7 @@ The `async_chunk` feature enables asynchronous, chunked processing of data acros
 
 For qwen3-omni:
 - **Thinker → Talker**: Per decode step (typically chunk_size=1)
-- **Talker → Code2Wav**: Accumulated to code2wav chunk_size(default=25, currently only support default, will support chunk_size soon) before sending
+- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. Set `initial_codec_chunk_frames` to emit smaller chunks during the initial phase for reduced TTFA
 - **Code2Wav**: Streaming decode with code2wav chunk_size
 
 With `async_chunk`:

docs/design/feature/sequence_parallel.md

@@ -240,7 +240,8 @@ class TransformerWithRoPE(nn.Module):
 
 ## Approach 2: Intrusive Modification (For Complex Cases)
 
-For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
+For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. Importantly, when taking this approach, be careful to ensure that you correctly manage the `_sp_shard_depth`; if the sequence parallel shard depth is 0, Ulysses will not be used.
+
 
 **When to use:**
 - Dynamic/conditional sharding logic
@@ -252,15 +253,18 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
 
 def forward(self, hidden_states, ...):
     if self.parallel_config.sequence_parallel_size > 1:
+        # <Increment the _sp_shard_depth on the fwd context>
         hidden_states = sp_shard(hidden_states, dim=1)
 
     # ... computation ...
 
     if self.parallel_config.sequence_parallel_size > 1:
         output = sp_gather(output, dim=1)
+        # <Decrement the _sp_shard_depth on the fwd context>
 
     return output
 ```
+Note that currently, `sp_shard` / `sp_gather` do *not* automatically manage the `_sp_shard_depth`; you need to be careful to manage it yourself.
 
 ---
 

docs/design/feature/vae_parallel.md

@@ -0,0 +1,289 @@
+# VAE Patch Parallelism
+
+This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
+We use **Qwen-Image** as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is Vae Patch parallel?
+
+**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
+
++ Split into multiple spatial tiles
+
++ Distributed across multiple ranks
+
++ Decoded in parallel
+
++ Merged to reconstruct the final output
+
+This approach:
+
++ Distributes computation across multiple devices
+
++ Reduces peak memory usage per device
+
++ Accelerates decoding latency
+
+### Architecture
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
+
+The executor is model-agnostic and accepts three function parameters:
+
++ split – Partition the latent into tiles
+
++ exec – Decode a single tile
+
++ merge – Combine decoded tiles into the final output
+
+#### Execution Flow
+
++ Call split(z) to generate a list of TileTask and a GridSpec
+
++ Dispatch tasks across ranks using workload-based balancing
+
++ Each rank executes exec(task) on its assigned tiles
+
++ Gather decoded tile results to rank 0
+
++ Rank 0 performs merge(...)
+
++ (Optional) Broadcast final result to all ranks
+
+This design separates:
+
++ Distributed execution logic
+
++ Model-specific tiling and merging logic
+
+#### Why split / exec / merge is necessary?
+
+The latent tensor cannot be arbitrarily partitioned.
+
+During decoding:
+
++ Each output pixel may depend on neighboring pixels
+
++ The receptive field is model-dependent
+
+Therefore:
+
++ Tiles must include overlap
+
++ Merge must perform blending to avoid seams
+
+## Step-by-Step Implementation
+
+### Step 1: Implement DistributedAutoencoderKLQwenImage
+`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
+
+
+```
+class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+```
+**Key points**:
++ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
++ Call init_distributed() after loading weights
+
+### Step 2: Implement split/exec/merge
+Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`:
+```
+class GridSpec:
+    split_dims: tuple[int, ...]  # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W))
+    grid_shape: tuple[int, ...]  # Tile grid layout (num_rows, num_cols)
+    tile_spec: dict = field(default_factory=dict) # Metadata required for merging
+    output_dtype: torch.dtype | None = None # Final output dtype
+```
+```
+class TileTask:
+    tile_id: int # task id
+    grid_coord: tuple[int, ...]  # Tile position in grid
+    tensor: torch.Tensor | list[torch.Tensor]  # The tile tensor
+    workload: int | float = 1 # Used for load balancing (e.g., tile area)
+```
+And tiled base split/exec/merge as follow:
+```
+def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+    # mostly copy from AutoencoderKL
+    _, _, num_frames, height, width = z.shape
+    sample_height = height * self.spatial_compression_ratio
+    sample_width = width * self.spatial_compression_ratio
+
+    tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+    tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+    tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+    tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+    blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+    blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+    # Split z into overlapping tiles and decode them separately.
+    # The tiles have an overlap to avoid seams between tiles.
+    tiletask_list = []
+    for i in range(0, height, tile_latent_stride_height):
+        for j in range(0, width, tile_latent_stride_width):
+            time_list = []
+            for k in range(num_frames):
+                self._conv_idx = [0]
+                tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                time_list.append(tile)
+            tiletask_list.append(
+                TileTask(
+                    len(tiletask_list),
+                    (i // tile_latent_stride_height, j // tile_latent_stride_width),
+                    time_list,
+                    workload=time_list[0].shape[3] * time_list[0].shape[4],
+                )
+            )
+    tile_spec = {
+        "sample_height": sample_height,
+        "sample_width": sample_width,
+        "blend_height": blend_height,
+        "blend_width": blend_width,
+    }
+    grid_spec = GridSpec(
+        split_dims=(3, 4),
+        grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+        tile_spec=tile_spec,
+        output_dtype=self.dtype,
+    )
+    return tiletask_list, grid_spec
+
+def tile_exec(self, task: TileTask) -> torch.Tensor:
+    """Decode a single latent tile into RGB space."""
+    self.clear_cache()
+    time = []
+    for k in range(len(task.tensor)):
+        self._conv_idx = [0]
+        tile = self.post_quant_conv(task.tensor[k])
+        decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
+        time.append(decoded)
+    result = torch.cat(time, dim=2)
+    return result
+
+def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+    """Merge decoded tiles into a full image."""
+    grid_h, grid_w = grid_spec.grid_shape
+    result_rows = []
+    self.clear_cache()
+
+    result_rows = []
+    for i in range(grid_h):
+        result_row = []
+        for j in range(grid_w):
+            tile = coord_tensor_map[(i, j)]
+            if i > 0:
+                tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+            if j > 0:
+                tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+            result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
+        result_rows.append(torch.cat(result_row, dim=-1))
+    dec = torch.cat(result_rows, dim=3)[
+        :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
+    ]
+    return dec
+```
+
+### Step 3: Override tiled_decode
+We need to override tiled_decode, the main logic is:
++ check distributed is enabled
++ select split/exec/merge
++ Invoke self.distributed_decoder.execute to decode
+```
+def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
+    if not self.is_distributed_enabled():
+        return super().tiled_decode(z, return_dict=return_dict)
+
+    logger.info("Decode run with distributed executor")
+    result = self.distributed_decoder.execute(
+        z,
+        DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
+        broadcast_result=True,
+    )
+    if not return_dict:
+        return (result,)
+
+    return DecoderOutput(sample=result)
+```
+`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0.
+
+### Step 4: Modify Pipeline
+Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage
+```
+class YourModelPipeline(nn.Module):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        ...
+-       self.vae = AutoencoderKL.from_pretrained(
+-           model, subfolder="vae", local_files_only=local_files_only).to(self.device)
++       self.vae = DistributedAutoencoderKL.from_pretrained(
++           model, subfolder="vae", local_files_only=local_files_only
++       ).to(self.device)
+```
+
+## Testing
+Verify numerical consistency between:
++ vae_patch_parallel_size = 1
+
++ vae_patch_parallel_size = N
+
+Example:
+torch.allclose(output_1, output_n, atol=1e-5)
+
+Testing requirements:
++ Fix random seed
++ Use identical tiling strategy
+
+```python
+m = Omni(
+        model=model_name,
+        vae_use_tiling=True,
+        parallel_config=DiffusionParallelConfig(
+            tensor_parallel_size=2,
+            vae_patch_parallel_size=1, # or 2
+        ),
+    )
+```
+When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead.
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Notes |
+|-------|------|-------|
+| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
+| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
+| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
+
+---
+
+## Summary
+
+Adding Vae Patch Parallel support to diffusion model:
+
+1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
+2. **Change vae model in pipeline to Distributed Vae**
+3. **Test** - Verify with `tensor_parallel_size=N` quality

docs/getting_started/installation/gpu.md

@@ -30,7 +30,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

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

@@ -40,7 +40,7 @@ If you do not need to modify source code of vLLM, you can directly install the s
 uv pip install vllm==0.16.0 --torch-backend=auto
 ```
 
-The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
+The release 0.16.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -91,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.14.0 \
+    vllm/vllm-omni:v0.16.0 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

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

@@ -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.14.0 \
+  vllm/vllm-omni-rocm:v0.16.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.14.0
+  vllm/vllm-omni-rocm:v0.16.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/models/supported_models.md

@@ -30,10 +30,15 @@ th {
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
 | `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
+| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` |
+| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` |
+| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` |
 |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
 |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
+|`CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` |
+|`MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
 |`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
 |`OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` |

docs/serving/speech_api.md

@@ -96,6 +96,7 @@ Content-Type: application/json
 | `language` | string | "Auto" | Language (see supported languages below) |
 | `instructions` | string | "" | Voice style/emotion instructions |
 | `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
+| `initial_codec_chunk_frames` | integer | null | Initial chunk size for reduced TTFA (overrides stage config) |
 
 **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
 
@@ -103,7 +104,7 @@ Content-Type: application/json
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `ref_audio` | string | null | Reference audio (URL or base64 data URL) |
+| `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
 | `ref_text` | string | null | Transcript of reference audio |
 | `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
 
@@ -124,6 +125,81 @@ Lists available voices for the loaded model.
     "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
 }
 ```
+```
+POST /v1/audio/voices
+Content-Type: multipart/form-data
+```
+
+Upload a new voice sample for voice cloning in Base task TTS requests.
+
+**Form Parameters:**
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) |
+| `consent` | string | Yes | Consent recording ID |
+| `name` | string | Yes | Name for the new voice |
+
+**Response Example:**
+
+```json
+{
+  "success": true,
+  "voice": {
+    "name": "custom_voice_1",
+    "consent": "user_consent_id",
+    "created_at": 1738660000,
+    "mime_type": "audio/wav",
+    "file_size": 1024000
+  }
+}
+```
+
+**Usage Example:**
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/voices \
+  -F "audio_sample=@/path/to/voice_sample.wav" \
+  -F "consent=user_consent_id" \
+  -F "name=custom_voice_1"
+```
+
+
+```bash
+DELETE /v1/audio/voices/{name}
+```
+
+Delete an uploaded voice sample.
+
+**Path Parameters:**
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `name` | string | Yes | Name of the voice to delete |
+
+**Response Example:**
+
+```json
+{
+  "success": true,
+  "message": "Voice 'custom_voice_1' deleted successfully"
+}
+```
+
+**Error Response (404 Not Found):**
+
+```json
+{
+  "success": false,
+  "error": "Voice 'unknown_voice' not found"
+}
+```
+
+**Usage Example:**
+
+```bash
+curl -X DELETE http://localhost:8091/v1/audio/voices/custom_voice_1
+```
 
 ## Examples
 
@@ -184,6 +260,25 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     }' --output cloned.wav
 ```
 
+upload voice
+```bash
+curl -X POST http://localhost:8091/v1/audio/voices \
+  -F "audio_sample=@/path/to/voice_sample.wav" \
+  -F "consent=user_consent_id" \
+  -F "name=custom_voice_1"
+```
+
+use upload voice
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, this is a cloned voice",
+        "task_type": "Base",
+        "voice": "custom_voice_1"
+    }' --output cloned.wav
+```
+
 ## Supported Models
 
 | Model | Task Type | Description |

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -27,7 +27,7 @@ The following table shows which models are currently supported by parallelism me
 | **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |
 | **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |
 | **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |         ❌          |
-| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
+| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ✅          |
 | **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
 | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
 | **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
@@ -51,9 +51,10 @@ The following table shows which models are currently supported by parallelism me
 
 ### VideoGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP |
-|-------|------------------|:----------:|:--------------:|:---------------:|:----:|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ |
+| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel |
+|-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **LTX-2** | `Lightricks/LTX-2` | ✅ | ✅ | ✅ | ❌ | ❌ |
 
 ### Tensor Parallelism
 
@@ -85,7 +86,7 @@ outputs = omni.generate(
 VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP).
 
 !!! note "Enablement and feature gate"
-    - VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`, `stabilityai/stable-diffusion-3.5`).
+    - VAE patch parallelism is currently **enabled only for validated pipelines** (check [ImageGen](#imagegen) and [VideoGen](#videogen) for more information).
     - If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.)
 
 #### Offline Inference

docs/user_guide/diffusion/quantization/fp8.md

@@ -62,6 +62,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
 |-------|-----------|---------------|------------------|
 | 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 |
 
 ## Combining with Other Features
 

docs/user_guide/diffusion_acceleration.md

@@ -58,8 +58,8 @@ The following table shows which models are currently supported by each accelerat
 | **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
 | **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
 | **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
 | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
 | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
@@ -67,13 +67,15 @@ The following table shows which models are currently supported by each accelerat
 | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
 | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ |
+| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
+| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
 
 ### VideoGen
 
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel |
+|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
+| **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 
 ### Quantization
 

docs/user_guide/examples/offline_inference/qwen3_tts.md

@@ -11,7 +11,7 @@ Qwen3 TTS provides multiple task variants for speech generation:
 
 - **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction.
 - **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice.
-- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection.
+- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection. The `ref_audio` field accepts a local file path, HTTP URL, or base64 data URL.
 
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
@@ -98,7 +98,8 @@ Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires
 python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
 ```
 
-Each 25-frame Code2Wav chunk is logged as it arrives. The final WAV file is written once generation
+Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
+and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation
 completes. This demonstrates that audio data is available progressively rather than only at the end.
 
 > **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.

docs/user_guide/examples/online_serving/qwen3_tts.md

@@ -112,7 +112,7 @@ The Python client supports the following command-line arguments:
 - `--voice`: Speaker/voice name (default: `vivian`). Options: `vivian`, `ryan`, `aiden`, etc.
 - `--language`: Language. Options: `Auto`, `Chinese`, `English`, `Japanese`, `Korean`, `German`, `French`, `Russian`, `Portuguese`, `Spanish`, `Italian`
 - `--instructions`: Voice style/emotion instructions
-- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task)
+- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task). Local paths are automatically base64-encoded by the client before sending to the server.
 - `--ref-text`: Reference audio transcript for voice cloning (Base task)
 - `--response-format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`, `aac`, `opus`
 - `--output` (or `-o`): Output audio file path (default: `tts_output.wav`)
@@ -220,7 +220,7 @@ Lists available voices for the loaded model:
     "task_type": "CustomVoice",
     "language": "Auto",
     "instructions": "Optional style instructions",
-    "ref_audio": "URL or base64 for voice cloning",
+    "ref_audio": "HTTP URL, base64 data URL, or file:// URI for voice cloning",
     "ref_text": "Reference audio transcript",
     "x_vector_only_mode": false,
     "max_new_tokens": 2048
@@ -261,14 +261,14 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 
 | Parameter            | Type   | Required | Description                                  |
 | -------------------- | ------ | -------- | -------------------------------------------- |
-| `ref_audio`          | string | **Yes**  | Reference audio (URL or base64 data URL)     |
+| `ref_audio`          | string | **Yes**  | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
 | `ref_text`           | string | No       | Transcript of reference audio (for ICL mode) |
 | `x_vector_only_mode` | bool   | No       | Use speaker embedding only (no ICL)          |
 
 ## Streaming
 
 Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded
-(one chunk per 25-frame Code2Wav window):
+(one chunk per Code2Wav window, default 25 frames; configurable in the stage config):
 
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \

docs/user_guide/examples/online_serving/text_to_image.md

@@ -23,6 +23,29 @@ Or use the startup script:
 bash run_server.sh
 ```
 
+### Start with Parallelism Acceleration
+
+Enable Tensor Parallelism and VAE Patch Parallelism for faster inference:
+
+```bash
+# With Tensor Parallelism (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2
+
+# With Tensor Parallelism and VAE Patch Parallelism (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2 --vae-patch-parallel-size 2 --vae-use-tiling
+
+# With Sequence Parallelism (Ulysses-SP, requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
+
+# With Ring-Attention (requires >= 2 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
+
+# Combined: Ulysses + Ring (requires >= 4 GPUs)
+vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
+```
+
+For more details on parallelism acceleration, see the [Parallelism Acceleration Guide](../../diffusion/parallelism_acceleration.md).
+
 ## API Calls
 
 ### Method 1: Using curl

examples/offline_inference/bagel/end2end.py

@@ -49,7 +49,14 @@ def parse_args():
     parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)")
     parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)")
     parser.add_argument(
-        "--negative-prompt", type=str, default=None, help="Negative prompt (not yet supported, reserved for future)"
+        "--negative-prompt", type=str, default=None, help="Negative prompt for CFG (default: empty prompt)"
+    )
+    parser.add_argument(
+        "--cfg-parallel-size",
+        type=int,
+        default=1,
+        choices=[1, 2, 3],
+        help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
     )
 
     args = parser.parse_args()
@@ -82,12 +89,14 @@ def main():
     from PIL import Image
 
     if args.modality == "img2img":
-        from PIL import Image
-
         from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
 
-        print("[Info] Running in img2img mode (Stage 1 only)")
-        client = OmniDiffusion(model=model_name)
+        print(f"[Info] Running in {args.modality} mode (Stage 1 only, cfg_parallel_size={args.cfg_parallel_size})")
+
+        client = OmniDiffusion(
+            model=model_name,
+            parallel_config={"cfg_parallel_size": args.cfg_parallel_size},
+        )
 
         if args.image_path:
             if os.path.exists(args.image_path):
@@ -162,6 +171,8 @@ def main():
                 # text2img
                 final_prompt_text = f"<|im_start|>{p}<|im_end|>"
                 prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
+                if args.negative_prompt is not None:
+                    prompt_dict["negative_prompt"] = args.negative_prompt
                 formatted_prompts.append(prompt_dict)
 
         params_list = omni.default_sampling_params_list
@@ -170,10 +181,13 @@ def main():
             if len(params_list) > 1:
                 diffusion_params = params_list[1]
                 diffusion_params.num_inference_steps = args.steps  # type: ignore
-                diffusion_params.extra_args = {  # type: ignore
+                extra = {
                     "cfg_text_scale": args.cfg_text_scale,
                     "cfg_img_scale": args.cfg_img_scale,
                 }
+                if args.negative_prompt is not None:
+                    extra["negative_prompt"] = args.negative_prompt
+                diffusion_params.extra_args = extra  # type: ignore
 
         omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 

examples/offline_inference/glm_image/README.md

@@ -40,8 +40,7 @@ GLM-Image uses a 2-stage pipeline:
 
 ```bash
 python end2end.py \
-    --model-path /path/to/glm-image \
-    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
     --prompt "A beautiful sunset over the ocean with sailing boats" \
     --height 1024 \
     --width 1024 \
@@ -52,8 +51,7 @@ python end2end.py \
 
 ```bash
 python end2end.py \
-    --model-path /path/to/glm-image \
-    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
     --prompt "Transform this scene into a winter wonderland" \
     --image input.png \
     --output output_i2i.png
@@ -64,7 +62,7 @@ python end2end.py \
 ```bash
 python end2end.py \
     --model-path /path/to/glm-image \
-    --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
+    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
     --prompt "A photorealistic cat sitting on a window sill" \
     --height 1024 \
     --width 1024 \

examples/offline_inference/glm_image/end2end.py

@@ -178,9 +178,6 @@ def main(args: argparse.Namespace) -> None:
     print("=" * 60)
 
     # Validate arguments
-    if not args.model_path:
-        raise ValueError("--model-path is required")
-
     if not args.prompt:
         raise ValueError("--prompt is required")
 
@@ -365,7 +362,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--model-path",
         type=str,
-        required=True,
+        default="zai-org/GLM-Image",
         help="Path to GLM-Image model directory or HuggingFace model ID",
     )
     parser.add_argument(

examples/offline_inference/helios/README.md

@@ -0,0 +1,193 @@
+# Helios Video Generation
+
+Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
+
+| Variant | Description | Key Features |
+|---------|-------------|--------------|
+| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
+| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
+| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
+
+## Setup
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+
+## Run Examples
+
+Get into the example folder:
+```bash
+cd examples/offline_inference/helios
+```
+
+### Text-to-Video (T2V)
+
+**Helios-Base** (Stage 1 only):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Base \
+    --sample-type t2v \
+    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
+    --guidance-scale 5.0 \
+    --output helios_t2v_base.mp4
+```
+
+**Helios-Mid** (Stage 2 + CFG-Zero*):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Mid \
+    --sample-type t2v \
+    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
+    --guidance-scale 5.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 20 20 20 \
+    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
+    --output helios_t2v_mid.mp4
+```
+
+**Helios-Distilled** (Stage 2 pyramid + DMD):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Distilled \
+    --sample-type t2v \
+    --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
+    --num-frames 240 \
+    --guidance-scale 1.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 2 2 2 \
+    --is-amplify-first-chunk \
+    --output helios_t2v_distilled.mp4
+```
+
+### Image-to-Video (I2V)
+
+**Helios-Base**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Base \
+    --sample-type i2v \
+    --image-path /path/to/image.jpg \
+    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
+    --guidance-scale 5.0 \
+    --output helios_i2v_base.mp4
+```
+
+**Helios-Mid** (Stage 2 + CFG-Zero*):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Mid \
+    --sample-type i2v \
+    --image-path /path/to/image.jpg \
+    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
+    --guidance-scale 5.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 20 20 20 \
+    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
+    --output helios_i2v_mid.mp4
+```
+
+**Helios-Distilled**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Distilled \
+    --sample-type i2v \
+    --image-path /path/to/image.jpg \
+    --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
+    --num-frames 240 \
+    --guidance-scale 1.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 2 2 2 \
+    --is-amplify-first-chunk \
+    --output helios_i2v_distilled.mp4
+```
+
+### Video-to-Video (V2V)
+
+**Helios-Base**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Base \
+    --sample-type v2v \
+    --video-path /path/to/video.mp4 \
+    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
+    --guidance-scale 5.0 \
+    --output helios_v2v_base.mp4
+```
+
+**Helios-Mid** (Stage 2 + CFG-Zero*):
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Mid \
+    --sample-type v2v \
+    --video-path /path/to/video.mp4 \
+    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
+    --guidance-scale 5.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 20 20 20 \
+    --use-cfg-zero-star --use-zero-init --zero-steps 1 \
+    --output helios_v2v_mid.mp4
+```
+
+**Helios-Distilled**:
+```bash
+python end2end.py \
+    --model BestWishYsh/Helios-Distilled \
+    --sample-type v2v \
+    --video-path /path/to/video.mp4 \
+    --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
+    --num-frames 240 \
+    --guidance-scale 1.0 \
+    --is-enable-stage2 \
+    --pyramid-num-inference-steps-list 2 2 2 \
+    --is-amplify-first-chunk \
+    --output helios_v2v_distilled.mp4
+```
+
+## Common Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
+| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
+| `--prompt` | — | Text prompt describing the video |
+| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
+| `--image-path` | — | Input image (required for `i2v`) |
+| `--video-path` | — | Input video (required for `v2v`) |
+| `--height` | `384` | Video height in pixels |
+| `--width` | `640` | Video width in pixels |
+| `--num-frames` | `99` | Number of output frames |
+| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
+| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
+| `--seed` | `42` | Random seed |
+| `--fps` | `16` | Output video frame rate |
+| `--output` | `helios_output.mp4` | Output file path |
+
+### Stage 2 / Pyramid Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
+| `--pyramid-num-stages` | `3` | Number of pyramid stages |
+| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
+| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
+
+### CFG-Zero* Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
+| `--use-zero-init` | off | Zero init for first steps |
+| `--zero-steps` | `1` | Number of zero-init steps |
+
+### Memory & Parallelism
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--vae-use-slicing` | off | Enable VAE slicing |
+| `--vae-use-tiling` | off | Enable VAE tiling |
+| `--enforce-eager` | off | Disable torch.compile |
+| `--enable-cpu-offload` | off | CPU offloading |
+| `--enable-layerwise-offload` | off | Layerwise offloading |
+| `--ulysses-degree` | `1` | Ulysses SP degree |
+| `--ring-degree` | `1` | Ring SP degree |
+| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
+| `--tensor-parallel-size` | `1` | Tensor parallelism size |

examples/offline_inference/helios/end2end.py

@@ -0,0 +1,343 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Helios video generation example supporting T2V, I2V, and V2V.
+
+Usage (T2V, Helios-Base, Stage 1 only):
+    python end2end.py \
+        --model /path/to/Helios-Base --sample-type t2v \
+        --prompt "A serene lakeside sunrise with mist over the water." \
+        --height 384 --width 640 --num-frames 99 \
+        --num-inference-steps 50 --guidance-scale 5.0
+
+Usage (I2V, Helios-Base):
+    python end2end.py \
+        --model /path/to/Helios-Base --sample-type i2v \
+        --image-path /path/to/image.jpg \
+        --prompt "Description of desired animation." \
+        --guidance-scale 5.0
+
+Usage (V2V, Helios-Base):
+    python end2end.py \
+        --model /path/to/Helios-Base --sample-type v2v \
+        --video-path /path/to/video.mp4 \
+        --prompt "Description of desired transformation." \
+        --guidance-scale 5.0
+
+Usage (Helios-Mid, Stage 2 + CFG-Zero*):
+    python end2end.py \
+        --model /path/to/Helios-Mid --sample-type t2v \
+        --prompt "A serene lakeside sunrise with mist over the water." \
+        --guidance-scale 5.0 \
+        --is-enable-stage2 \
+        --pyramid-num-inference-steps-list 20 20 20 \
+        --use-cfg-zero-star --use-zero-init --zero-steps 1
+
+Usage (Helios-Distilled, Stage 2 pyramid + DMD):
+    python end2end.py \
+        --model /path/to/Helios-Distilled --sample-type t2v \
+        --prompt "A serene lakeside sunrise with mist over the water." \
+        --guidance-scale 1.0 \
+        --is-enable-stage2 \
+        --pyramid-num-inference-steps-list 2 2 2 \
+        --is-amplify-first-chunk
+"""
+
+import argparse
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+
+def load_image_as_tensor(image_path: str, height: int, width: int) -> torch.Tensor:
+    """Load an image and return a (1, C, H, W) tensor normalized to [-1, 1]."""
+    from PIL import Image
+    from torchvision import transforms
+
+    image = Image.open(image_path).convert("RGB").resize((width, height))
+    transform = transforms.Compose(
+        [
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ]
+    )
+    return transform(image).unsqueeze(0)
+
+
+def load_video_as_tensor(video_path: str, height: int, width: int) -> torch.Tensor:
+    """Load a video and return a (1, C, T, H, W) tensor normalized to [-1, 1]."""
+    from torchvision import transforms
+    from torchvision.io import read_video
+
+    video_frames, _, _ = read_video(video_path, output_format="TCHW")
+    video_frames = video_frames.float() / 255.0
+
+    transform = transforms.Compose(
+        [
+            transforms.Resize((height, width), antialias=True),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ]
+    )
+    video_frames = torch.stack([transform(f) for f in video_frames])
+    return video_frames.permute(1, 0, 2, 3).unsqueeze(0)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate a video with Helios (T2V / I2V / V2V).")
+    parser.add_argument(
+        "--model",
+        default="BestWishYsh/Helios-Base",
+        help="Helios model ID or local path (e.g. Helios-Base, Helios-Mid, Helios-Distilled).",
+    )
+    parser.add_argument(
+        "--sample-type",
+        choices=["t2v", "i2v", "v2v"],
+        default="t2v",
+        help="Generation mode: t2v (text-to-video), i2v (image-to-video), v2v (video-to-video).",
+    )
+    parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
+    parser.add_argument(
+        "--negative-prompt",
+        default=(
+            "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, "
+            "images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, "
+            "incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, "
+            "misshapen limbs, fused fingers, still picture, messy background, three legs, many people "
+            "in the background, walking backwards"
+        ),
+        help="Negative prompt.",
+    )
+    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+    parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
+    parser.add_argument("--height", type=int, default=384, help="Video height.")
+    parser.add_argument("--width", type=int, default=640, help="Video width.")
+    parser.add_argument("--num-frames", type=int, default=99, help="Number of video frames.")
+    parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps (Stage 1 only).")
+    parser.add_argument("--output", type=str, default="helios_output.mp4", help="Output video path.")
+    parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
+
+    # I2V / V2V
+    parser.add_argument("--image-path", type=str, default=None, help="Input image path for I2V mode.")
+    parser.add_argument("--video-path", type=str, default=None, help="Input video path for V2V mode.")
+
+    # Stage 2 (pyramid multi-stage denoising)
+    parser.add_argument(
+        "--is-enable-stage2",
+        action="store_true",
+        help="Enable pyramid multi-stage denoising (Stage 2). Required for Helios-Distilled.",
+    )
+    parser.add_argument(
+        "--pyramid-num-stages",
+        type=int,
+        default=3,
+        help="Number of pyramid stages for Stage 2.",
+    )
+    parser.add_argument(
+        "--pyramid-num-inference-steps-list",
+        type=int,
+        nargs="+",
+        default=[10, 10, 10],
+        help="Inference steps per pyramid stage.",
+    )
+
+    # DMD
+    parser.add_argument(
+        "--is-amplify-first-chunk",
+        action="store_true",
+        help="Enable DMD amplification for the first chunk (Helios-Distilled).",
+    )
+
+    # CFG Zero Star
+    parser.add_argument(
+        "--use-cfg-zero-star",
+        action="store_true",
+        help="Enable CFG Zero Star guidance (recommended for Helios-Mid).",
+    )
+    parser.add_argument(
+        "--use-zero-init",
+        action="store_true",
+        help="Use zero initialization for the first denoising steps with CFG-Zero*.",
+    )
+    parser.add_argument(
+        "--zero-steps",
+        type=int,
+        default=1,
+        help="Number of initial denoising steps using zero prediction (default: 1).",
+    )
+
+    # Memory & parallelism
+    parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
+    parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
+    parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
+    parser.add_argument("--enable-cpu-offload", action="store_true", help="Enable CPU offloading.")
+    parser.add_argument("--enable-layerwise-offload", action="store_true", help="Enable layerwise offloading.")
+    parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
+    parser.add_argument("--ring-degree", type=int, default=1, help="Ring SP degree.")
+    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
+    parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.")
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+    )
+
+    omni = Omni(
+        model=args.model,
+        enable_layerwise_offload=args.enable_layerwise_offload,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
+        enable_cpu_offload=args.enable_cpu_offload,
+        parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
+    )
+
+    # Validate I2V / V2V arguments
+    if args.sample_type == "i2v" and not args.image_path:
+        raise ValueError("--image-path is required for I2V mode.")
+    if args.sample_type == "v2v" and not args.video_path:
+        raise ValueError("--video-path is required for V2V mode.")
+
+    # Build extra_args for Helios-specific parameters
+    extra_args = {}
+    if args.is_enable_stage2:
+        extra_args["is_enable_stage2"] = True
+        extra_args["pyramid_num_stages"] = args.pyramid_num_stages
+        extra_args["pyramid_num_inference_steps_list"] = args.pyramid_num_inference_steps_list
+    if args.is_amplify_first_chunk:
+        extra_args["is_amplify_first_chunk"] = True
+    if args.use_cfg_zero_star:
+        extra_args["use_cfg_zero_star"] = True
+    if args.use_zero_init:
+        extra_args["use_zero_init"] = True
+        extra_args["zero_steps"] = args.zero_steps
+
+    if args.sample_type == "i2v":
+        image_tensor = load_image_as_tensor(args.image_path, args.height, args.width)
+        extra_args["image"] = image_tensor
+    elif args.sample_type == "v2v":
+        video_tensor = load_video_as_tensor(args.video_path, args.height, args.width)
+        extra_args["video"] = video_tensor
+
+    # Print generation configuration
+    print(f"\n{'=' * 60}")
+    print("Helios Generation Configuration:")
+    print(f"  Model: {args.model}")
+    print(f"  Sample type: {args.sample_type.upper()}")
+    if args.sample_type == "i2v":
+        print(f"  Image: {args.image_path}")
+    elif args.sample_type == "v2v":
+        print(f"  Video: {args.video_path}")
+    print(f"  Prompt: {args.prompt[:80]}{'...' if len(args.prompt) > 80 else ''}")
+    print(f"  Video size: {args.width}x{args.height}, {args.num_frames} frames")
+    print(f"  Inference steps: {args.num_inference_steps}")
+    print(f"  Guidance scale: {args.guidance_scale}")
+    if args.is_enable_stage2:
+        print(f"  Stage 2: enabled (stages={args.pyramid_num_stages}, steps={args.pyramid_num_inference_steps_list})")
+        if args.is_amplify_first_chunk:
+            print("  DMD amplify first chunk: enabled")
+        if args.use_cfg_zero_star:
+            print(f"  CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
+    else:
+        if args.use_cfg_zero_star:
+            print(f"  CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
+        print("  Stage 2: disabled (Stage 1 only)")
+    print(f"{'=' * 60}\n")
+
+    generation_start = time.perf_counter()
+    frames = omni.generate(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+        },
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+            extra_args=extra_args,
+        ),
+    )
+    generation_end = time.perf_counter()
+    generation_time = generation_end - generation_start
+
+    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
+
+    # Extract video frames from OmniRequestOutput
+    if isinstance(frames, list) and len(frames) > 0:
+        first_item = frames[0]
+
+        if hasattr(first_item, "final_output_type"):
+            if first_item.final_output_type != "image":
+                raise ValueError(
+                    f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
+                )
+
+            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
+                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
+                    inner_output = first_item.request_output[0]
+                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
+                        frames = inner_output.images[0] if inner_output.images else None
+                        if frames is None:
+                            raise ValueError("No video frames found in output.")
+            elif hasattr(first_item, "images") and first_item.images:
+                frames = first_item.images
+            else:
+                raise ValueError("No video frames found in OmniRequestOutput.")
+
+    # Unwrap batch list from postprocess_video: [numpy(T,H,W,C)] -> numpy(T,H,W,C)
+    if isinstance(frames, list) and len(frames) > 0 and isinstance(frames[0], np.ndarray):
+        frames = frames[0]
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        from diffusers.utils import export_to_video
+    except ImportError:
+        raise ImportError("diffusers is required for export_to_video.")
+
+    if isinstance(frames, torch.Tensor):
+        video_tensor = frames.detach().cpu()
+        if video_tensor.dim() == 5:
+            if video_tensor.shape[1] in (3, 4):
+                video_tensor = video_tensor[0].permute(1, 2, 3, 0)
+            else:
+                video_tensor = video_tensor[0]
+        elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
+            video_tensor = video_tensor.permute(1, 2, 3, 0)
+        if video_tensor.is_floating_point():
+            video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
+        video_array = video_tensor.float().numpy()
+    elif isinstance(frames, np.ndarray):
+        video_array = frames
+    else:
+        video_array = frames
+
+    if isinstance(video_array, np.ndarray) and video_array.ndim == 5:
+        video_array = video_array[0]
+
+    export_to_video(video_array, str(output_path), fps=args.fps)
+    print(f"Saved generated video to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/hunyuan_image3/README.md

@@ -0,0 +1,25 @@
+# HunyuanImage-3.0 Image-to-Text Inference
+
+This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni.
+
+## Local CLI Usage
+
+Download the example image:
+
+```bash
+wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+```
+
+Run example:
+
+```bash
+python image_to_text.py \
+  --image cherry_blossom.jpg \
+  --prompt "<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
+```
+
+Key arguments:
+
+- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional).
+- `--image`: Path to input image (required).
+- `--prompt`: Text description used to guide image understanding (required).

examples/offline_inference/hunyuan_image3/image_to_text.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import os
+
+from PIL import Image
+
+from vllm_omni.entrypoints.omni import Omni
+
+"""
+The tencent/HunyuanImage-3.0-Instruct base model is built on the Hunyuan v1 architecture, specifically the tencent/Hunyuan-A13B-Instruct model. It utilizes two tokenizer delimiter templates:
+
+1) Pretrained template (default for gen_text mode), which concatenates system, image
+   tokens, and user question WITHOUT role delimiters:
+"<|startoftext|>{system_prompt}{image_tokens}{user_question}"
+
+   Example (before image token expansion):
+"<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
+
+2) Instruct template, which uses explicit role prefixes and separators.
+"""
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate text from image using HunyuanImage-3.0-Instruct.")
+    parser.add_argument(
+        "--model",
+        default="tencent/HunyuanImage-3.0-Instruct",
+        help="Model name or local path.",
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        required=True,
+        help="Path to input image file (PNG, JPG, etc.).",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        required=True,
+        help="Pretrain template prompt: <|startoftext|>{system}<img>{question}",
+    )
+    return parser.parse_args()
+
+
+def load_image(image_path: str) -> Image.Image:
+    """Load an image from file path."""
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+    return Image.open(image_path).convert("RGB")
+
+
+def main(args: argparse.Namespace) -> None:
+    omni = Omni(model=args.model)
+
+    prompt_dict = {
+        "prompt": args.prompt,
+        "modalities": ["text"],
+    }
+
+    # Add image input if provided
+    if args.image:
+        if not os.path.exists(args.image):
+            raise FileNotFoundError(f"Input image not found: {args.image}")
+
+        input_image = load_image(args.image)
+        prompt_dict["multi_modal_data"] = {"image": input_image}
+
+    prompts = [prompt_dict]
+    omni_outputs = omni.generate(prompts=prompts)
+
+    prompt_text = omni_outputs[0].request_output[0].prompt
+    generated_text = omni_outputs[0].request_output[0].outputs[0].text
+    print(f"Prompt: {prompt_text}")
+    print(f"Text: {generated_text}")
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/offline_inference/image_to_video/README.md

@@ -65,6 +65,7 @@ Key arguments:
 - `--vae-use-slicing`: Enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: Enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 - `--use-hsdp`: Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.
 - `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.

examples/offline_inference/image_to_video/image_to_video.py

@@ -2,11 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Image-to-Video generation example using Wan2.2 I2V/TI2V models.
+Image-to-Video generation example using Wan2.2 I2V/TI2V models or LTX2.
 
 Supports:
 - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
 - Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
+- LTX2 image-to-video pipeline
 
 Usage:
     # I2V-A14B (MoE)
@@ -16,10 +17,18 @@ Usage:
     # TI2V-5B (unified)
     python image_to_video.py --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
         --image input.jpg --prompt "A cat playing with yarn"
+
+    # LTX2 image-to-video
+    python image_to_video.py --model /path/to/LTX-2 \
+        --model-class-name LTX2ImageToVideoPipeline \
+        --image input.jpg --prompt "A cinematic dolly shot of a boat" \
+        --num-frames 121 --num_inference_steps 40 --guidance_scale 4.0 \
+        --frame-rate 24 --fps 24 --output ltx2_i2v.mp4
 """
 
 import argparse
 import os
+import time
 from pathlib import Path
 
 import numpy as np
@@ -34,12 +43,17 @@ from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 I2V/TI2V.")
+    parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 or LTX2.")
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
         help="Diffusers Wan2.2 I2V model ID or local path.",
     )
+    parser.add_argument(
+        "--model-class-name",
+        default=None,
+        help="Override model class name (e.g., LTX2ImageToVideoPipeline).",
+    )
     parser.add_argument("--image", required=True, help="Path to input image.")
     parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
     parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
@@ -55,11 +69,17 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
     parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps.")
     parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
+    parser.add_argument(
+        "--frame-rate",
+        type=float,
+        default=None,
+        help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
+    )
     parser.add_argument(
         "--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
     )
     parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
-    parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
+    parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
     parser.add_argument(
         "--vae-use-slicing",
         action="store_true",
@@ -105,11 +125,34 @@ def parse_args() -> argparse.Namespace:
         choices=[1, 2],
         help="Number of GPUs used for classifier free guidance parallel size.",
     )
+    parser.add_argument(
+        "--vae-patch-parallel-size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for VAE patch/tile parallelism (decode).",
+    )
     parser.add_argument(
         "--enforce-eager",
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=24000,
+        help="Sample rate for audio output when saved (default: 24000 for LTX2).",
+    )
+    parser.add_argument(
+        "--cache-backend",
+        type=str,
+        default=None,
+        choices=["cache_dit", "tea_cache"],
+        help=(
+            "Cache backend to use for acceleration. "
+            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
+            "Default: None (no cache acceleration)."
+        ),
+    )
     parser.add_argument(
         "--use-hsdp",
         action="store_true",
@@ -136,10 +179,13 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> tuple[int, int]:
+def calculate_dimensions(
+    image: PIL.Image.Image,
+    max_area: int = 480 * 832,
+    mod_value: int = 16,
+) -> tuple[int, int]:
     """Calculate output dimensions maintaining aspect ratio."""
     aspect_ratio = image.height / image.width
-    mod_value = 16  # Must be divisible by 16
 
     height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
     width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
@@ -150,22 +196,66 @@ def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> t
 def main():
     args = parse_args()
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    model_name = str(args.model).lower() if args.model is not None else ""
+    model_class_name = args.model_class_name
+    is_ltx2 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower())
+    if model_class_name is None and is_ltx2:
+        model_class_name = "LTX2ImageToVideoPipeline"
 
     # Load input image
     image = PIL.Image.open(args.image).convert("RGB")
 
+    fps = args.fps if args.fps is not None else (24 if is_ltx2 else 16)
+    frame_rate = args.frame_rate if args.frame_rate is not None else float(fps)
+    guidance_scale = args.guidance_scale if args.guidance_scale is not None else (4.0 if is_ltx2 else 5.0)
+    num_frames = args.num_frames if args.num_frames is not None else (121 if is_ltx2 else 81)
+    num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else (40 if is_ltx2 else 50)
+
     # Calculate dimensions if not provided
     height = args.height
     width = args.width
     if height is None or width is None:
-        # Default to 480P area for I2V
-        calc_height, calc_width = calculate_dimensions(image, max_area=480 * 832)
+        # Default to 480P area for Wan2.2 I2V, 512x768 area for LTX2
+        max_area = 512 * 768 if is_ltx2 else 480 * 832
+        mod_value = 32 if is_ltx2 else 16
+        calc_height, calc_width = calculate_dimensions(image, max_area=max_area, mod_value=mod_value)
         height = height or calc_height
         width = width or calc_width
 
     # Resize image to target dimensions
     image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
 
+    # Configure cache based on backend type
+    cache_config = None
+    if args.cache_backend == "cache_dit":
+        if is_ltx2:
+            cache_config = {
+                "Fn_compute_blocks": 2,
+                "Bn_compute_blocks": 0,
+                "max_warmup_steps": 8,
+                "residual_diff_threshold": 0.12,
+                "max_continuous_cached_steps": 1,
+                "max_cached_steps": 20,
+                "enable_taylorseer": False,
+                "scm_steps_mask_policy": None,
+            }
+        else:
+            cache_config = {
+                "Fn_compute_blocks": 1,
+                "Bn_compute_blocks": 0,
+                "max_warmup_steps": 4,
+                "residual_diff_threshold": 0.24,
+                "max_continuous_cached_steps": 3,
+                "enable_taylorseer": False,
+                "taylorseer_order": 1,
+                "scm_steps_mask_policy": None,
+                "scm_steps_policy": "dynamic",
+            }
+    elif args.cache_backend == "tea_cache":
+        cache_config = {
+            "rel_l1_thresh": 0.2,
+        }
+
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
     parallel_config = DiffusionParallelConfig(
@@ -173,6 +263,7 @@ def main():
         ring_degree=args.ring_degree,
         cfg_parallel_size=args.cfg_parallel_size,
         tensor_parallel_size=args.tensor_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
         use_hsdp=args.use_hsdp,
         hsdp_shard_size=args.hsdp_shard_size,
         hsdp_replicate_size=args.hsdp_replicate_size,
@@ -187,6 +278,9 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        model_class_name=model_class_name,
+        cache_backend=args.cache_backend,
+        cache_config=cache_config,
     )
 
     if profiler_enabled:
@@ -199,10 +293,14 @@ def main():
     print(f"  Model: {args.model}")
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Frames: {args.num_frames}")
-    print(f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}")
+    print(
+        f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
+        f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"
+    )
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
 
+    generation_start = time.perf_counter()
     # omni.generate() returns Generator[OmniRequestOutput, None, None]
     frames = omni.generate(
         {
@@ -214,38 +312,69 @@ def main():
             height=height,
             width=width,
             generator=generator,
-            guidance_scale=args.guidance_scale,
+            guidance_scale=guidance_scale,
             guidance_scale_2=args.guidance_scale_high,
-            num_inference_steps=args.num_inference_steps,
-            num_frames=args.num_frames,
+            num_inference_steps=num_inference_steps,
+            num_frames=num_frames,
+            frame_rate=frame_rate,
         ),
     )
+    generation_end = time.perf_counter()
+    generation_time = generation_end - generation_start
 
-    # Extract video frames from OmniRequestOutput
-    if isinstance(frames, list) and len(frames) > 0:
-        first_item = frames[0]
+    # Print profiling results
+    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
-        # Check if it's an OmniRequestOutput
-        if hasattr(first_item, "final_output_type"):
-            if first_item.final_output_type != "image":
-                raise ValueError(
-                    f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
-                )
+    audio = None
+    if isinstance(frames, list):
+        frames = frames[0] if frames else None
 
-            # Pipeline mode: extract from nested request_output
-            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
-                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
-                    inner_output = first_item.request_output[0]
-                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
-                        frames = inner_output.images[0] if inner_output.images else None
-                        if frames is None:
-                            raise ValueError("No video frames found in output.")
-            # Diffusion mode: use direct images field
-            elif hasattr(first_item, "images") and first_item.images:
-                frames = first_item.images
+    if isinstance(frames, OmniRequestOutput):
+        if frames.final_output_type != "image":
+            raise ValueError(
+                f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
+            )
+        if frames.multimodal_output and "audio" in frames.multimodal_output:
+            audio = frames.multimodal_output["audio"]
+        if frames.is_pipeline_output and frames.request_output is not None:
+            inner_output = frames.request_output
+            if isinstance(inner_output, list):
+                inner_output = inner_output[0] if inner_output else None
+            if isinstance(inner_output, OmniRequestOutput):
+                if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
+                    audio = inner_output.multimodal_output["audio"]
+                frames = inner_output
+        if isinstance(frames, OmniRequestOutput):
+            if frames.images:
+                if len(frames.images) == 1 and isinstance(frames.images[0], tuple) and len(frames.images[0]) == 2:
+                    frames, audio = frames.images[0]
+                elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
+                    audio = frames.images[0].get("audio")
+                    frames = frames.images[0].get("frames") or frames.images[0].get("video")
+                else:
+                    frames = frames.images
             else:
                 raise ValueError("No video frames found in OmniRequestOutput.")
 
+    if isinstance(frames, list) and frames:
+        first_item = frames[0]
+        if isinstance(first_item, tuple) and len(first_item) == 2:
+            frames, audio = first_item
+        elif isinstance(first_item, dict):
+            audio = first_item.get("audio")
+            frames = first_item.get("frames") or first_item.get("video")
+        elif isinstance(first_item, list):
+            frames = first_item
+
+    if isinstance(frames, tuple) and len(frames) == 2:
+        frames, audio = frames
+    elif isinstance(frames, dict):
+        audio = frames.get("audio")
+        frames = frames.get("frames") or frames.get("video")
+
+    if frames is None:
+        raise ValueError("No video frames found in output.")
+
     output_path = Path(args.output)
     output_path.parent.mkdir(parents=True, exist_ok=True)
 
@@ -254,32 +383,128 @@ def main():
     except ImportError:
         raise ImportError("diffusers is required for export_to_video.")
 
-    # frames may be np.ndarray (preferred) or torch.Tensor
+    def _normalize_frame(frame):
+        if isinstance(frame, torch.Tensor):
+            frame_tensor = frame.detach().cpu()
+            if frame_tensor.dim() == 4 and frame_tensor.shape[0] == 1:
+                frame_tensor = frame_tensor[0]
+            if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
+                frame_tensor = frame_tensor.permute(1, 2, 0)
+            if frame_tensor.is_floating_point():
+                frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
+            return frame_tensor.float().numpy()
+        if isinstance(frame, np.ndarray):
+            frame_array = frame
+            if frame_array.ndim == 4 and frame_array.shape[0] == 1:
+                frame_array = frame_array[0]
+            if np.issubdtype(frame_array.dtype, np.integer):
+                frame_array = frame_array.astype(np.float32) / 255.0
+            return frame_array
+        try:
+            from PIL import Image
+        except ImportError:
+            Image = None
+        if Image is not None and isinstance(frame, Image.Image):
+            return np.asarray(frame).astype(np.float32) / 255.0
+        return frame
+
+    def _ensure_frame_list(video_array):
+        if isinstance(video_array, list):
+            if len(video_array) == 0:
+                return video_array
+            first_item = video_array[0]
+            if isinstance(first_item, np.ndarray):
+                if first_item.ndim == 5:
+                    return list(first_item[0])
+                if first_item.ndim == 4:
+                    if len(video_array) == 1:
+                        return list(first_item)
+                    return list(first_item)
+                if first_item.ndim == 3:
+                    return video_array
+            return video_array
+        if isinstance(video_array, np.ndarray):
+            if video_array.ndim == 5:
+                return list(video_array[0])
+            if video_array.ndim == 4:
+                return list(video_array)
+            if video_array.ndim == 3:
+                return [video_array]
+        return video_array
+
+    # frames may be np.ndarray, torch.Tensor, or list of tensors/arrays/images
     # export_to_video expects a list of frames with values in [0, 1]
     if isinstance(frames, torch.Tensor):
         video_tensor = frames.detach().cpu()
         if video_tensor.dim() == 5:
-            # [B, C, F, H, W] or [B, F, H, W, C]
             if video_tensor.shape[1] in (3, 4):
                 video_tensor = video_tensor[0].permute(1, 2, 3, 0)
             else:
                 video_tensor = video_tensor[0]
         elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
             video_tensor = video_tensor.permute(1, 2, 3, 0)
-        # If float, assume [-1,1] and normalize to [0,1]
         if video_tensor.is_floating_point():
             video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
         video_array = video_tensor.float().numpy()
+    elif isinstance(frames, np.ndarray):
+        video_array = frames
+        if video_array.ndim == 5:
+            video_array = video_array[0]
+        if np.issubdtype(video_array.dtype, np.integer):
+            video_array = video_array.astype(np.float32) / 255.0
+    elif isinstance(frames, list):
+        if len(frames) == 0:
+            raise ValueError("No video frames found in output.")
+        video_array = [_normalize_frame(frame) for frame in frames]
     else:
         video_array = frames
-        if hasattr(video_array, "shape") and video_array.ndim == 5:
-            video_array = video_array[0]
 
-    # Convert 4D array (frames, H, W, C) to list of frames for export_to_video
-    if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
-        video_array = list(video_array)
+    video_array = _ensure_frame_list(video_array)
 
-    export_to_video(video_array, str(output_path), fps=args.fps)
+    use_ltx2_export = is_ltx2
+    encode_video = None
+    if use_ltx2_export:
+        try:
+            from diffusers.pipelines.ltx2.export_utils import encode_video
+        except ImportError:
+            encode_video = None
+
+    if use_ltx2_export and encode_video is not None:
+        if isinstance(video_array, list):
+            frames_np = np.stack(video_array, axis=0)
+        elif isinstance(video_array, np.ndarray):
+            frames_np = video_array
+        else:
+            frames_np = np.asarray(video_array)
+
+        if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
+            frames_np = frames_np[..., :3]
+
+        frames_np = np.clip(frames_np, 0.0, 1.0)
+        frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
+        video_tensor = torch.from_numpy(frames_u8)
+
+        audio_out = None
+        if audio is not None:
+            if isinstance(audio, list):
+                audio = audio[0] if audio else None
+            if isinstance(audio, np.ndarray):
+                audio = torch.from_numpy(audio)
+            if isinstance(audio, torch.Tensor):
+                audio_out = audio
+                if audio_out.dim() > 1:
+                    audio_out = audio_out[0]
+                audio_out = audio_out.float().cpu()
+
+        encode_video(
+            video_tensor,
+            fps=fps,
+            audio=audio_out,
+            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
+            output_path=str(output_path),
+        )
+    else:
+        export_to_video(video_array, str(output_path), fps=fps)
     print(f"Saved generated video to {output_path}")
 
     if profiler_enabled:

examples/offline_inference/mammothmodal2_preview/README.md

@@ -0,0 +1,32 @@
+# MammothModa2-Preview
+
+## Run examples (MammothModa2-Preview)
+
+Download model
+```bash
+hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
+```
+
+### Text-to-Image (T2I)
+
+```bash
+python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py \
+  --model ./MammothModa2-Preview \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
+  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+  --height 1024 \
+  --width 1024 \
+  --num-inference-steps 50 \
+  --text-guidance-scale 4.0 \
+  --out output.png
+```
+
+### Image Summary
+
+```bash
+python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
+  --model ./MammothModa2-Preview \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+  --question "Summarize this image." \
+  --image ./image.png
+```

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -0,0 +1,115 @@
+"""
+Offline inference example: MammothModa2 image summarization (single AR stage).
+
+Example:
+  uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
+    --model path/to/MammothModa2-Preview \
+    --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+    --image /path/to/input.jpg \
+    --question "Please summarize the content of this image."
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+from PIL import Image
+from vllm import SamplingParams
+from vllm.multimodal.image import convert_image_mode
+
+from vllm_omni import Omni
+
+DEFAULT_SYSTEM = "You are a helpful assistant."
+DEFAULT_QUESTION = "Please summarize the content of this image."
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
+    parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
+    parser.add_argument(
+        "--stage-config", type=str, required=True, help="Path to stage config yaml (single-stage AR->text)."
+    )
+    parser.add_argument("--image", type=str, required=True, help="Path to input image.")
+    parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
+    parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=512,
+        help="Max new tokens to generate.",
+    )
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--trust-remote-code", action="store_true")
+    return parser.parse_args()
+
+
+def build_prompt(system: str, question: str) -> str:
+    return (
+        f"<|im_start|>system\n{system}<|im_end|>\n"
+        "<|im_start|>user\n"
+        "<|vision_start|><|image_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+
+def main() -> None:
+    args = parse_args()
+
+    if not os.path.exists(args.image):
+        raise FileNotFoundError(f"Image file not found: {args.image}")
+
+    pil_image = Image.open(args.image)
+    image_data = convert_image_mode(pil_image, "RGB")
+    prompt = build_prompt(args.system, args.question)
+
+    omni = Omni(
+        model=args.model,
+        stage_configs_path=args.stage_config,
+        trust_remote_code=args.trust_remote_code,
+    )
+    try:
+        sp = SamplingParams(
+            temperature=float(args.temperature),
+            top_p=float(args.top_p),
+            top_k=-1,
+            max_tokens=int(args.max_tokens),
+            seed=int(args.seed),
+            detokenize=True,
+        )
+        # NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
+        # Consume it inside the try block so the worker isn't closed early.
+        outputs = list(
+            omni.generate(
+                [
+                    {
+                        "prompt": prompt,
+                        "multi_modal_data": {"image": image_data},
+                        "additional_information": {"omni_task": ["chat"]},
+                    }
+                ],
+                [sp],
+            )
+        )
+    finally:
+        omni.close()
+
+    lines: list[str] = []
+    for stage_outputs in outputs:
+        req_outputs = getattr(stage_outputs, "request_output", stage_outputs)
+        req_outputs = req_outputs if isinstance(req_outputs, list) else [req_outputs]
+        for ro in req_outputs:
+            text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
+            lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
+            lines.append("answer:\n")
+            lines.append(text.strip() + "\n")
+            lines.append("\n")
+
+    print("\n".join(lines))
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -0,0 +1,251 @@
+"""
+Offline inference example for MammothModa2 Text-to-Image (T2I) generation.
+This script uses the vllm_omni.Omni pipeline with a multi-stage configuration.
+
+Workflow:
+1. Stage 0 (AR): Generates visual tokens and their corresponding hidden states.
+2. Stage 1 (DiT): Consumes the hidden states as conditions to perform diffusion
+   and VAE decoding to produce the final image.
+
+Example Usage:
+    uv run python examples/offline_inference/run_mammothmoda2_t2i.py \
+        --model path/to/MammothModa2-Preview \
+        --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
+        --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+        --out output.png
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+from pathlib import Path
+from typing import NamedTuple
+
+import torch
+from PIL import Image
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni import Omni
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+_PATCH_SIZE = 16  # AR image grid patch size (pixels per token)
+
+
+class T2IGenConfig(NamedTuple):
+    eol_token_id: int
+    visual_token_start_id: int
+    visual_token_end_id: int
+    top_k: int  # AR sampling top-k (covers the full visual generation vocabulary)
+    # Qwen2.5-VL special vision tokens: <|image_pad|>, <|video_pad|>, <|vision_start|>, <|vision_end|>
+    visual_ids: list[int]
+
+
+def load_t2i_generation_config(model_dir: str) -> T2IGenConfig:
+    """Load T2I token IDs from t2i_generation_config.json and config.json.
+
+    Supports both local directory paths and HuggingFace Hub model IDs.
+    """
+    model_path = Path(model_dir)
+
+    def _read_json(filename: str) -> dict:
+        local = model_path / filename
+        if local.exists():
+            with local.open(encoding="utf-8") as f:
+                return json.load(f)
+        # Fall back to HuggingFace Hub when model_dir is a repo ID.
+        try:
+            from huggingface_hub import hf_hub_download
+        except ImportError as exc:
+            raise ImportError(
+                "huggingface_hub is required to load configs from HF Hub. Install it with: pip install huggingface_hub"
+            ) from exc
+        cached = hf_hub_download(repo_id=model_dir, filename=filename)
+        with open(cached, encoding="utf-8") as f:
+            return json.load(f)
+
+    gen_cfg = _read_json("t2i_generation_config.json")
+    llm_cfg = _read_json("config.json").get("llm_config", {})
+
+    return T2IGenConfig(
+        eol_token_id=int(gen_cfg["eol_token_id"]),
+        visual_token_start_id=int(gen_cfg["visual_token_start_id"]),
+        visual_token_end_id=int(gen_cfg["visual_token_end_id"]),
+        top_k=int(gen_cfg["top_k"]),
+        visual_ids=[
+            int(llm_cfg["image_token_id"]),
+            int(llm_cfg["video_token_id"]),
+            int(llm_cfg["vision_start_token_id"]),
+            int(llm_cfg["vision_end_token_id"]),
+        ],
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Run MammothModa2 T2I (AR -> DiT) with vLLM-Omni.")
+    p.add_argument("--model", type=str, required=True, help="Path to the model directory.")
+    p.add_argument("--stage-config", type=str, required=True, help="Path to the multi-stage YAML configuration.")
+    p.add_argument(
+        "--prompt",
+        type=str,
+        action="append",
+        default=None,
+        help=(
+            "Text prompt for image generation. Can be provided multiple times "
+            "to generate multiple images with shared height/width/CFG settings."
+        ),
+    )
+    p.add_argument("--height", type=int, default=1024, help="Output image height (must be a multiple of 16).")
+    p.add_argument("--width", type=int, default=1024, help="Output image width (must be a multiple of 16).")
+    p.add_argument("--num-inference-steps", type=int, default=50, help="Number of diffusion steps for the DiT stage.")
+    p.add_argument(
+        "--text-guidance-scale", type=float, default=9.0, help="Classifier-Free Guidance (CFG) scale for DiT."
+    )
+    p.add_argument(
+        "--cfg-range",
+        type=float,
+        nargs=2,
+        default=(0.0, 1.0),
+        help="Relative step range [start, end] where CFG is active.",
+    )
+    p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
+    p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
+    args = p.parse_args()
+    if not args.prompt:
+        args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
+    return args
+
+
+def tensor_to_pil(image: torch.Tensor) -> Image.Image:
+    """Convert a normalized torch tensor [-1, 1] to a PIL Image."""
+    if image.ndim == 4:
+        image = image[0]
+    image = image.detach().to("cpu")
+    image = (image / 2 + 0.5).clamp(0, 1)
+    image = (image * 255).to(torch.uint8)
+    image = image.permute(1, 2, 0).contiguous().numpy()
+    return Image.fromarray(image)
+
+
+def _format_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str:
+    """Build the AR-stage prompt string including the image grid header."""
+    return (
+        "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
+        f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+        f"<|image start|>{ar_width}*{ar_height}<|image token|>"
+    )
+
+
+def _collect_images(outputs: list) -> list[torch.Tensor]:
+    """Extract all image tensors produced by the final (DiT) stage."""
+    images: list[torch.Tensor] = []
+    for out in outputs:
+        ro_list = getattr(out, "request_output", out)
+        if not isinstance(ro_list, list):
+            ro_list = [ro_list]
+        for ro_item in ro_list:
+            for completion in getattr(ro_item, "outputs", None) or []:
+                mm = getattr(completion, "multimodal_output", None)
+                if not isinstance(mm, dict) or "image" not in mm:
+                    raise RuntimeError(f"Missing image in multimodal output: {mm}")
+                payload = mm["image"]
+                for tensor in payload if isinstance(payload, list) else [payload]:
+                    if not isinstance(tensor, torch.Tensor):
+                        raise TypeError(f"Expected image tensor, got {type(tensor)}")
+                    images.append(tensor)
+    return images
+
+
+def _save_images(images: list[torch.Tensor], out_path: str) -> list[str]:
+    """Save image tensors to disk.
+
+    Single image: written to *out_path* exactly.
+    Multiple images: suffixed as ``<base>_0<ext>``, ``<base>_1<ext>``, …
+    """
+    if not images:
+        raise RuntimeError("No images to save.")
+    base, ext = os.path.splitext(out_path)
+    ext = ext or ".png"
+    paths = []
+    for i, tensor in enumerate(images):
+        path = out_path if len(images) == 1 else f"{base}_{i}{ext}"
+        tensor_to_pil(tensor).save(path)
+        paths.append(path)
+    return paths
+
+
+def main() -> None:
+    args = parse_args()
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+
+    if args.height <= 0 or args.width <= 0:
+        raise ValueError(f"Height and width must be positive, got {args.height}x{args.width}")
+    if args.height % _PATCH_SIZE != 0 or args.width % _PATCH_SIZE != 0:
+        raise ValueError(f"Height and width must be multiples of {_PATCH_SIZE}, got {args.height}x{args.width}")
+
+    ar_height = args.height // _PATCH_SIZE
+    ar_width = args.width // _PATCH_SIZE
+    gen_cfg = load_t2i_generation_config(args.model)
+    expected_grid_tokens = ar_height * (ar_width + 1)
+
+    logger.info("Initializing Omni pipeline...")
+    omni = Omni(model=args.model, stage_configs_path=args.stage_config, trust_remote_code=args.trust_remote_code)
+    try:
+        ar_sampling = SamplingParams(
+            temperature=1.0,
+            top_p=1.0,
+            top_k=gen_cfg.top_k,
+            max_tokens=max(1, expected_grid_tokens + 1),  # +1 for hidden state of eoi
+            detokenize=False,
+        )
+        dit_sampling = SamplingParams(
+            temperature=0.0,
+            top_p=1.0,
+            top_k=-1,
+            max_tokens=1,
+            detokenize=False,
+        )
+
+        additional_information = {
+            "omni_task": ["t2i"],
+            "ar_width": [ar_width],
+            "ar_height": [ar_height],
+            "eol_token_id": [gen_cfg.eol_token_id],
+            "visual_token_start_id": [gen_cfg.visual_token_start_id],
+            "visual_token_end_id": [gen_cfg.visual_token_end_id],
+            "image_height": [args.height],
+            "image_width": [args.width],
+            "num_inference_steps": [args.num_inference_steps],
+            "text_guidance_scale": [args.text_guidance_scale],
+            "cfg_range": [args.cfg_range[0], args.cfg_range[1]],
+            "visual_ids": gen_cfg.visual_ids,
+        }
+        inputs = [
+            {
+                "prompt": _format_prompt(p, ar_width, ar_height),
+                "additional_information": dict(additional_information),
+            }
+            for p in args.prompt
+        ]
+
+        logger.info("Starting generation...")
+        # omni.generate() returns a Generator; consume it to run the full pipeline.
+        outputs = list(omni.generate(inputs, [ar_sampling, dit_sampling]))
+
+        logger.info("Post-processing and saving image(s)...")
+        for path in _save_images(_collect_images(outputs), args.out):
+            logger.info(f"Saved: {path}")
+    finally:
+        omni.close()
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/qwen3_omni/README.md

@@ -64,6 +64,51 @@ If media file paths are not provided, the script will use default assets. Suppor
 - `multi_audios`: Multiple audio inputs
 - `mixed_modalities`: Combination of video, image, and audio inputs
 
+### Async-chunk (offline)
+
+For true stage-level concurrency -- where downstream stages (Talker, Code2Wav)
+start **before** the upstream stage (Thinker) finishes -- use the async_chunk
+example. This requires:
+
+1. A stage config YAML with ``async_chunk: true`` (e.g.
+   ``qwen3_omni_moe_async_chunk.yaml``).
+2. Hardware that matches the config (e.g. 2x H100 for the default 3-stage
+   config).
+
+The async_chunk example uses ``AsyncOmni`` instead of the synchronous ``Omni``
+class, which enables the async orchestrator to receive stage-0 intermediate
+outputs and trigger downstream stages early. Chunk data flows directly between
+stage workers via the in-worker ``OmniChunkTransferAdapter`` / connector,
+**not** through the orchestrator.
+
+#### Single prompt
+```bash
+cd examples/offline_inference/qwen3_omni
+bash run_single_prompt_async_chunk.sh
+```
+
+#### Multiple prompts with concurrency control
+```bash
+bash run_multiple_prompts_async_chunk.sh --max-in-flight 4
+```
+
+#### Text-only output (skip audio generation)
+```bash
+python end2end_async_chunk.py --query-type text --modalities text
+```
+
+#### Custom stage config
+```bash
+python end2end_async_chunk.py \
+    --query-type use_audio \
+    --stage-configs-path /path/to/your_async_chunk.yaml
+```
+
+> **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the
+> recommended entry point for non-async-chunk workflows. Only use the
+> async_chunk example when you need the stage-level concurrency semantics
+> described in PR #962 / #1151.
+
 ### FAQ
 
 If you encounter error about backend of librosa, try to install ffmpeg with command below.

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -0,0 +1,614 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline inference with async_chunk enabled via AsyncOmni.
+
+This script uses AsyncOmni (the async orchestrator) to run offline inference
+with async_chunk semantics: downstream stages (Talker, Code2Wav) start
+*before* upstream stages finish, consuming chunks as they arrive via
+the in-worker OmniChunkTransferAdapter / connector.
+
+Compared to the synchronous ``end2end.py`` (which uses ``Omni``), this
+entry point achieves true stage-level concurrency -- stage-1/2 are
+actively processing while stage-0 is still generating.
+
+Usage
+-----
+    python end2end_async_chunk.py --query-type use_audio \
+        --stage-configs-path <path-to-async-chunk-yaml>
+
+See ``--help`` for all options.
+"""
+
+import asyncio
+import logging
+import os
+import time
+import uuid
+from typing import NamedTuple
+
+import numpy as np
+import soundfile as sf
+import torch
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import librosa
+from PIL import Image
+from vllm import SamplingParams
+from vllm.assets.audio import AudioAsset
+from vllm.assets.image import ImageAsset
+from vllm.assets.video import VideoAsset, video_to_ndarrays
+from vllm.multimodal.image import convert_image_mode
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Query builders (reuse the patterns from end2end.py)
+# ---------------------------------------------------------------------------
+
+default_system = (
+    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
+    "Group, capable of perceiving auditory and visual inputs, as well as "
+    "generating text and speech."
+)
+
+
+class QueryResult(NamedTuple):
+    inputs: dict
+    limit_mm_per_prompt: dict[str, int]
+
+
+def get_text_query(question: str = None) -> QueryResult:
+    if question is None:
+        question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+    return QueryResult(inputs={"prompt": prompt}, limit_mm_per_prompt={})
+
+
+def get_audio_query(
+    question: str = None,
+    audio_path: str | None = None,
+    sampling_rate: int = 16000,
+) -> QueryResult:
+    if question is None:
+        question = "What is the content of this audio?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|audio_start|><|audio_pad|><|audio_end|>"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+    if audio_path:
+        if not os.path.exists(audio_path):
+            raise FileNotFoundError(f"Audio file not found: {audio_path}")
+        audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {"audio": audio_data},
+        },
+        limit_mm_per_prompt={"audio": 1},
+    )
+
+
+def get_image_query(question: str = None, image_path: str | None = None) -> QueryResult:
+    if question is None:
+        question = "What is the content of this image?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+    if image_path:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Image file not found: {image_path}")
+        pil_image = Image.open(image_path)
+        image_data = convert_image_mode(pil_image, "RGB")
+    else:
+        image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {"image": image_data},
+        },
+        limit_mm_per_prompt={"image": 1},
+    )
+
+
+def get_video_query(
+    question: str = None,
+    video_path: str | None = None,
+    num_frames: int = 16,
+) -> QueryResult:
+    if question is None:
+        question = "Why is this video funny?"
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+    if video_path:
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
+    else:
+        video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {"video": video_frames},
+        },
+        limit_mm_per_prompt={"video": 1},
+    )
+
+
+query_map = {
+    "text": get_text_query,
+    "use_audio": get_audio_query,
+    "use_image": get_image_query,
+    "use_video": get_video_query,
+}
+
+# ---------------------------------------------------------------------------
+# Core async routine
+# ---------------------------------------------------------------------------
+
+
+def clone_prompt_for_request(template: dict) -> dict:
+    """Shallow-clone prompt dict so concurrent requests own independent containers."""
+    cloned = dict(template)
+    for key in ("multi_modal_data", "mm_processor_kwargs", "additional_information"):
+        value = template.get(key)
+        if isinstance(value, dict):
+            cloned[key] = dict(value)
+        elif isinstance(value, list):
+            cloned[key] = list(value)
+    return cloned
+
+
+def _default_async_chunk_stage_configs_path() -> str | None:
+    """Best-effort default stage config for running Qwen3-Omni with async_chunk.
+
+    When this example is executed from within the repository, we resolve the
+    default YAML path relative to this file. When installed elsewhere, the
+    file may not exist and callers should pass --stage-configs-path explicitly.
+    """
+    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
+    candidate = os.path.join(
+        repo_root,
+        "vllm_omni",
+        "model_executor",
+        "stage_configs",
+        "qwen3_omni_moe_async_chunk.yaml",
+    )
+    return candidate if os.path.exists(candidate) else None
+
+
+async def run_single_request(
+    async_omni: AsyncOmni,
+    prompt: dict,
+    request_id: str,
+    sampling_params_list: list[SamplingParams] | None,
+    output_dir: str,
+    output_modalities: list[str] | None = None,
+    stream_audio_to_disk: bool = False,
+) -> dict:
+    """Run one request through AsyncOmni and collect outputs.
+
+    Returns a dict with timing information and saved file paths.
+    """
+    t_start = time.perf_counter()
+    text_parts: list[str] = []
+    audio_chunks: list[torch.Tensor] = []
+    audio_sr: int | None = None
+    first_audio_ts: float | None = None
+    audio_list_consumed: int = 0
+    audio_last_tensor: torch.Tensor | None = None
+    stage_0_first_output_ts: float | None = None
+
+    samplerate = 24000
+    wav_file = os.path.join(output_dir, f"output_{request_id}.wav")
+    sf_writer: sf.SoundFile | None = None
+    audio_samples_written: int = 0
+
+    try:
+        async for omni_output in async_omni.generate(
+            prompt=prompt,
+            request_id=request_id,
+            sampling_params_list=sampling_params_list,
+            output_modalities=output_modalities,
+        ):
+            if not isinstance(omni_output.request_output, list):
+                outputs_list = [omni_output.request_output]
+            else:
+                outputs_list = omni_output.request_output
+
+            for output in outputs_list:
+                if omni_output.final_output_type == "text":
+                    if stage_0_first_output_ts is None:
+                        stage_0_first_output_ts = time.perf_counter()
+                    text_output = output.outputs[0].text
+                    if output.finished:
+                        text_parts.append(text_output)
+                elif omni_output.final_output_type == "audio":
+                    mm_out = output.outputs[0].multimodal_output
+                    if mm_out and "audio" in mm_out:
+                        if first_audio_ts is None:
+                            first_audio_ts = time.perf_counter()
+                        if audio_sr is None and "sr" in mm_out:
+                            sr_val = mm_out["sr"]
+                            audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
+                            samplerate = audio_sr
+                        audio_data = mm_out["audio"]
+                        if isinstance(audio_data, list):
+                            new_chunks = audio_data[audio_list_consumed:]
+                            audio_list_consumed = len(audio_data)
+                        elif isinstance(audio_data, torch.Tensor):
+                            new_chunks = [audio_data]
+                            audio_last_tensor = audio_data
+                        else:
+                            new_chunks = []
+
+                        if stream_audio_to_disk and new_chunks:
+                            if sf_writer is None:
+                                sf_writer = sf.SoundFile(
+                                    wav_file,
+                                    mode="w",
+                                    samplerate=samplerate,
+                                    channels=1,
+                                    subtype="FLOAT",
+                                )
+                            for chunk in new_chunks:
+                                chunk_np = chunk.float().detach().cpu().numpy().flatten()
+                                sf_writer.write(chunk_np)
+                                audio_samples_written += len(chunk_np)
+                        else:
+                            audio_chunks.extend(new_chunks)
+    finally:
+        if sf_writer is not None:
+            sf_writer.close()
+
+    t_end = time.perf_counter()
+    result = {
+        "request_id": request_id,
+        "e2e_latency_s": t_end - t_start,
+        "saved_files": [],
+    }
+
+    # Save text output
+    if text_parts:
+        text_file = os.path.join(output_dir, f"{request_id}.txt")
+        with open(text_file, "w", encoding="utf-8") as f:
+            f.write("\n".join(text_parts))
+        result["saved_files"].append(text_file)
+        print(
+            f"[Request {request_id}] Text saved to {text_file} "
+            f"(stage-0 first output at {stage_0_first_output_ts - t_start:.3f}s)"
+        )
+
+    # Save audio output
+    if stream_audio_to_disk and audio_samples_written > 0:
+        result["saved_files"].append(wav_file)
+        result["audio_duration_s"] = audio_samples_written / samplerate
+        result["num_audio_chunks"] = audio_list_consumed
+        ttfa = (first_audio_ts - t_start) if first_audio_ts else None
+        result["time_to_first_audio_s"] = ttfa
+        ttfa_str = f"{ttfa:.3f}s" if ttfa is not None else "N/A"
+        print(
+            f"[Request {request_id}] Audio streamed to {wav_file} "
+            f"(duration={result['audio_duration_s']:.2f}s, "
+            f"TTFA={ttfa_str}, "
+            f"e2e={result['e2e_latency_s']:.3f}s)"
+        )
+    elif audio_chunks or audio_last_tensor is not None:
+        if audio_chunks:
+            if len(audio_chunks) > 1:
+                audio_tensor = torch.cat(audio_chunks, dim=-1)
+            else:
+                audio_tensor = audio_chunks[0]
+        else:
+            audio_tensor = audio_last_tensor
+        audio_numpy = audio_tensor.float().detach().cpu().numpy()
+        if audio_numpy.ndim > 1:
+            audio_numpy = audio_numpy.flatten()
+        sf.write(wav_file, audio_numpy, samplerate=samplerate, format="WAV")
+        result["saved_files"].append(wav_file)
+        result["audio_duration_s"] = len(audio_numpy) / samplerate
+        result["num_audio_chunks"] = len(audio_chunks)
+        ttfa = (first_audio_ts - t_start) if first_audio_ts else None
+        result["time_to_first_audio_s"] = ttfa
+        ttfa_str = f"{ttfa:.3f}s" if ttfa is not None else "N/A"
+        print(
+            f"[Request {request_id}] Audio saved to {wav_file} "
+            f"({len(audio_chunks)} chunks, "
+            f"duration={result['audio_duration_s']:.2f}s, "
+            f"TTFA={ttfa_str}, "
+            f"e2e={result['e2e_latency_s']:.3f}s)"
+        )
+
+    return result
+
+
+async def run_all(args):
+    """Main async entry: build prompts, create AsyncOmni, run requests."""
+    # Build query
+    query_func = query_map[args.query_type]
+    if args.query_type == "use_video":
+        query_result = query_func(
+            video_path=getattr(args, "video_path", None),
+            num_frames=getattr(args, "num_frames", 16),
+        )
+    elif args.query_type == "use_image":
+        query_result = query_func(image_path=getattr(args, "image_path", None))
+    elif args.query_type == "use_audio":
+        query_result = query_func(
+            audio_path=getattr(args, "audio_path", None),
+            sampling_rate=getattr(args, "sampling_rate", 16000),
+        )
+    else:
+        query_result = query_func()
+
+    # Build prompt list
+    if args.txt_prompts is not None:
+        assert args.query_type == "text", "txt-prompts is only supported for text query type"
+        with open(args.txt_prompts, encoding="utf-8") as f:
+            lines = [ln.strip() for ln in f if ln.strip()]
+        prompts = [get_text_query(ln).inputs for ln in lines]
+        print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
+    else:
+        prompts = [clone_prompt_for_request(query_result.inputs) for _ in range(args.num_prompts)]
+
+    # Inject output modalities if specified
+    output_modalities = None
+    if args.modalities is not None:
+        output_modalities = args.modalities.split(",")
+        for prompt in prompts:
+            prompt["modalities"] = output_modalities
+
+    # Create AsyncOmni
+    print(f"[Info] Creating AsyncOmni with stage_configs_path={args.stage_configs_path}")
+    async_omni = None
+    try:
+        async_omni = AsyncOmni(
+            model=args.model,
+            stage_configs_path=args.stage_configs_path,
+            log_stats=args.log_stats,
+            stage_init_timeout=args.stage_init_timeout,
+        )
+
+        # Use default sampling params from stage config (they are pre-configured
+        # in the YAML for each stage).
+        sampling_params_list = None
+
+        output_dir = args.output_dir
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Run requests with concurrency control
+        semaphore = asyncio.Semaphore(args.max_in_flight)
+        request_timeout = getattr(args, "request_timeout_s", None)
+        stream_audio = getattr(args, "stream_audio_to_disk", False)
+
+        async def _run_one(idx: int, prompt: dict):
+            async with semaphore:
+                request_id = f"req_{idx}_{uuid.uuid4().hex[:8]}"
+                coro = run_single_request(
+                    async_omni=async_omni,
+                    prompt=prompt,
+                    request_id=request_id,
+                    sampling_params_list=sampling_params_list,
+                    output_dir=output_dir,
+                    output_modalities=output_modalities,
+                    stream_audio_to_disk=stream_audio,
+                )
+                if request_timeout and request_timeout > 0:
+                    return await asyncio.wait_for(coro, timeout=request_timeout)
+                return await coro
+
+        wall_start = time.perf_counter()
+        tasks = [_run_one(i, p) for i, p in enumerate(prompts)]
+        all_results = await asyncio.gather(*tasks, return_exceptions=True)
+        wall_end = time.perf_counter()
+
+        # Print summary
+        print("\n" + "=" * 60)
+        print("Summary")
+        print("=" * 60)
+        success_count = 0
+        total_audio_dur = 0.0
+        for r in all_results:
+            if isinstance(r, Exception):
+                print(f"  [ERROR] {type(r).__name__}: {r}")
+            else:
+                success_count += 1
+                total_audio_dur += r.get("audio_duration_s", 0.0)
+                print(f"  [{r['request_id']}] e2e={r['e2e_latency_s']:.3f}s  files={r['saved_files']}")
+        wall_time = wall_end - wall_start
+        print(f"\nTotal: {success_count}/{len(prompts)} succeeded")
+        print(f"Wall time: {wall_time:.3f}s")
+        if total_audio_dur > 0:
+            print(f"Total audio duration: {total_audio_dur:.2f}s")
+            print(f"Real-time factor: {total_audio_dur / wall_time:.2f}x")
+        print("=" * 60)
+    finally:
+        if async_omni is not None:
+            async_omni.shutdown()
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(
+        description=(
+            "Offline inference with async_chunk enabled via AsyncOmni. "
+            "Downstream stages start before upstream stages finish, "
+            "achieving true stage-level concurrency."
+        )
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        help="Model name or path.",
+    )
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="use_audio",
+        choices=query_map.keys(),
+        help="Query type.",
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=_default_async_chunk_stage_configs_path(),
+        help=(
+            "Path to an async_chunk stage config YAML. "
+            "If not set, uses the model's default config "
+            "(make sure it has async_chunk: true)."
+        ),
+    )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics.",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing a single stage (seconds).",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="output_audio_async_chunk",
+        help="Directory to save output files.",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of prompts to generate (duplicated from query).",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line.",
+    )
+    parser.add_argument(
+        "--max-in-flight",
+        type=int,
+        default=1,
+        help="Maximum concurrent requests (default: 1).",
+    )
+    parser.add_argument(
+        "--request-timeout-s",
+        type=float,
+        default=None,
+        help=(
+            "Per-request timeout in seconds. When set, a request that "
+            "exceeds this duration is cancelled and reported as an error. "
+            "Default: None (no timeout)."
+        ),
+    )
+    parser.add_argument(
+        "--batch-timeout-s",
+        type=float,
+        default=None,
+        help=(
+            "Global timeout for the entire batch in seconds. When set, "
+            "the whole run_all() is cancelled if it exceeds this duration. "
+            "Default: None (no global timeout)."
+        ),
+    )
+    parser.add_argument(
+        "--stream-audio-to-disk",
+        action="store_true",
+        default=False,
+        help=(
+            "Write audio chunks to WAV incrementally instead of "
+            "accumulating in memory. Useful for very long audio or "
+            "high --max-in-flight to reduce memory footprint."
+        ),
+    )
+    parser.add_argument(
+        "--modalities",
+        type=str,
+        default=None,
+        help="Comma-separated output modalities filter (e.g. 'text', 'audio', 'text,audio').",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file.",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Path to local image file.",
+    )
+    parser.add_argument(
+        "--video-path",
+        "-v",
+        type=str,
+        default=None,
+        help="Path to local video file.",
+    )
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=16,
+        help="Number of frames to extract from video.",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Sampling rate for audio loading.",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    async def _main():
+        batch_timeout = getattr(args, "batch_timeout_s", None)
+        if batch_timeout and batch_timeout > 0:
+            await asyncio.wait_for(run_all(args), timeout=batch_timeout)
+        else:
+            await run_all(args)
+
+    try:
+        asyncio.run(_main())
+    except asyncio.TimeoutError:
+        print(
+            f"\n[TIMEOUT] Batch exceeded --batch-timeout-s="
+            f"{args.batch_timeout_s}s. AsyncOmni shutdown was handled "
+            f"by finally block."
+        )
+    except KeyboardInterrupt:
+        print("\nInterrupted by user. AsyncOmni shutdown was handled by finally block.")

examples/offline_inference/qwen3_omni/run_multiple_prompts_async_chunk.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+# Run multiple Qwen3-Omni requests with async_chunk enabled.
+#
+# Uses AsyncOmni with --max-in-flight to control request-level
+# concurrency (each request still gets true stage-level concurrency
+# via async_chunk).
+#
+# Usage:
+#   bash run_multiple_prompts_async_chunk.sh
+#   bash run_multiple_prompts_async_chunk.sh --max-in-flight 4
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+
+python "${SCRIPT_DIR}/end2end_async_chunk.py" \
+    --query-type text \
+    --txt-prompts "${SCRIPT_DIR}/text_prompts_10.txt" \
+    --stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \
+    --output-dir output_audio_async_chunk \
+    --max-in-flight 2 \
+    "$@"

examples/offline_inference/qwen3_omni/run_single_prompt_async_chunk.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+# Run a single Qwen3-Omni request with async_chunk enabled.
+#
+# This uses AsyncOmni (async orchestrator) so that downstream stages
+# (Talker, Code2Wav) start *before* stage-0 (Thinker) finishes,
+# achieving true stage-level concurrency via chunk-level streaming.
+#
+# Prerequisites:
+#   - An async_chunk stage config YAML (e.g. qwen3_omni_moe_async_chunk.yaml)
+#   - Hardware matching the config (e.g. 2x H100 for the default 3-stage config)
+#
+# Usage:
+#   bash run_single_prompt_async_chunk.sh
+#   bash run_single_prompt_async_chunk.sh --query-type text --modalities text
+#   bash run_single_prompt_async_chunk.sh --stage-configs-path /path/to/custom.yaml
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+
+python "${SCRIPT_DIR}/end2end_async_chunk.py" \
+    --query-type use_audio \
+    --stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \
+    --output-dir output_audio_async_chunk \
+    "$@"

examples/offline_inference/qwen3_tts/README.md

@@ -95,7 +95,8 @@ Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires
 python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
 ```
 
-Each 25-frame Code2Wav chunk is logged as it arrives. The final WAV file is written once generation
+Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
+and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation
 completes. This demonstrates that audio data is available progressively rather than only at the end.
 
 > **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.

examples/offline_inference/qwen3_tts/end2end.py

@@ -7,6 +7,7 @@ tasks, then runs Omni generation and saves output wav files.
 import asyncio
 import logging
 import os
+import time
 from typing import Any, NamedTuple
 
 import soundfile as sf
@@ -78,10 +79,16 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
     task_type = "CustomVoice"
     model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
     if use_batch_sample:
-        texts = ["其实我真的有发现，我是一个特别善于观察别人情绪的人。", "She said she would be here by noon."]
-        instructs = ["", "Very happy."]
-        languages = ["Chinese", "English"]
-        speakers = ["Vivian", "Ryan"]
+        texts = [
+            "其实我真的有发现，我是一个特别善于观察别人情绪的人。",
+            "She said she would be here by noon.",
+            "I like you very much.",
+            "Really, you do?",
+            "Yes, absolutely.",
+        ]
+        instructs = ["", "Very happy.", "Very happy.", "Very happy.", "Very happy."]
+        languages = ["Chinese", "English", "English", "English", "English"]
+        speakers = ["Vivian", "Ryan", "Ryan", "Ryan", "Ryan"]
         inputs = []
         for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
             additional_information = {
@@ -337,13 +344,27 @@ async def main_streaming(args):
 
     for i, prompt in enumerate(inputs):
         request_id = str(i)
+        t_start = time.perf_counter()
+        t_prev = t_start
+        chunk_idx = 0
         async for stage_output in omni.generate(prompt, request_id=request_id):
             mm = stage_output.request_output.outputs[0].multimodal_output
             if not stage_output.finished:
+                t_now = time.perf_counter()
                 audio = mm.get("audio")
                 n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
-                logger.info(f"Request {request_id}: received chunk {n}")
+                dt_ms = (t_now - t_prev) * 1000
+                ttfa_ms = (t_now - t_start) * 1000
+                if chunk_idx == 0:
+                    logger.info(f"Request {request_id}: chunk {chunk_idx} samples={n} TTFA={ttfa_ms:.1f}ms")
+                else:
+                    logger.info(f"Request {request_id}: chunk {chunk_idx} samples={n} inter_chunk={dt_ms:.1f}ms")
+                t_prev = t_now
+                chunk_idx += 1
             else:
+                t_end = time.perf_counter()
+                total_ms = (t_end - t_start) * 1000
+                logger.info(f"Request {request_id}: done total={total_ms:.1f}ms chunks={chunk_idx}")
                 _save_wav(output_dir, request_id, mm)
 
 

examples/offline_inference/text_to_video/text_to_video.md

@@ -1,6 +1,7 @@
 # Text-To-Video
 
-The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts.
+The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts. This script can also be used
+for `Lightricks/LTX-2` to generate video+audio.
 
 ## Local CLI Usage
 
@@ -19,6 +20,22 @@ python text_to_video.py \
   --output t2v_out.mp4
 ```
 
+LTX2 example:
+
+```bash
+python text_to_video.py \
+  --model "Lightricks/LTX-2" \
+  --prompt "A cinematic close-up of ocean waves at golden hour." \
+  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
+  --height 512 \
+  --width 768 \
+  --num-frames 121 \
+  --num-inference-steps 40 \
+  --guidance-scale 4.0 \
+  --frame-rate 24 \
+  --output ltx2_out.mp4
+```
+
 Key arguments:
 
 - `--prompt`: text description (string).
@@ -32,6 +49,9 @@ Key arguments:
 - `--vae-use-slicing`: enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
+- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/text_to_video/text_to_video.py

@@ -31,6 +31,12 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--height", type=int, default=720, help="Video height.")
     parser.add_argument("--width", type=int, default=1280, help="Video width.")
     parser.add_argument("--num-frames", type=int, default=81, help="Number of frames (Wan default is 81).")
+    parser.add_argument(
+        "--frame-rate",
+        type=float,
+        default=None,
+        help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
+    )
     parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.")
     parser.add_argument(
         "--boundary-ratio",
@@ -109,12 +115,25 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
+    parser.add_argument(
+        "--audio-sample-rate",
+        type=int,
+        default=24000,
+        help="Sample rate for audio output when saved (default: 24000 for LTX2).",
+    )
+    parser.add_argument(
+        "--vae-patch-parallel-size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for VAE patch/tile parallelism (decode).",
+    )
     return parser.parse_args()
 
 
 def main():
     args = parse_args()
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    frame_rate = args.frame_rate if args.frame_rate is not None else float(args.fps)
 
     # Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
     cache_config = None
@@ -134,13 +153,13 @@ def main():
             "scm_steps_mask_policy": None,  # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
             "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
         }
-    # Configure parallel settings (only SP is supported for Wan)
-    # Note: cfg_parallel and tensor_parallel are not implemented for Wan models
+    # Configure parallel settings
     parallel_config = DiffusionParallelConfig(
         ulysses_degree=args.ulysses_degree,
         ring_degree=args.ring_degree,
         cfg_parallel_size=args.cfg_parallel_size,
         tensor_parallel_size=args.tensor_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
     )
 
     # Check if profiling is requested via environment variable
@@ -153,12 +172,12 @@ def main():
         vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
-        cache_backend=args.cache_backend,
-        cache_config=cache_config,
         enable_cache_dit_summary=args.enable_cache_dit_summary,
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        cache_backend=args.cache_backend,
+        cache_config=cache_config,
     )
 
     if profiler_enabled:
@@ -172,7 +191,9 @@ def main():
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Frames: {args.num_frames}")
     print(
-        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}"
+        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
+        f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
+        f" vae_patch_parallel_size={args.vae_patch_parallel_size}"
     )
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
@@ -191,6 +212,7 @@ def main():
             guidance_scale_2=args.guidance_scale_high,
             num_inference_steps=args.num_inference_steps,
             num_frames=args.num_frames,
+            frame_rate=frame_rate,
         ),
     )
     generation_end = time.perf_counter()
@@ -199,31 +221,56 @@ def main():
     # Print profiling results
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
-    # Extract video frames from OmniRequestOutput
-    if isinstance(frames, list) and len(frames) > 0:
-        first_item = frames[0]
+    audio = None
+    if isinstance(frames, list):
+        frames = frames[0] if frames else None
 
-        # Check if it's an OmniRequestOutput
-        if hasattr(first_item, "final_output_type"):
-            if first_item.final_output_type != "image":
-                raise ValueError(
-                    f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
-                )
-
-            # Pipeline mode: extract from nested request_output
-            if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
-                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
-                    inner_output = first_item.request_output[0]
-                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
-                        frames = inner_output.images[0] if inner_output.images else None
-                        if frames is None:
-                            raise ValueError("No video frames found in output.")
-            # Diffusion mode: use direct images field
-            elif hasattr(first_item, "images") and first_item.images:
-                frames = first_item.images
+    if isinstance(frames, OmniRequestOutput):
+        if frames.final_output_type != "image":
+            raise ValueError(
+                f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
+            )
+        if frames.multimodal_output and "audio" in frames.multimodal_output:
+            audio = frames.multimodal_output["audio"]
+        if frames.is_pipeline_output and frames.request_output is not None:
+            inner_output = frames.request_output
+            if isinstance(inner_output, list):
+                inner_output = inner_output[0] if inner_output else None
+            if isinstance(inner_output, OmniRequestOutput):
+                if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
+                    audio = inner_output.multimodal_output["audio"]
+                frames = inner_output
+        if isinstance(frames, OmniRequestOutput):
+            if frames.images:
+                if len(frames.images) == 1 and isinstance(frames.images[0], tuple) and len(frames.images[0]) == 2:
+                    frames, audio = frames.images[0]
+                elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
+                    audio = frames.images[0].get("audio")
+                    frames = frames.images[0].get("frames") or frames.images[0].get("video")
+                else:
+                    frames = frames.images
             else:
                 raise ValueError("No video frames found in OmniRequestOutput.")
 
+    if isinstance(frames, list) and frames:
+        first_item = frames[0]
+        if isinstance(first_item, tuple) and len(first_item) == 2:
+            frames, audio = first_item
+        elif isinstance(first_item, dict):
+            audio = first_item.get("audio")
+            frames = first_item.get("frames") or first_item.get("video")
+        elif isinstance(first_item, list):
+            frames = first_item
+
+    if isinstance(frames, tuple) and len(frames) == 2:
+        frames, audio = frames
+    elif isinstance(frames, dict):
+        audio = frames.get("audio")
+        frames = frames.get("frames") or frames.get("video")
+
+    if frames is None:
+        raise ValueError("No video frames found in output.")
+
     output_path = Path(args.output)
     output_path.parent.mkdir(parents=True, exist_ok=True)
     try:
@@ -231,32 +278,127 @@ def main():
     except ImportError:
         raise ImportError("diffusers is required for export_to_video.")
 
-    # frames may be np.ndarray (preferred) or torch.Tensor
+    def _normalize_frame(frame):
+        if isinstance(frame, torch.Tensor):
+            frame_tensor = frame.detach().cpu()
+            if frame_tensor.dim() == 4 and frame_tensor.shape[0] == 1:
+                frame_tensor = frame_tensor[0]
+            if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
+                frame_tensor = frame_tensor.permute(1, 2, 0)
+            if frame_tensor.is_floating_point():
+                frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
+            return frame_tensor.float().numpy()
+        if isinstance(frame, np.ndarray):
+            frame_array = frame
+            if frame_array.ndim == 4 and frame_array.shape[0] == 1:
+                frame_array = frame_array[0]
+            if np.issubdtype(frame_array.dtype, np.integer):
+                frame_array = frame_array.astype(np.float32) / 255.0
+            return frame_array
+        try:
+            from PIL import Image
+        except ImportError:
+            Image = None
+        if Image is not None and isinstance(frame, Image.Image):
+            return np.asarray(frame).astype(np.float32) / 255.0
+        return frame
+
+    def _ensure_frame_list(video_array):
+        if isinstance(video_array, list):
+            if len(video_array) == 0:
+                return video_array
+            first_item = video_array[0]
+            if isinstance(first_item, np.ndarray):
+                if first_item.ndim == 5:
+                    return list(first_item[0])
+                if first_item.ndim == 4:
+                    if len(video_array) == 1:
+                        return list(first_item)
+                    return list(first_item)
+                if first_item.ndim == 3:
+                    return video_array
+            return video_array
+        if isinstance(video_array, np.ndarray):
+            if video_array.ndim == 5:
+                return list(video_array[0])
+            if video_array.ndim == 4:
+                return list(video_array)
+            if video_array.ndim == 3:
+                return [video_array]
+        return video_array
+
+    # frames may be np.ndarray, torch.Tensor, or list of tensors/arrays/images
     # export_to_video expects a list of frames with values in [0, 1]
     if isinstance(frames, torch.Tensor):
         video_tensor = frames.detach().cpu()
         if video_tensor.dim() == 5:
-            # [B, C, F, H, W] or [B, F, H, W, C]
             if video_tensor.shape[1] in (3, 4):
                 video_tensor = video_tensor[0].permute(1, 2, 3, 0)
             else:
                 video_tensor = video_tensor[0]
         elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
             video_tensor = video_tensor.permute(1, 2, 3, 0)
-        # If float, assume [-1,1] and normalize to [0,1]
         if video_tensor.is_floating_point():
             video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
         video_array = video_tensor.float().numpy()
+    elif isinstance(frames, np.ndarray):
+        video_array = frames
+        if video_array.ndim == 5:
+            video_array = video_array[0]
+        if np.issubdtype(video_array.dtype, np.integer):
+            video_array = video_array.astype(np.float32) / 255.0
+    elif isinstance(frames, list):
+        if len(frames) == 0:
+            raise ValueError("No video frames found in output.")
+        video_array = [_normalize_frame(frame) for frame in frames]
     else:
         video_array = frames
-        if hasattr(video_array, "shape") and video_array.ndim == 5:
-            video_array = video_array[0]
 
-    # Convert 4D array (frames, H, W, C) to list of frames for export_to_video
-    if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
-        video_array = list(video_array)
+    video_array = _ensure_frame_list(video_array)
 
-    export_to_video(video_array, str(output_path), fps=args.fps)
+    use_ltx2_export = False
+    if args.model and "ltx" in str(args.model).lower():
+        use_ltx2_export = True
+    if audio is not None:
+        use_ltx2_export = True
+
+    if use_ltx2_export:
+        try:
+            from diffusers.pipelines.ltx2.export_utils import encode_video
+        except ImportError:
+            raise ImportError("diffusers is required for LTX2 encode_video.")
+
+        if isinstance(video_array, list):
+            frames_np = np.stack(video_array, axis=0)
+        elif isinstance(video_array, np.ndarray):
+            frames_np = video_array
+        else:
+            frames_np = np.asarray(video_array)
+
+        frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
+        video_tensor = torch.from_numpy(frames_u8)
+
+        audio_out = None
+        if audio is not None:
+            if isinstance(audio, list):
+                audio = audio[0] if audio else None
+            if isinstance(audio, np.ndarray):
+                audio = torch.from_numpy(audio)
+            if isinstance(audio, torch.Tensor):
+                audio_out = audio
+                if audio_out.dim() > 1:
+                    audio_out = audio_out[0]
+                audio_out = audio_out.float().cpu()
+
+        encode_video(
+            video_tensor,
+            fps=args.fps,
+            audio=audio_out,
+            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
+            output_path=str(output_path),
+        )
+    else:
+        export_to_video(video_array, str(output_path), fps=args.fps)
     print(f"Saved generated video to {output_path}")
 
     if profiler_enabled:

examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh

@@ -160,10 +160,7 @@ esac
 echo "Running query type: $QUERY_TYPE"
 echo ""
 
-
-output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d @- <<EOF
+request_body=$(cat <<EOF
 {
   "model": "Qwen/Qwen2.5-Omni-7B",
   "sampling_params_list": $sampling_params_list,
@@ -186,7 +183,12 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
   ]
 }
 EOF
-  )
+)
+
+output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
+    -X POST http://localhost:8091/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d "$request_body")
 
 # Here it only shows the text content of the first choice. Audio content has many binaries, so it's not displayed here.
 echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"

examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh

@@ -138,10 +138,7 @@ esac
 echo "Running query type: $QUERY_TYPE"
 echo ""
 
-
-output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
-    -H "Content-Type: application/json" \
-    -d @- <<EOF
+request_body=$(cat <<EOF
 {
   "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
   "sampling_params_list": $sampling_params_list,
@@ -164,7 +161,12 @@ output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \
   ]
 }
 EOF
-  )
+)
+
+output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
+    -X POST http://localhost:8091/v1/chat/completions \
+    -H "Content-Type: application/json" \
+    -d "$request_body")
 
 # Here it only shows the text content of the first choice. Audio content has many binaries, so it's not displayed here.
 echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')"

examples/online_serving/qwen3_tts/README.md

@@ -184,6 +184,59 @@ sudo apt install ffmpeg
 
 ## API Reference
 
+### Voices Endpoint
+
+#### GET /v1/audio/voices
+
+List all available voices/speakers from the loaded model, including both built-in model voices and uploaded custom voices.
+
+**Response Example:**
+```json
+{
+  "voices": ["vivian", "ryan", "custom_voice_1"],
+  "uploaded_voices": [
+    {
+      "name": "custom_voice_1",
+      "consent": "user_consent_id",
+      "created_at": 1738660000,
+      "file_size": 1024000,
+      "mime_type": "audio/wav"
+    }
+  ]
+}
+```
+
+#### POST /v1/audio/voices
+
+Upload a new voice sample for voice cloning in Base task TTS requests.
+
+**Form Parameters:**
+- `audio_sample` (required): Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4)
+- `consent` (required): Consent recording ID
+- `name` (required): Name for the new voice
+
+**Response Example:**
+```json
+{
+  "success": true,
+  "voice": {
+    "name": "custom_voice_1",
+    "consent": "user_consent_id",
+    "created_at": 1738660000,
+    "mime_type": "audio/wav",
+    "file_size": 1024000
+  }
+}
+```
+
+**Usage Example:**
+```bash
+curl -X POST http://localhost:8000/v1/audio/voices \
+  -F "audio_sample=@/path/to/voice_sample.wav" \
+  -F "consent=user_consent_id" \
+  -F "name=custom_voice_1"
+```
+
 ### Endpoint
 
 ```
@@ -193,20 +246,6 @@ Content-Type: application/json
 
 This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
 
-### Voices Endpoint
-
-```
-GET /v1/audio/voices
-```
-
-Lists available voices for the loaded model:
-
-```json
-{
-    "voices": ["aiden", "dylan", "eric", "one_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
-}
-```
-
 ### Request Body
 
 ```json
@@ -250,6 +289,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 | `language`       | string | "Auto"        | Language (see supported languages below)                                     |
 | `instructions`   | string | ""            | Voice style/emotion instructions                                             |
 | `max_new_tokens` | int    | 2048          | Maximum tokens to generate                                                   |
+| `initial_codec_chunk_frames` | int | null | Initial chunk size for reduced TTFA (overrides stage config)                 |
 | `stream`         | bool   | false         | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
 
 **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
@@ -265,7 +305,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 ## Streaming
 
 Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded
-(one chunk per 25-frame Code2Wav window):
+(one chunk per Code2Wav window, default 25 frames; configurable in the stage config):
 
 ```bash
 curl -X POST http://localhost:8091/v1/audio/speech \

pyproject.toml

@@ -49,7 +49,8 @@ dev = [
     "imageio[ffmpeg]>=0.6.0",
     "opencv-python>=4.12.0.88",
     "mooncake-transfer-engine==0.3.8.post1",
-    "av" # for ComfyUI tests
+    "av", # for ComfyUI tests
+    "openpyxl>=3.0.0", # for nightly CI
 ]
 
 docs = [
@@ -202,3 +203,5 @@ extend-ignore-identifiers-re = [
     ".*[Oo]no_[Aa]nna.*",
     ".*cann.*",
 ]
+[tool.typos.default.extend-words]
+ue = "ue"

requirements/common.txt

@@ -12,6 +12,8 @@ torchsde>=0.2.6
 openai-whisper>=20250625
 imageio[ffmpeg]>=2.37.2
 sox>=1.5.0
+x-transformers>=2.12.2
+einops>=0.8.1
 prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0

tests/diffusion/attention/test_attention_sp.py

@@ -414,7 +414,8 @@ def ulysses_attention_on_test_model(
             # Generate and save full input data with fixed seed
             # Reinitialize RNG to ensure reproducibility
             torch.manual_seed(42)
-            current_omni_platform.seed_everything(42)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(42)
             full_hidden_states = torch.randn(
                 (batch_size, seq_len, hidden_size),
                 dtype=dtype,

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -0,0 +1,165 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedVaeExecutor,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+
+class E2EOperator:
+    """tiles with (2, 3) -- (H,W)"""
+
+    def split(self, z):
+        rows_num = 2
+        cols_num = 3
+
+        h_size = z.shape[0]
+        w_size = z.shape[1]
+
+        tasks = []
+        for i in range(rows_num):
+            for j in range(cols_num):
+                tasks.append(
+                    TileTask(
+                        tile_id=len(tasks),
+                        grid_coord=(i, j),
+                        tensor=z[
+                            ((i * h_size) // rows_num) : (((i + 1) * h_size) // rows_num),
+                            ((j * w_size) // cols_num) : (((j + 1) * w_size) // cols_num),
+                        ],
+                    )
+                )
+
+        grid_spec = GridSpec(
+            split_dims=(0, 1),
+            grid_shape=(2, 3),
+        )
+        return tasks, grid_spec
+
+    def exec(self, task: TileTask):
+        return torch.full_like(task.tensor, fill_value=task.tile_id)
+
+    def merge(self, coord_tensor_map, grid_spec):
+        tiles = []
+        for r in range(grid_spec.grid_shape[0]):
+            row_tiles = []
+            for c in range(grid_spec.grid_shape[1]):
+                coord = (r, c)
+                row_tiles.append(coord_tensor_map[coord])
+            tiles.append(torch.cat(row_tiles, dim=1))
+        return torch.cat(tiles, dim=0)
+
+
+class DummyMixin(DistributedVaeMixin):
+    def __init__(self):
+        self.use_tiling = True
+        self.distributed_decoder = MagicMock()
+        self.distributed_decoder.parallel_size = 2
+        self.distributed_decoder.group = None
+
+
+@pytest.fixture(autouse=True)
+def mock_dist():
+    with (
+        patch.object(dist, "get_world_size", return_value=2),
+        patch.object(dist, "get_rank", return_value=0),
+        patch.object(dist, "is_initialized", return_value=True),
+        patch.object(dist, "gather", return_value=None),
+        patch.object(dist, "broadcast", return_value=None),
+    ):
+        yield
+
+
+@pytest.fixture(autouse=True)
+def mock_dit_group():
+    with patch(
+        "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group",
+        new=MagicMock(return_value=None),
+    ):
+        yield
+
+
+@pytest.fixture(autouse=True)
+def mock_dist_vae_executor():
+    with (
+        patch.object(DistributedVaeExecutor, "gather_tensors", side_effect=lambda x: [x]),
+        patch.object(DistributedVaeExecutor, "broadcast_tensor", side_effect=lambda x: x),
+    ):
+        yield
+
+
+# ============================
+# Unitest
+# ============================
+
+
+def test_balance_tasks():
+    executor = DistributedVaeExecutor()
+    workloads = [2, 5, 13, 8, 2, 4]  # each is 17
+    tasks = [TileTask(0, (i,), torch.tensor([i]), workload=load) for i, load in enumerate(workloads)]
+    assigned = executor._balance_tasks(tasks, num_rank=2)
+    assert len(assigned) == 2
+    total_work = [sum(t.workload for t in group) for group in assigned]
+    assert total_work[0] == total_work[1]
+
+
+def test_compute_global_padding_shape():
+    executor = DistributedVaeExecutor()
+    executor.rank = 0
+
+    local_results = [(0, torch.zeros((2, 3))), (1, torch.zeros((4, 2)))]
+    shape = executor._compute_global_padding_shape(local_results, 2, "cpu")
+
+    assert shape == [2, 4, 3]
+
+
+def test_pack_and_unpack():
+    executor = DistributedVaeExecutor()
+    executor.world_size = 1
+
+    grid_spec = GridSpec(split_dims=(0, 1), grid_shape=(2, 2))
+
+    # ======================
+    # pack
+    # ======================
+    local_results = [(0, torch.tensor([[1, 2], [3, 4]]))]
+
+    global_shape = [3, 3, 3]  # (tiles, H, W)
+
+    tile_tensor, meta_tensor = executor._pack_local_tiles(
+        local_results, global_shape, grid_spec, device="cpu", dtype=torch.int64
+    )
+
+    # check pack
+    assert tile_tensor.shape == torch.Size(global_shape)
+    assert meta_tensor.shape == (global_shape[0], len(grid_spec.split_dims) + 1)
+    assert meta_tensor[0, 0] == 0
+    assert meta_tensor[0, 1] == 2
+    assert meta_tensor[0, 2] == 2
+
+    # ======================
+    # unpack
+    # ======================
+    meta_gather = [meta_tensor]
+    tile_gather = [tile_tensor]
+
+    tid_coord_map = {0: (0, 0)}
+
+    coord_tensor_map = executor._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map)
+
+    # check unpack
+    assert torch.equal(coord_tensor_map[(0, 0)], torch.tensor([[1, 2], [3, 4]]))
+
+
+def test_is_distributed_enabled():
+    mixin = DummyMixin()
+    assert mixin.is_distributed_enabled() is True
+
+    mixin.use_tiling = False
+    assert mixin.is_distributed_enabled() is False

tests/diffusion/distributed/test_hsdp.py

@@ -156,13 +156,22 @@ class TestDiffusionParallelConfigHSDP:
 
     def test_hsdp_cannot_use_with_tp(self):
         """Test that HSDP and Tensor Parallelism cannot be used together."""
-        with pytest.raises(ValueError, match="cannot be used with Tensor Parallelism"):
+        with pytest.raises(ValueError, match="cannot be used with TP or DP"):
             DiffusionParallelConfig(
                 tensor_parallel_size=2,
                 use_hsdp=True,
                 hsdp_shard_size=4,
             )
 
+    def test_hsdp_cannot_use_with_dp(self):
+        """Test that HSDP and Data Parallelism cannot be used together."""
+        with pytest.raises(ValueError, match="cannot be used with TP or DP"):
+            DiffusionParallelConfig(
+                data_parallel_size=2,
+                use_hsdp=True,
+                hsdp_shard_size=4,
+            )
+
     def test_from_dict_with_hsdp(self):
         """Test creating config from dict with HSDP settings."""
         config = DiffusionParallelConfig.from_dict(
@@ -177,6 +186,158 @@ class TestDiffusionParallelConfigHSDP:
         assert config.hsdp_shard_size == 2  # auto: 4 // 2
 
 
+class TestStandaloneHSDPDetection:
+    """Tests for standalone HSDP detection and dit_parallel_size calculation.
+
+    These tests verify the logic used in initialize_model_parallel() to detect
+    standalone HSDP mode and calculate effective parallel sizes.
+
+    Standalone HSDP is when all non-HSDP parallelism dimensions are 1.
+    """
+
+    @staticmethod
+    def compute_standalone_hsdp_params(
+        data_parallel_size: int = 1,
+        cfg_parallel_size: int = 1,
+        sequence_parallel_size: int = 1,
+        pipeline_parallel_size: int = 1,
+        tensor_parallel_size: int = 1,
+        fully_shard_degree: int = 1,
+        hsdp_replicate_size: int = 1,
+    ) -> dict:
+        """Compute standalone HSDP detection parameters.
+
+        This mirrors the logic in initialize_model_parallel().
+        """
+        dit_parallel_size = (
+            data_parallel_size
+            * cfg_parallel_size
+            * sequence_parallel_size
+            * pipeline_parallel_size
+            * tensor_parallel_size
+        )
+
+        # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
+        is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
+
+        # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size)
+        if is_standalone_hsdp:
+            effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size
+        else:
+            effective_dit_parallel_size = dit_parallel_size
+
+        effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
+
+        return {
+            "original_dit_parallel_size": dit_parallel_size,
+            "is_standalone_hsdp": is_standalone_hsdp,
+            "effective_dit_parallel_size": effective_dit_parallel_size,
+            "effective_dp_size": effective_dp_size,
+        }
+
+    def test_standalone_hsdp_basic(self):
+        """Test basic standalone HSDP detection (shard_size=4, replicate=1)."""
+        result = self.compute_standalone_hsdp_params(
+            fully_shard_degree=4,
+            hsdp_replicate_size=1,
+        )
+        assert result["original_dit_parallel_size"] == 1
+        assert result["is_standalone_hsdp"] is True
+        assert result["effective_dit_parallel_size"] == 4
+        assert result["effective_dp_size"] == 4
+
+    def test_standalone_hsdp_with_replicate(self):
+        """Test standalone HSDP with replication (shard_size=4, replicate=2)."""
+        result = self.compute_standalone_hsdp_params(
+            fully_shard_degree=4,
+            hsdp_replicate_size=2,
+        )
+        assert result["original_dit_parallel_size"] == 1
+        assert result["is_standalone_hsdp"] is True
+        assert result["effective_dit_parallel_size"] == 8  # 4 * 2
+        assert result["effective_dp_size"] == 8
+
+    def test_combined_hsdp_sp_not_standalone(self):
+        """Test HSDP combined with SP is NOT detected as standalone.
+
+        This is a regression test for the bug where the condition
+        `dit_parallel_size == fully_shard_degree` incorrectly matched
+        combined modes like SP=4 + HSDP=4.
+        """
+        result = self.compute_standalone_hsdp_params(
+            sequence_parallel_size=4,
+            fully_shard_degree=4,
+            hsdp_replicate_size=1,
+        )
+        assert result["original_dit_parallel_size"] == 4
+        assert result["is_standalone_hsdp"] is False
+        # Should NOT override dp_size for combined mode
+        assert result["effective_dp_size"] == 1  # original data_parallel_size
+
+    def test_combined_hsdp_cfg_not_standalone(self):
+        """Test HSDP combined with CFG is NOT detected as standalone."""
+        result = self.compute_standalone_hsdp_params(
+            cfg_parallel_size=2,
+            fully_shard_degree=4,
+            hsdp_replicate_size=1,
+        )
+        assert result["original_dit_parallel_size"] == 2
+        assert result["is_standalone_hsdp"] is False
+        assert result["effective_dp_size"] == 1
+
+    def test_combined_hsdp_pp_not_standalone(self):
+        """Test HSDP combined with PP is NOT detected as standalone."""
+        result = self.compute_standalone_hsdp_params(
+            pipeline_parallel_size=2,
+            fully_shard_degree=4,
+            hsdp_replicate_size=1,
+        )
+        assert result["original_dit_parallel_size"] == 2
+        assert result["is_standalone_hsdp"] is False
+        assert result["effective_dp_size"] == 1
+
+    def test_no_hsdp_not_standalone(self):
+        """Test that no HSDP (fully_shard_degree=1) is NOT standalone."""
+        result = self.compute_standalone_hsdp_params(
+            fully_shard_degree=1,
+        )
+        assert result["original_dit_parallel_size"] == 1
+        assert result["is_standalone_hsdp"] is False
+        assert result["effective_dp_size"] == 1
+
+    def test_combined_multiple_parallelism_not_standalone(self):
+        """Test HSDP combined with multiple parallelism is NOT standalone."""
+        result = self.compute_standalone_hsdp_params(
+            sequence_parallel_size=2,
+            cfg_parallel_size=2,
+            fully_shard_degree=4,
+            hsdp_replicate_size=1,
+        )
+        assert result["original_dit_parallel_size"] == 4  # 2 * 2
+        assert result["is_standalone_hsdp"] is False
+        assert result["effective_dp_size"] == 1
+
+    def test_standalone_hsdp_large_shard(self):
+        """Test standalone HSDP with large shard size."""
+        result = self.compute_standalone_hsdp_params(
+            fully_shard_degree=8,
+            hsdp_replicate_size=1,
+        )
+        assert result["is_standalone_hsdp"] is True
+        assert result["effective_dit_parallel_size"] == 8
+        assert result["effective_dp_size"] == 8
+
+    def test_standalone_hsdp_large_replicate(self):
+        """Test standalone HSDP with large replicate size."""
+        result = self.compute_standalone_hsdp_params(
+            fully_shard_degree=4,
+            hsdp_replicate_size=4,
+        )
+        assert result["is_standalone_hsdp"] is True
+        assert result["effective_dit_parallel_size"] == 16  # 4 * 4
+        assert result["effective_dp_size"] == 16
+
+
 class TestHSDPShardConditions:
     """Tests for _hsdp_shard_conditions matching logic."""
 

tests/diffusion/lora/test_lora_manager.py

@@ -182,7 +182,11 @@ def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch
     monkeypatch.setattr(manager_mod, "replace_submodule", _fake_replace_submodule)
 
     pipeline = torch.nn.Module()
-    pipeline.packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+    pipeline.stacked_params_mapping = [
+        (".to_qkv.", ".to_q.", "q"),
+        (".to_qkv.", ".to_k.", "k"),
+        (".to_qkv.", ".to_v.", "v"),
+    ]
     pipeline.transformer = torch.nn.Module()
     pipeline.transformer.to_qkv = _FakeLinearBase()
 
@@ -194,7 +198,7 @@ def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch
     )
 
     # Treat the dummy layer as a packed 3-slice projection so the manager uses
-    # `packed_modules_mapping` to decide replacement based on target_modules.
+    # `stacked_params_mapping` to decide replacement based on target_modules.
     monkeypatch.setattr(manager, "_get_packed_modules_list", lambda _module: ["q", "k", "v"])
 
     peft_helper = type("_PH", (), {"r": 1, "target_modules": ["to_q"]})()
@@ -254,7 +258,11 @@ def test_lora_manager_activates_fused_lora_on_packed_layer():
 
 def test_lora_manager_activates_packed_lora_from_sublayers():
     pipeline = torch.nn.Module()
-    pipeline.packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+    pipeline.stacked_params_mapping = [
+        (".to_qkv", ".to_q", "q"),
+        (".to_qkv", ".to_k", "k"),
+        (".to_qkv", ".to_v", "v"),
+    ]
     manager = DiffusionLoRAManager(
         pipeline=pipeline,
         device=torch.device("cpu"),

tests/diffusion/models/bagel/test_combine_cfg.py

@@ -0,0 +1,316 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for Bagel._combine_cfg logic."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.bagel.bagel_transformer import Bagel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestCombineCfg:
+    """Tests for the _combine_cfg static method."""
+
+    def _make_tensors(self, shape=(10, 64), seed=42):
+        """Create deterministic test tensors."""
+        gen = torch.Generator().manual_seed(seed)
+        v_t = torch.randn(shape, generator=gen)
+        cfg_text_v_t = torch.randn(shape, generator=gen)
+        cfg_img_v_t = torch.randn(shape, generator=gen)
+        return v_t, cfg_text_v_t, cfg_img_v_t
+
+    def test_text_channel_renorm_preserves_direction(self):
+        """text_channel renorm should change direction but constrain magnitude."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        # Result norm per token should be <= original v_t norm (clamp max=1.0)
+        result_norm = torch.norm(result, dim=-1)
+        v_t_norm = torch.norm(v_t, dim=-1)
+        assert torch.all(result_norm <= v_t_norm + 1e-6), "text_channel renorm should not increase per-token norm"
+
+    def test_scale_1_returns_v_t(self):
+        """cfg_text_scale=1.0 means no CFG: result should equal v_t."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=1.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        # scale=1 → v_t_text_ = cfg_text + 1*(v_t - cfg_text) = v_t
+        # renorm scale = norm(v_t)/norm(v_t) = 1.0, so result = v_t
+        assert torch.allclose(result, v_t, atol=1e-6)
+
+    def test_img_cfg_applied_when_scale_gt_1(self):
+        """When cfg_img_scale > 1.0, result should differ from text-only CFG."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        text_only = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        with_img = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert not torch.allclose(text_only, with_img, atol=1e-6), (
+            "Image CFG should produce different result from text-only CFG"
+        )
+
+    def test_img_cfg_none_ignored(self):
+        """cfg_img_v_t=None should be equivalent to cfg_img_scale <= 1.0."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result_none = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        result_low_scale = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=0.5,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert torch.allclose(result_none, result_low_scale, atol=1e-6), (
+            "cfg_img_v_t=None and cfg_img_scale<=1.0 should give same result"
+        )
+
+    def test_global_renorm(self):
+        """global renorm should produce valid output without error."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="global",
+            cfg_renorm_min=0.0,
+        )
+
+        assert result.shape == v_t.shape
+        assert not torch.any(torch.isnan(result))
+
+    def test_channel_renorm(self):
+        """channel renorm should produce valid output without error."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert result.shape == v_t.shape
+        assert not torch.any(torch.isnan(result))
+
+    def test_invalid_renorm_type_raises(self):
+        """Unknown renorm type should raise NotImplementedError."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        with pytest.raises(NotImplementedError):
+            Bagel._combine_cfg(
+                v_t,
+                cfg_text_v_t,
+                None,
+                cfg_text_scale=4.0,
+                cfg_img_scale=1.0,
+                cfg_renorm_type="unknown",
+                cfg_renorm_min=0.0,
+            )
+
+    def test_renorm_min_clamps_scale(self):
+        """cfg_renorm_min should prevent scale from going too low."""
+        v_t = torch.ones(10, 64)
+        # Make cfg_text_v_t very different so CFG amplifies heavily
+        cfg_text_v_t = torch.zeros(10, 64)
+
+        result_no_min = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=100.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        result_with_min = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            None,
+            cfg_text_scale=100.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.5,
+        )
+
+        # With higher renorm_min, result magnitude should be larger
+        # (scale is clamped to at least 0.5 instead of going near 0)
+        norm_no_min = torch.norm(result_no_min)
+        norm_with_min = torch.norm(result_with_min)
+        assert norm_with_min >= norm_no_min - 1e-6, "Higher cfg_renorm_min should preserve more magnitude"
+
+    def test_global_renorm_with_img_cfg(self):
+        """global renorm + img CFG should produce valid, different output."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        text_only = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="global",
+            cfg_renorm_min=0.0,
+        )
+
+        with_img = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            cfg_img_v_t.clone(),
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="global",
+            cfg_renorm_min=0.0,
+        )
+
+        assert not torch.allclose(text_only, with_img, atol=1e-6), (
+            "global renorm + img CFG should differ from text-only"
+        )
+        assert not torch.any(torch.isnan(with_img))
+
+    def test_channel_renorm_with_img_cfg(self):
+        """channel renorm + img CFG should produce valid, different output."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        text_only = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            None,
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.0,
+            cfg_renorm_type="channel",
+            cfg_renorm_min=0.0,
+        )
+
+        with_img = Bagel._combine_cfg(
+            v_t.clone(),
+            cfg_text_v_t.clone(),
+            cfg_img_v_t.clone(),
+            cfg_text_scale=4.0,
+            cfg_img_scale=1.5,
+            cfg_renorm_type="channel",
+            cfg_renorm_min=0.0,
+        )
+
+        assert not torch.allclose(text_only, with_img, atol=1e-6), (
+            "channel renorm + img CFG should differ from text-only"
+        )
+        assert not torch.any(torch.isnan(with_img))
+
+    def test_global_channel_renorm_constrains_norm(self):
+        """global and channel renorm should not increase overall norm."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        for renorm_type in ["global", "channel"]:
+            result = Bagel._combine_cfg(
+                v_t.clone(),
+                cfg_text_v_t.clone(),
+                cfg_img_v_t.clone(),
+                cfg_text_scale=4.0,
+                cfg_img_scale=1.5,
+                cfg_renorm_type=renorm_type,
+                cfg_renorm_min=0.0,
+            )
+            # Global norm of result should be <= global norm of v_t (clamp max=1.0)
+            assert torch.norm(result) <= torch.norm(v_t) + 1e-5, f"{renorm_type} renorm should not increase global norm"
+
+    def test_text_channel_img_cfg_no_second_renorm(self):
+        """text_channel mode: img CFG is applied AFTER renorm, without a second renorm.
+        So the result norm can exceed v_t norm when img_scale > 1."""
+        v_t, cfg_text_v_t, cfg_img_v_t = self._make_tensors()
+
+        result = Bagel._combine_cfg(
+            v_t,
+            cfg_text_v_t,
+            cfg_img_v_t,
+            cfg_text_scale=4.0,
+            cfg_img_scale=2.0,
+            cfg_renorm_type="text_channel",
+            cfg_renorm_min=0.0,
+        )
+
+        # text_channel renorms after text CFG, then applies img CFG without renorm
+        # So result norm CAN exceed v_t norm — this is expected behavior
+        assert result.shape == v_t.shape
+        assert not torch.any(torch.isnan(result))
+
+    def test_all_renorm_types_consistent_direction(self):
+        """All renorm types should guide in the same general direction."""
+        v_t, cfg_text_v_t, _ = self._make_tensors()
+
+        results = {}
+        for renorm_type in ["text_channel", "global", "channel"]:
+            results[renorm_type] = Bagel._combine_cfg(
+                v_t.clone(),
+                cfg_text_v_t.clone(),
+                None,
+                cfg_text_scale=4.0,
+                cfg_img_scale=1.0,
+                cfg_renorm_type=renorm_type,
+                cfg_renorm_min=0.0,
+            )
+
+        # All results should have positive cosine similarity with each other
+        for a_name, a in results.items():
+            for b_name, b in results.items():
+                cos_sim = torch.nn.functional.cosine_similarity(a.flatten().unsqueeze(0), b.flatten().unsqueeze(0))
+                assert cos_sim > 0.5, (
+                    f"{a_name} and {b_name} should point in similar direction, "
+                    f"but cosine similarity = {cos_sim.item():.4f}"
+                )

tests/diffusion/test_diffusion_model_runner.py

@@ -56,7 +56,10 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
         enable_cache_dit_summary=enable_cache_dit_summary,
         parallel_config=SimpleNamespace(use_hsdp=False),
     )
-    runner.kv_transfer_manager = SimpleNamespace(receive_kv_cache=lambda req, target_device: None)
+    runner.kv_transfer_manager = SimpleNamespace(
+        receive_kv_cache=lambda req, target_device=None: None,
+        receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None,
+    )
     return runner
 
 
@@ -105,8 +108,8 @@ def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeyp
 
 def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
     class _DummyLoader:
-        def __init__(self, load_config):
-            del load_config
+        def __init__(self, load_config, od_config=None):
+            del load_config, od_config
 
         def load_model(self, **kwargs):
             del kwargs

tests/diffusion/test_worker_wrapper_base.py

@@ -96,6 +96,7 @@ class TestWorkerWrapperBaseInitialization:
             local_rank=0,
             rank=0,
             od_config=mock_od_config,
+            skip_load_model=False,
         )
 
 

tests/distributed/omni_connectors/test_basic_connectors.py

@@ -66,14 +66,13 @@ def shm_connector():
 
 
 def test_put_get_inline(shm_connector):
-    """Test inline transfer for small data."""
+    """Test transfer for small data (currently always uses SHM)."""
     data = {"small": "data"}
-    # Ensure data is smaller than threshold (100 bytes)
 
     success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data)
     assert success is True
-    assert "inline_bytes" in metadata
-    assert "shm" not in metadata
+    assert "shm" in metadata
+    assert "size" in metadata
 
     # Retrieve
     retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata)
@@ -87,7 +86,7 @@ def test_put_get_shm(mocker: MockerFixture, shm_connector, monkeypatch: pytest.M
     data = {"large": "x" * 200}
 
     # Mock SHM return values
-    mock_handle = {"name": "test_shm", "size": 200}
+    mock_handle = {"name": "req_2", "size": 200}
     mock_write = mocker.MagicMock(return_value=mock_handle)
     monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_write_bytes", mock_write)
 

tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml

@@ -4,6 +4,7 @@
 stage_args:
   - stage_id: 0
     stage_type: llm
+    prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
     runtime:
       devices: "0"
       max_batch_size: 1
@@ -39,6 +40,7 @@ stage_args:
       to_stage_1: mooncake_connector
   - stage_id: 1
     stage_type: diffusion
+    cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches
     runtime:
       devices: "0"
       max_batch_size: 1

tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml

@@ -4,6 +4,7 @@
 stage_args:
   - stage_id: 0
     stage_type: llm
+    prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts
     runtime:
       devices: "0"
       max_batch_size: 1
@@ -38,6 +39,7 @@ stage_args:
 
   - stage_id: 1
     stage_type: diffusion
+    cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches
     runtime:
       devices: "0"
       max_batch_size: 1

tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml

@@ -21,6 +21,7 @@ stage_args:
       trust_remote_code: true
       engine_output_type: latent
       enable_prefix_caching: false
+      mm_processor_cache_gb: 0
     is_comprehension: true
     final_output: true
     final_output_type: text

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -37,16 +37,16 @@ from vllm_omni.entrypoints.omni import Omni
 # "Generated with seed=52, num_inference_steps=15,
 # prompt='A futuristic city skyline at twilight, cyberpunk style'"
 REFERENCE_PIXELS = [
-    {"position": (100, 100), "rgb": (68, 107, 134)},
-    {"position": (400, 50), "rgb": (95, 139, 166)},
-    {"position": (700, 100), "rgb": (99, 122, 151)},
-    {"position": (150, 400), "rgb": (111, 125, 153)},
-    {"position": (512, 512), "rgb": (97, 107, 131)},
-    {"position": (700, 400), "rgb": (48, 64, 98)},
-    {"position": (100, 700), "rgb": (79, 63, 84)},
-    {"position": (400, 700), "rgb": (40, 58, 79)},
-    {"position": (700, 700), "rgb": (60, 75, 103)},
-    {"position": (256, 256), "rgb": (97, 128, 156)},
+    {"position": (100, 100), "rgb": (49, 96, 134)},
+    {"position": (400, 50), "rgb": (63, 127, 167)},
+    {"position": (700, 100), "rgb": (70, 101, 141)},
+    {"position": (150, 400), "rgb": (115, 90, 150)},
+    {"position": (512, 512), "rgb": (98, 86, 119)},
+    {"position": (700, 400), "rgb": (29, 42, 91)},
+    {"position": (100, 700), "rgb": (47, 50, 88)},
+    {"position": (400, 700), "rgb": (36, 52, 91)},
+    {"position": (700, 700), "rgb": (45, 58, 99)},
+    {"position": (256, 256), "rgb": (62, 94, 135)},
 ]
 
 # Maximum allowed difference per color channel
@@ -80,6 +80,10 @@ def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_st
     params_list[0].max_tokens = max_tokens  # type: ignore
     if len(params_list) > 1:
         params_list[1].num_inference_steps = num_inference_steps  # type: ignore
+        params_list[1].extra_args = {  # type: ignore
+            "cfg_text_scale": 4.0,
+            "cfg_img_scale": 1.5,
+        }
     return params_list
 
 

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -1,3 +1,4 @@
+import gc
 import sys
 from pathlib import Path
 
@@ -20,8 +21,10 @@ models = ["riverclouds/qwen_image_random"]
 
 
 def inference(model_name: str, offload: bool = True):
+    gc.collect()
     current_omni_platform.empty_cache()
     device_index = torch.cuda.current_device()
+    torch.cuda.reset_peak_memory_stats(device=device_index)
     monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
     m = Omni(model=model_name, enable_cpu_offload=offload)
@@ -42,6 +45,10 @@ def inference(model_name: str, offload: bool = True):
     peak = monitor.peak_used_mb
     monitor.stop()
 
+    del m
+    gc.collect()
+    current_omni_platform.empty_cache()
+
     return peak
 
 
@@ -51,9 +58,9 @@ def inference(model_name: str, offload: bool = True):
 @pytest.mark.parametrize("model_name", models)
 def test_cpu_offload_diffusion_model(model_name: str):
     try:
-        no_offload_peak_memory = inference(model_name, offload=False)
-        cleanup_dist_env_and_memory()
         offload_peak_memory = inference(model_name, offload=True)
+        cleanup_dist_env_and_memory()
+        no_offload_peak_memory = inference(model_name, offload=False)
     except Exception:
         pytest.fail("Inference failed")
     print(f"Offload peak memory: {offload_peak_memory} MB")

tests/e2e/offline_inference/test_diffusion_lora.py

@@ -9,6 +9,7 @@ from safetensors.torch import save_file
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -89,7 +90,7 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
                 width=width,
                 num_inference_steps=2,
                 guidance_scale=0.0,
-                generator=torch.Generator("cuda").manual_seed(42),
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
                 num_outputs_per_prompt=1,
             ),
         )
@@ -119,7 +120,7 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
                     width=width,
                     num_inference_steps=2,
                     guidance_scale=0.0,
-                    generator=torch.Generator("cuda").manual_seed(42),
+                    generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
                     num_outputs_per_prompt=1,
                     lora_request=lora_request,
                     lora_scale=2.0,

tests/e2e/offline_inference/test_mammoth_moda2.py

@@ -0,0 +1,215 @@
+"""
+End-to-end test for MammothModa2 text-to-image generation.
+
+Verifies that the AR->DiT pipeline produces an image tensor whose pixel values
+match a golden reference.
+
+Model weights: env var ``MAMMOTHMODA2_MODEL_PATH``
+  (default: <repo_root>/MammothModa2-Preview)
+Stage config:  env var ``MAMMOTHMODA2_T2I_STAGE_CONFIG``
+  (default: vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml)
+
+Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json``
+  Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2.py``
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import pytest
+import torch
+from vllm.sampling_params import SamplingParams
+
+from tests.utils import hardware_test
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+_IMAGE_TOKEN_ID = 151655  # "<|image_pad|>"
+_VIDEO_TOKEN_ID = 151656  # "<|video_pad|>"
+_VISION_START_TOKEN_ID = 151652  # "<|vision_start|>"
+_VISION_END_TOKEN_ID = 151653  # "<|vision_end|>"
+_AR_PATCH_SIZE = 16
+
+_STAGE_CONFIGS_DIR = Path(__file__).resolve().parents[3] / "vllm_omni" / "model_executor" / "stage_configs"
+MODEL_PATH = os.environ.get(
+    "MAMMOTHMODA2_MODEL_PATH",
+    str(Path(__file__).resolve().parents[3] / "MammothModa2-Preview"),
+)
+T2I_STAGE_CONFIG = os.environ.get(
+    "MAMMOTHMODA2_T2I_STAGE_CONFIG",
+    str(_STAGE_CONFIGS_DIR / "mammoth_moda2.yaml"),
+)
+
+# Golden pixel reference file.  Set UPDATE_GOLDEN=1 to regenerate.
+_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
+# Fixed sampling coordinates: (channel, row_fraction, col_fraction)
+# Covers corners, centre, and mid-edges across all 3 channels.
+_PIXEL_SAMPLE_COORDS = [
+    (0, 0.0, 0.0),
+    (0, 0.5, 0.5),
+    (0, 1.0, 1.0),
+    (0, 0.25, 0.75),
+    (1, 0.0, 1.0),
+    (1, 0.5, 0.0),
+    (1, 0.75, 0.25),
+    (1, 1.0, 0.5),
+    (2, 0.0, 0.5),
+    (2, 0.5, 1.0),
+    (2, 0.75, 0.75),
+    (2, 1.0, 0.0),
+]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+def _load_t2i_gen_config(model_dir: str) -> dict:
+    cfg_path = Path(model_dir) / "t2i_generation_config.json"
+    if not cfg_path.exists():
+        pytest.skip(f"t2i_generation_config.json not found at {cfg_path}")
+    with cfg_path.open() as f:
+        return json.load(f)
+
+
+def _format_t2i_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str:
+    return (
+        "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
+        f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+        f"<|image start|>{ar_width}*{ar_height}<|image token|>"
+    )
+
+
+def _sample_pixels(img_tensor: torch.Tensor) -> list[float]:
+    """Sample pixel values at fixed fractional coordinates from a (C, H, W) tensor."""
+    t = img_tensor.float().clamp(0.0, 1.0)
+    if t.ndim == 4:
+        t = t[0]  # unbatch
+    C, H, W = t.shape
+    values = []
+    for c, rh, rw in _PIXEL_SAMPLE_COORDS:
+        ri = min(int(rh * (H - 1)), H - 1)
+        ci = min(int(rw * (W - 1)), W - 1)
+        values.append(round(float(t[c, ri, ci]), 6))
+    return values
+
+
+# ---------------------------------------------------------------------------
+# End-to-end test
+# ---------------------------------------------------------------------------
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"})
+def test_mammothmoda2_t2i_e2e():
+    """
+    End-to-end text-to-image generation with MammothModa2 (AR -> DiT).
+
+    Verifies:
+      - Omni pipeline initialises with the two-stage YAML config.
+      - DiT stage outputs an image tensor with the correct shape.
+      - A fixed set of pixel values matches a golden reference
+        (regenerate with ``UPDATE_GOLDEN=1``).
+    """
+    from vllm_omni import Omni
+
+    if not Path(MODEL_PATH).exists():
+        pytest.skip(f"Model weights not found at {MODEL_PATH}")
+    if not Path(T2I_STAGE_CONFIG).exists():
+        pytest.skip(f"Stage config not found at {T2I_STAGE_CONFIG}")
+
+    gen_cfg = _load_t2i_gen_config(MODEL_PATH)
+    eol_token_id = int(gen_cfg["eol_token_id"])
+    visual_start = int(gen_cfg["visual_token_start_id"])
+    visual_end = int(gen_cfg["visual_token_end_id"])
+
+    height, width = 256, 256  # small for CI speed
+    ar_height, ar_width = height // _AR_PATCH_SIZE, width // _AR_PATCH_SIZE
+    expected_grid_tokens = ar_height * (ar_width + 1)
+
+    prompt_text = "A cat sitting on a laptop keyboard"
+    formatted_prompt = _format_t2i_prompt(prompt_text, ar_width, ar_height)
+
+    omni = Omni(model=MODEL_PATH, stage_configs_path=T2I_STAGE_CONFIG, trust_remote_code=True)
+    try:
+        # Greedy / deterministic sampling so pixel values are reproducible.
+        ar_sampling = SamplingParams(
+            temperature=0.0,
+            top_k=1,
+            max_tokens=max(1, expected_grid_tokens + 1),
+            detokenize=False,
+        )
+        dit_sampling = SamplingParams(temperature=0.0, max_tokens=1, detokenize=False)
+
+        outputs = list(
+            omni.generate(
+                [
+                    {
+                        "prompt": formatted_prompt,
+                        "additional_information": {
+                            "omni_task": ["t2i"],
+                            "ar_width": [ar_width],
+                            "ar_height": [ar_height],
+                            "eol_token_id": [eol_token_id],
+                            "visual_token_start_id": [visual_start],
+                            "visual_token_end_id": [visual_end],
+                            "image_height": [height],
+                            "image_width": [width],
+                            "num_inference_steps": [2],
+                            "text_guidance_scale": [1.0],
+                            "cfg_range": [0.0, 1.0],
+                            "visual_ids": [
+                                _IMAGE_TOKEN_ID,
+                                _VIDEO_TOKEN_ID,
+                                _VISION_START_TOKEN_ID,
+                                _VISION_END_TOKEN_ID,
+                            ],
+                        },
+                    }
+                ],
+                [ar_sampling, dit_sampling],
+            )
+        )
+
+        assert len(outputs) > 0, "Pipeline produced no outputs"
+
+        # Find the image tensor produced by the DiT stage and compare pixels.
+        found_image = False
+        for out in outputs:
+            ro_list = getattr(out, "request_output", out)
+            if not isinstance(ro_list, list):
+                ro_list = [ro_list]
+            for ro in ro_list:
+                completion_outputs = getattr(ro, "outputs", None)
+                if not isinstance(completion_outputs, list):
+                    continue
+                for completion in completion_outputs:
+                    mm = getattr(completion, "multimodal_output", None)
+                    if not (isinstance(mm, dict) and "image" in mm):
+                        continue
+                    img_list = mm["image"] if isinstance(mm["image"], list) else [mm["image"]]
+                    for img_tensor in img_list:
+                        assert isinstance(img_tensor, torch.Tensor), f"Expected image tensor, got {type(img_tensor)}"
+                        assert img_tensor.ndim in (3, 4), f"Expected 3D or 4D image tensor, got {img_tensor.ndim}D"
+
+                        sampled = _sample_pixels(img_tensor)
+
+                        if os.environ.get("UPDATE_GOLDEN"):
+                            _GOLDEN_T2I_PATH.parent.mkdir(parents=True, exist_ok=True)
+                            _GOLDEN_T2I_PATH.write_text(json.dumps({"pixels": sampled}, indent=2))
+                            print(f"\nGolden file written to {_GOLDEN_T2I_PATH}")
+                        elif _GOLDEN_T2I_PATH.exists():
+                            golden = json.loads(_GOLDEN_T2I_PATH.read_text())["pixels"]
+                            for i, (got, exp) in enumerate(zip(sampled, golden)):
+                                assert abs(got - exp) < 1e-4, f"Pixel {i} mismatch: got {got}, expected {exp}"
+
+                        found_image = True
+
+        assert found_image, "No image tensor found in pipeline output"
+    finally:
+        omni.close()

tests/e2e/offline_inference/test_stable_audio_model.py

@@ -8,6 +8,7 @@ import torch
 from tests.utils import hardware_test
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -41,7 +42,7 @@ def test_stable_audio_model(model_name: str):
         sampling_params_list=OmniDiffusionSamplingParams(
             num_inference_steps=4,  # Minimal steps for speed
             guidance_scale=7.0,
-            generator=torch.Generator("cuda").manual_seed(42),
+            generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
             num_outputs_per_prompt=1,
             extra_args={
                 "audio_start_in_s": audio_start_in_s,

tests/e2e/offline_inference/test_teacache.py

@@ -17,6 +17,7 @@ import torch
 
 from tests.utils import hardware_test
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -63,7 +64,7 @@ def test_teacache(model_name: str):
                 width=width,
                 num_inference_steps=num_inference_steps,
                 guidance_scale=0.0,
-                generator=torch.Generator("cuda").manual_seed(42),
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
                 num_outputs_per_prompt=1,  # Single output for speed
             ),
         )

tests/e2e/offline_inference/test_vae_decode_parallelism.py

@@ -0,0 +1,154 @@
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+import time
+
+from vllm_omni import Omni
+from vllm_omni.platforms import current_omni_platform
+
+# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+all_model_configs = [
+    {
+        "model_name": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+        "out_width": 1280,
+        "out_height": 720,
+        "out_frames": 16,
+    },
+    {
+        "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+        "out_width": 1280,
+        "out_height": 720,
+        "out_frames": 16,
+    },
+    {
+        "model_name": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+        "out_width": 1280,
+        "out_height": 720,
+        "out_frames": 16,
+    },
+    {
+        "model_name": "Qwen/Qwen-Image",
+        "out_width": 3072,
+        "out_height": 3072,
+        "out_frames": 1,
+    },
+    # {
+    #     "model_name": "stepfun-ai/NextStep-1.1",
+    #     "out_width": 3072,
+    #     "out_height": 3072,
+    #     "out_frames": 1,
+    # },
+]
+
+
+def is_nextstep_model(model_name: str) -> bool:
+    """Check if the model is a NextStep model by reading its config."""
+    from vllm.transformers_utils.config import get_hf_file_to_dict
+
+    try:
+        cfg = get_hf_file_to_dict("config.json", model_name)
+        if cfg and cfg.get("model_type") == "nextstep":
+            return True
+    except Exception:
+        pass
+    return False
+
+
+def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
+    try:
+        parallel_config = DiffusionParallelConfig(
+            tensor_parallel_size=tp,
+            vae_patch_parallel_size=vae_patch_parallel_size,
+        )
+
+        omni_kwargs = {
+            "model": model_configs["model_name"],
+            "vae_use_tiling": using_tile,
+            "parallel_config": parallel_config,
+        }
+        use_nextstep = is_nextstep_model(model_configs["model_name"])
+        if use_nextstep:
+            # NextStep-1.1 requires explicit pipeline class
+            omni_kwargs["model_class_name"] = "NextStep11Pipeline"
+        m = Omni(**omni_kwargs)
+        image = Image.new("RGB", (out_width, out_height), (0, 0, 0))
+        start = time.perf_counter()
+        outputs = m.generate(
+            {
+                "prompt": "A cat sitting on a table",
+                "multi_modal_data": {"image": image},
+            },
+            sampling_params_list=OmniDiffusionSamplingParams(
+                height=model_configs["out_height"],
+                width=model_configs["out_width"],
+                num_frames=model_configs["out_frames"],
+                num_inference_steps=2,
+                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
+            ),
+        )
+        end = time.perf_counter()
+        first_output = outputs[0]
+        req_out = first_output.request_output[0]
+        frames = req_out.images[0]
+        if isinstance(frames, torch.Tensor):
+            frames = frames.detach().cpu().numpy()
+        elif isinstance(frames, Image.Image):
+            frames = np.array(frames)
+        # frames shape: (batch, num_frames, height, width, channels)
+        cost = (end - start) * 1000
+        return frames, cost
+    finally:
+        m.close()
+        cleanup_dist_env_and_memory()
+
+
+@pytest.mark.parametrize("model_configs", all_model_configs)
+def test_vae_parallel_model(model_configs: str):
+    parallel_result, parallel_time = model_run(
+        model_configs=model_configs,
+        tp=2,
+        out_width=1280,
+        out_height=704,
+        out_frames=5,
+        using_tile=True,
+        vae_patch_parallel_size=2,
+    )
+    non_parallel_result, non_parallel_time = model_run(
+        model_configs=model_configs,
+        tp=2,
+        out_width=1280,
+        out_height=704,
+        out_frames=5,
+        using_tile=True,
+        vae_patch_parallel_size=1,
+    )
+    result_diff = np.abs(non_parallel_result - parallel_result)
+
+    mean_threshold = 3e-2
+    max_threshold = 3e-2  # they should be totally same
+    print(
+        f"{model_configs['model_name']} TP = 2 (tile + parallel vs tile): "
+        f"mean_abs_diff={result_diff.mean():.6e}, max_abs_diff={result_diff.max():.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
+        f"parallel generate take time: {parallel_time:.2f} ms, non-parallel take time: {non_parallel_time:.2f} ms"
+    )
+    # assert non_parallel_time > parallel_time
+    assert result_diff.mean() < mean_threshold
+    assert result_diff.max() < max_threshold

tests/e2e/online_serving/test_qwen3_tts.py

@@ -23,14 +23,12 @@ from tests.utils import hardware_test
 MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
 
 
-def get_stage_config():
+def get_stage_config(name: str = "qwen3_tts.yaml"):
     """Get the stage config path for Qwen3-TTS."""
-    return str(
-        Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "qwen3_tts.yaml"
-    )
+    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="class")
 def omni_server():
     """Start vLLM-Omni server with CustomVoice model."""
     stage_config_path = get_stage_config()
@@ -225,3 +223,63 @@ class TestQwen3TTSAPIEndpoints:
         data = response.json()
         assert "data" in data
         assert len(data["data"]) > 0
+
+
+@pytest.fixture(scope="class")
+def omni_server_no_async_chunk():
+    """Start vLLM-Omni server with non-async-chunk config."""
+    stage_config_path = get_stage_config("qwen3_tts_no_async_chunk.yaml")
+
+    with OmniServer(
+        MODEL,
+        [
+            "--stage-configs-path",
+            stage_config_path,
+            "--stage-init-timeout",
+            "120",
+            "--trust-remote-code",
+            "--enforce-eager",
+            "--disable-log-stats",
+        ],
+    ) as server:
+        yield server
+
+
+class TestQwen3TTSNoAsyncChunk:
+    """E2E tests for Qwen3-TTS in non-async-chunk (full decode) mode."""
+
+    @pytest.mark.core_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "L4"}, num_cards=4)
+    def test_speech_english(self, omni_server_no_async_chunk) -> None:
+        """Test English TTS with non-async-chunk pipeline."""
+        response = make_speech_request(
+            host=omni_server_no_async_chunk.host,
+            port=omni_server_no_async_chunk.port,
+            text="Hello, how are you?",
+            voice="vivian",
+            language="English",
+        )
+
+        assert response.status_code == 200, f"Request failed: {response.text}"
+        assert response.headers.get("content-type") == "audio/wav"
+        assert verify_wav_audio(response.content), "Response is not valid WAV audio"
+        assert len(response.content) > MIN_AUDIO_BYTES
+
+    @pytest.mark.core_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "L4"}, num_cards=4)
+    def test_speech_chinese(self, omni_server_no_async_chunk) -> None:
+        """Test Chinese TTS with non-async-chunk pipeline."""
+        response = make_speech_request(
+            host=omni_server_no_async_chunk.host,
+            port=omni_server_no_async_chunk.port,
+            text="你好，我是通义千问",
+            voice="vivian",
+            language="Chinese",
+        )
+
+        assert response.status_code == 200, f"Request failed: {response.text}"
+        assert response.headers.get("content-type") == "audio/wav"
+        assert verify_wav_audio(response.content), "Response is not valid WAV audio"
+        assert len(response.content) > MIN_AUDIO_BYTES

tests/e2e/online_serving/test_qwen3_tts_base.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E tests for Qwen3-TTS Base voice-clone model.
+
+Regression test for #1663: speech tokenizer loaded in bfloat16 produced
+all-silence PCM due to NaN overflow in SnakeBeta activation.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+import struct
+import tempfile
+from pathlib import Path
+
+import httpx
+import pytest
+
+from tests.conftest import (
+    OmniServer,
+    convert_audio_file_to_text,
+    cosine_similarity_text,
+)
+from tests.utils import hardware_test
+
+MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
+
+# Official Qwen3-TTS reference audio/text from examples/offline_inference/qwen3_tts/end2end.py
+REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/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."
+SYN_TEXT = "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye."
+
+
+def get_stage_config():
+    return str(
+        Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "qwen3_tts.yaml"
+    )
+
+
+@pytest.fixture(scope="module")
+def omni_server():
+    """Start vLLM-Omni server with Base voice-clone model."""
+    stage_config_path = get_stage_config()
+
+    with OmniServer(
+        MODEL,
+        [
+            "--stage-configs-path",
+            stage_config_path,
+            "--stage-init-timeout",
+            "120",
+            "--trust-remote-code",
+            "--enforce-eager",
+            "--disable-log-stats",
+        ],
+    ) as server:
+        yield server
+
+
+def make_base_speech_request(
+    host: str,
+    port: int,
+    text: str = SYN_TEXT,
+    ref_text: str = REF_TEXT,
+    ref_audio: str = REF_AUDIO_URL,
+    response_format: str = "wav",
+    timeout: float = 120.0,
+) -> httpx.Response:
+    url = f"http://{host}:{port}/v1/audio/speech"
+    payload = {
+        "model": MODEL,
+        "input": text,
+        "task_type": "Base",
+        "ref_text": ref_text,
+        "ref_audio": ref_audio,
+        "response_format": response_format,
+    }
+    with httpx.Client(timeout=timeout) as client:
+        return client.post(url, json=payload)
+
+
+def verify_wav_audio(content: bytes) -> bool:
+    if len(content) < 44:
+        return False
+    return content[:4] == b"RIFF" and content[8:12] == b"WAVE"
+
+
+def assert_not_silence(pcm_bytes: bytes):
+    """Assert PCM16 samples are not all identical (e.g. all-silence)."""
+    samples = struct.unpack(f"<{len(pcm_bytes) // 2}h", pcm_bytes)
+    unique = set(samples)
+    assert len(unique) > 1, (
+        f"All-silence detected: {len(samples)} samples, unique values: {unique}. "
+        "See https://github.com/vllm-project/vllm-omni/issues/1663"
+    )
+
+
+MIN_AUDIO_BYTES = 10000
+
+
+class TestQwen3TTSBaseVoiceClone:
+    """Regression tests for Base voice-clone (fix #1663)."""
+
+    @pytest.mark.core_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "L4"}, num_cards=4)
+    def test_base_voice_clone_not_silence(self, omni_server) -> None:
+        """PCM output must contain real audio, not all-silence."""
+        response = make_base_speech_request(
+            host=omni_server.host,
+            port=omni_server.port,
+            response_format="pcm",
+        )
+
+        assert response.status_code == 200, f"Request failed: {response.text}"
+        assert len(response.content) > MIN_AUDIO_BYTES, f"Audio too small: {len(response.content)} bytes"
+        assert_not_silence(response.content)
+
+    @pytest.mark.core_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "L4"}, num_cards=4)
+    def test_base_voice_clone_whisper_transcription(self, omni_server) -> None:
+        """Whisper must transcribe the output as intelligible speech."""
+        response = make_base_speech_request(
+            host=omni_server.host,
+            port=omni_server.port,
+            response_format="wav",
+        )
+
+        assert response.status_code == 200, f"Request failed: {response.text}"
+        assert verify_wav_audio(response.content), "Response is not valid WAV"
+        assert len(response.content) > MIN_AUDIO_BYTES
+
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+            f.write(response.content)
+            wav_path = f.name
+
+        try:
+            transcript = convert_audio_file_to_text(wav_path)
+            print(f"Whisper transcript: {transcript}")
+            assert len(transcript.strip()) > 0, "Whisper returned empty transcript — audio is likely silence"
+            similarity = cosine_similarity_text(transcript.lower(), SYN_TEXT.lower())
+            print(f"Cosine similarity: {similarity:.3f}")
+            assert similarity > 0.9, (
+                f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'"
+            )
+        finally:
+            os.unlink(wav_path)
+
+    @pytest.mark.core_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "L4"}, num_cards=4)
+    def test_base_voice_clone_wav_format(self, omni_server) -> None:
+        """WAV response must have valid headers and sufficient size."""
+        response = make_base_speech_request(
+            host=omni_server.host,
+            port=omni_server.port,
+            response_format="wav",
+        )
+
+        assert response.status_code == 200, f"Request failed: {response.text}"
+        assert response.headers.get("content-type") == "audio/wav"
+        assert verify_wav_audio(response.content), "Response is not valid WAV"
+        assert len(response.content) > MIN_AUDIO_BYTES

tests/e2e/stage_configs/qwen2_5_omni_ci.yaml

@@ -74,12 +74,13 @@ stage_args:
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      gpu_memory_utilization: 0.15
+      gpu_memory_utilization: 0.9 #increase the gpu memory utilization to enable the test on H800
       enforce_eager: true
       trust_remote_code: true
       enable_prefix_caching: false
       engine_output_type: audio
-      max_num_batched_tokens: 100000
+      max_num_batched_tokens: 4096
+      max_model_len: 4096
     engine_input_source: [1]
     final_output: true
     final_output_type: audio

tests/e2e/stage_configs/qwen2_5_omni_thinker_ci.yaml

@@ -18,6 +18,7 @@ stage_args:
       trust_remote_code: true
       engine_output_type: latent
       enable_prefix_caching: false
+      mm_processor_cache_gb: 0
     is_comprehension: true
     final_output: true
     final_output_type: text

tests/e2e/stage_configs/qwen3_omni_ci.yaml

@@ -22,6 +22,7 @@ stage_args:
     max_num_batched_tokens: 32768
     max_model_len: 32768
     enable_prefix_caching: false
+    mm_processor_cache_gb: 0
     hf_config_name: thinker_config
     tensor_parallel_size: 1
     load_format: dummy
@@ -47,7 +48,7 @@ stage_args:
     model_arch: Qwen3OmniMoeForConditionalGeneration
     worker_type: ar
     scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    gpu_memory_utilization: 0.6
+    gpu_memory_utilization: 0.5
     enforce_eager: false
     trust_remote_code: true
     engine_output_type: latent  # Output codec codes for code2wav

tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml

@@ -22,6 +22,7 @@ stage_args:
       trust_remote_code: true
       engine_output_type: latent
       enable_prefix_caching: false
+      mm_processor_cache_gb: 0
     is_comprehension: true
     final_output: true
     final_output_type: text

tests/e2e/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -20,6 +20,7 @@ stage_args:
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
+      mm_processor_cache_gb: 0
       hf_config_name: thinker_config
       tensor_parallel_size: 1
       load_format: dummy

tests/entrypoints/openai_api/test_serving_speech.py

@@ -1,11 +1,15 @@
 # tests/entrypoints/openai/test_serving_speech.py
 import logging
+import os
 from inspect import Signature, signature
+from pathlib import Path
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
 import torch
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException, UploadFile
+from fastapi.params import File, Form
 from fastapi.testclient import TestClient
 from pydantic import ValidationError
 from pytest_mock import MockerFixture
@@ -206,10 +210,52 @@ def test_app(mocker: MockerFixture):
     # Add list_voices endpoint
     async def list_voices():
         speakers = sorted(speech_server.supported_speakers) if speech_server.supported_speakers else []
-        return {"voices": speakers}
+        uploaded_voices = []
+        if hasattr(speech_server, "uploaded_speakers"):
+            for voice_name, info in speech_server.uploaded_speakers.items():
+                uploaded_voices.append(
+                    {
+                        "name": info.get("name", voice_name),
+                        "consent": info.get("consent", ""),
+                        "created_at": info.get("created_at", 0),
+                        "file_size": info.get("file_size", 0),
+                        "mime_type": info.get("mime_type", ""),
+                    }
+                )
+        return {"voices": speakers, "uploaded_voices": uploaded_voices}
 
     app.add_api_route("/v1/audio/voices", list_voices, methods=["GET"])
 
+    # Add upload_voice endpoint
+    async def upload_voice(audio_sample: UploadFile = File(...), consent: str = Form(...), name: str = Form(...)):
+        try:
+            result = await speech_server.upload_voice(audio_sample, consent, name)
+            return {"success": True, "voice": result}
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+        except Exception as e:
+            logger.exception(f"Failed to upload voice: {e}")
+            raise HTTPException(status_code=500, detail=f"Failed to upload voice: {str(e)}")
+
+    app.add_api_route("/v1/audio/voices", upload_voice, methods=["POST"])
+
+    # Add delete_voice endpoint
+    async def delete_voice(name: str):
+        try:
+            success = await speech_server.delete_voice(name)
+            if not success:
+                raise HTTPException(status_code=404, detail=f"Voice '{name}' not found")
+            return {"success": True, "message": f"Voice '{name}' deleted successfully"}
+        except HTTPException:
+            raise
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+        except Exception as e:
+            logger.exception(f"Failed to delete voice '{name}': {e}")
+            raise HTTPException(status_code=500, detail=f"Failed to delete voice: {str(e)}")
+
+    app.add_api_route("/v1/audio/voices/{name}", delete_voice, methods=["DELETE"])
+
     return app
 
 
@@ -284,6 +330,149 @@ class TestSpeechAPI:
         assert response.status_code == 200
         assert "voices" in response.json()
 
+    def test_upload_voice_success(self, client, tmp_path):
+        """Test successful voice upload."""
+        # Create a mock audio file
+        audio_content = b"fake audio content" * 1000  # ~17KB
+        files = {
+            "audio_sample": ("test.wav", audio_content, "audio/wav"),
+        }
+        data = {
+            "consent": "user_consent_123",
+            "name": "test_voice",
+        }
+
+        response = client.post("/v1/audio/voices", files=files, data=data)
+        assert response.status_code == 200
+        result = response.json()
+        assert result["success"] is True
+        assert "voice" in result
+        voice_info = result["voice"]
+        assert voice_info["name"] == "test_voice"
+        assert voice_info["consent"] == "user_consent_123"
+        assert "created_at" in voice_info
+        assert voice_info["mime_type"] == "audio/wav"
+        assert voice_info["file_size"] == len(audio_content)
+        response = client.delete("/v1/audio/voices/test_voice")
+
+    def test_upload_voice_file_too_large(self, client):
+        """Test voice upload with file exceeding size limit."""
+        # Create a file larger than 10MB
+        audio_content = b"x" * (11 * 1024 * 1024)  # 11MB
+        files = {
+            "audio_sample": ("test.wav", audio_content, "audio/wav"),
+        }
+        data = {
+            "consent": "user_consent_123",
+            "name": "test_voice",
+        }
+
+        response = client.post("/v1/audio/voices", files=files, data=data)
+        assert response.status_code == 400
+        result = response.json()
+        assert "detail" in result
+        assert "10MB" in result["detail"]
+
+    def test_upload_voice_invalid_mime_type(self, client):
+        """Test voice upload with invalid MIME type."""
+        audio_content = b"fake audio content"
+        files = {
+            "audio_sample": ("test.txt", audio_content, "text/plain"),
+        }
+        data = {
+            "consent": "user_consent_123",
+            "name": "test_voice",
+        }
+
+        response = client.post("/v1/audio/voices", files=files, data=data)
+        assert response.status_code == 400
+        result = response.json()
+        assert "detail" in result
+        assert "MIME type" in result["detail"]
+
+    def test_upload_voice_name_collision(self, client):
+        """Test voice upload with duplicate name."""
+        # First upload
+        audio_content = b"fake audio content"
+        files = {
+            "audio_sample": ("test.wav", audio_content, "audio/wav"),
+        }
+        data = {
+            "consent": "user_consent_123",
+            "name": "test_voice",
+        }
+
+        response = client.post("/v1/audio/voices", files=files, data=data)
+        assert response.status_code == 200
+
+        # Second upload with same name
+        response = client.post("/v1/audio/voices", files=files, data=data)
+        assert response.status_code == 400
+        result = response.json()
+        assert "detail" in result
+        assert "already exists" in result["detail"]
+        response = client.delete("/v1/audio/voices/test_voice")
+
+    def test_upload_voice_missing_parameters(self, client):
+        """Test voice upload with missing required parameters."""
+        audio_content = b"fake audio content"
+        files = {
+            "audio_sample": ("test.wav", audio_content, "audio/wav"),
+        }
+
+        # Missing consent
+        data = {"name": "test_voice5"}
+        response = client.post("/v1/audio/voices", files=files, data=data)
+        assert response.status_code == 422  # Validation error
+
+        # Missing name
+        data = {"consent": "user_consent_123"}
+        response = client.post("/v1/audio/voices", files=files, data=data)
+        assert response.status_code == 422  # Validation error
+
+        # Missing file
+        data = {
+            "consent": "user_consent_123",
+            "name": "test_voice6",
+        }
+        response = client.post("/v1/audio/voices", data=data)
+        assert response.status_code == 422  # Validation error
+
+    def test_delete_voice_success(self, client):
+        """Test successful voice deletion."""
+        # First upload a voice
+        audio_content = b"fake audio content"
+        files = {
+            "audio_sample": ("test.wav", audio_content, "audio/wav"),
+        }
+        data = {
+            "consent": "user_consent_123",
+            "name": "test_voice7",
+        }
+
+        response = client.post("/v1/audio/voices", files=files, data=data)
+        assert response.status_code == 200
+
+        # Then delete it
+        response = client.delete("/v1/audio/voices/test_voice7")
+        assert response.status_code == 200
+        result = response.json()
+        assert result["success"] is True
+        assert "deleted successfully" in result["message"]
+
+        # Verify it's gone by trying to delete again
+        response = client.delete("/v1/audio/voices/test_voice7")
+        assert response.status_code == 404
+        result = response.json()
+        assert "not found" in result["detail"]
+
+    def test_delete_voice_not_found(self, client):
+        """Test deleting a non-existent voice."""
+        response = client.delete("/v1/audio/voices/nonexistent")
+        assert response.status_code == 404
+        result = response.json()
+        assert "not found" in result["detail"]
+
 
 class TestTTSMethods:
     """Unit tests for TTS validation and parameter building."""
@@ -349,13 +538,13 @@ class TestTTSMethods:
         req = OpenAICreateSpeechRequest(input="Hello", language="InvalidLang")
         assert "Invalid language" in speech_server._validate_tts_request(req)
 
-        # When no speakers loaded, any voice is accepted (unconstrained)
+        # CustomVoice on model with no speakers -> rejected
         req = OpenAICreateSpeechRequest(input="Hello", voice="Invalid")
-        assert speech_server._validate_tts_request(req) is None
+        assert "does not support CustomVoice" in speech_server._validate_tts_request(req)
 
-        # Valid request
-        req = OpenAICreateSpeechRequest(input="Hello", voice="Vivian")
-        assert speech_server._validate_tts_request(req) is None
+        # CustomVoice without voice on model with no speakers -> also rejected
+        req = OpenAICreateSpeechRequest(input="Hello")
+        assert "does not support CustomVoice" in speech_server._validate_tts_request(req)
 
     def test_validate_tts_request_task_types(self, speech_server):
         """Test task-specific validation."""
@@ -367,9 +556,43 @@ class TestTTSMethods:
         req = OpenAICreateSpeechRequest(input="Hello", task_type="VoiceDesign")
         assert "instructions" in speech_server._validate_tts_request(req)
 
-        # ref_text only for Base
+        # ref_text without task_type auto-infers Base, then fails on missing ref_audio
         req = OpenAICreateSpeechRequest(input="Hello", ref_text="text")
-        assert "Base task" in speech_server._validate_tts_request(req)
+        assert "ref_audio" in speech_server._validate_tts_request(req)
+
+    def test_validate_tts_request_auto_infer_base(self, speech_server):
+        """Test auto-inference of Base task when ref_audio/ref_text is provided."""
+        # ref_audio without task_type -> infers Base, requires non-empty ref_text
+        req = OpenAICreateSpeechRequest(input="Hello", ref_audio="data:audio/wav;base64,abc")
+        result = speech_server._validate_tts_request(req)
+        assert "ref_text" in result
+        assert req.task_type == "Base"
+
+        # ref_text without task_type -> infers Base, requires ref_audio
+        req = OpenAICreateSpeechRequest(input="Hello", ref_text="some text")
+        result = speech_server._validate_tts_request(req)
+        assert "ref_audio" in result
+        assert req.task_type == "Base"
+
+    def test_validate_tts_request_base_empty_ref_text(self, speech_server):
+        """Empty ref_text on Base task returns 400 instead of crashing engine."""
+        req = OpenAICreateSpeechRequest(
+            input="Hello", task_type="Base", ref_audio="data:audio/wav;base64,abc", ref_text=""
+        )
+        result = speech_server._validate_tts_request(req)
+        assert "non-empty 'ref_text'" in result
+
+        # x_vector_only_mode bypasses ref_text requirement
+        req = OpenAICreateSpeechRequest(
+            input="Hello", task_type="Base", ref_audio="data:audio/wav;base64,abc", ref_text="", x_vector_only_mode=True
+        )
+        assert speech_server._validate_tts_request(req) is None
+
+    def test_validate_tts_request_customvoice_no_speakers(self, speech_server):
+        """CustomVoice on a model with no speakers returns 400 instead of crashing engine."""
+        req = OpenAICreateSpeechRequest(input="Hello", task_type="CustomVoice")
+        result = speech_server._validate_tts_request(req)
+        assert "does not support CustomVoice" in result
 
     def test_build_tts_params(self, speech_server):
         """Test TTS parameter building."""
@@ -404,6 +627,121 @@ class TestTTSMethods:
         # Verify speakers are normalized to lowercase
         assert server.supported_speakers == {"ryan", "vivian", "aiden"}
 
+    def test_build_tts_params_with_uploaded_voice(self, speech_server):
+        """Test _build_tts_params auto-sets ref_audio for uploaded voices."""
+        # Mock an uploaded speaker
+        speech_server.uploaded_speakers = {
+            "custom_voice": {
+                "name": "custom_voice",
+                "file_path": "/tmp/voice_samples/custom_voice_consent_123.wav",
+                "mime_type": "audio/wav",
+            }
+        }
+        speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"}
+
+        # Mock _get_uploaded_audio_data to return base64 data
+        with patch.object(speech_server, "_get_uploaded_audio_data") as mock_get_audio:
+            mock_get_audio.return_value = "data:audio/wav;base64,ZmFrZWF1ZGlv"
+
+            req = OpenAICreateSpeechRequest(input="Hello", voice="custom_voice", task_type="Base")
+
+            params = speech_server._build_tts_params(req)
+
+            # Verify ref_audio was auto-set
+            assert "ref_audio" in params
+            assert params["ref_audio"] == ["data:audio/wav;base64,ZmFrZWF1ZGlv"]
+            assert "x_vector_only_mode" in params
+            assert params["x_vector_only_mode"] == [True]
+            mock_get_audio.assert_called_once_with("custom_voice")
+
+    def test_build_tts_params_without_uploaded_voice(self, speech_server):
+        """Test _build_tts_params does not auto-set ref_audio for non-uploaded voices."""
+        # No uploaded speakers
+        speech_server.uploaded_speakers = {}
+        speech_server.supported_speakers = {"ryan", "vivian"}
+
+        req = OpenAICreateSpeechRequest(input="Hello", voice="ryan", task_type="Base")
+
+        params = speech_server._build_tts_params(req)
+
+        # Verify ref_audio was NOT auto-set
+        assert "ref_audio" not in params
+        assert "x_vector_only_mode" not in params
+
+    def test_build_tts_params_with_explicit_ref_audio(self, speech_server):
+        """Test _build_tts_params uses explicit ref_audio even for uploaded voices."""
+        # Mock an uploaded speaker
+        speech_server.uploaded_speakers = {
+            "custom_voice": {
+                "name": "custom_voice",
+                "file_path": "/tmp/voice_samples/custom_voice_consent_123.wav",
+                "mime_type": "audio/wav",
+            }
+        }
+        speech_server.supported_speakers = {"ryan", "vivian", "custom_voice"}
+
+        req = OpenAICreateSpeechRequest(
+            input="Hello", voice="custom_voice", task_type="Base", ref_audio="data:audio/wav;base64,ZXhwbGljaXQ="
+        )
+
+        params = speech_server._build_tts_params(req)
+
+        # _build_tts_params should NOT auto-set ref_audio when explicit ref_audio
+        # is provided (request.ref_audio is not None skips the auto-set branch).
+        # The explicit ref_audio is resolved later in create_speech() via
+        # _resolve_ref_audio(), not in _build_tts_params().
+        assert "ref_audio" not in params
+        # x_vector_only_mode should not be set when explicit ref_audio is provided
+        assert "x_vector_only_mode" not in params
+
+    def test_get_uploaded_audio_data(self, speech_server):
+        """Test _get_uploaded_audio_data function."""
+        # Mock file operations
+        with (
+            patch("builtins.open", create=True) as mock_open,
+            patch("base64.b64encode") as mock_b64encode,
+            patch("pathlib.Path.exists") as mock_exists,
+        ):
+            mock_exists.return_value = True
+            mock_b64encode.return_value = b"ZmFrZWF1ZGlv"
+
+            # Setup mock file
+            mock_file = MagicMock()
+            mock_file.read.return_value = b"fakeaudio"
+            mock_open.return_value.__enter__.return_value = mock_file
+
+            # Setup uploaded speaker
+            speech_server.uploaded_speakers = {
+                "test_voice": {"name": "test_voice", "file_path": "/tmp/test.wav", "mime_type": "audio/wav"}
+            }
+            result = speech_server._get_uploaded_audio_data("test_voice")
+
+            assert result == "data:audio/wav;base64,ZmFrZWF1ZGlv"
+            mock_open.assert_called_once_with(Path("/tmp/test.wav"), "rb")
+            mock_b64encode.assert_called_once_with(b"fakeaudio")
+
+    def test_get_uploaded_audio_data_missing_file(self, speech_server):
+        """Test _get_uploaded_audio_data when file is missing."""
+        with patch("pathlib.Path.exists") as mock_exists:
+            mock_exists.return_value = False
+
+            # Setup uploaded speaker
+            speech_server.uploaded_speakers = {
+                "test_voice": {"name": "test_voice", "file_path": "/tmp/test.wav", "mime_type": "audio/wav"}
+            }
+
+            result = speech_server._get_uploaded_audio_data("test_voice")
+
+            assert result is None
+
+    def test_get_uploaded_audio_data_voice_not_found(self, speech_server):
+        """Test _get_uploaded_audio_data when voice is not in uploaded_speakers."""
+        speech_server.uploaded_speakers = {}
+
+        result = speech_server._get_uploaded_audio_data("nonexistent")
+
+        assert result is None
+
     def test_max_instructions_length_default(self, speech_server):
         """Test default max instructions length (500) when no config provided."""
         # Fixture creates server with no CLI override and no TTS stage
@@ -510,6 +848,70 @@ class TestTTSMethods:
         assert "max 10 characters" in error
 
 
+class TestFileValidationFunctions:
+    """Unit tests for file validation helper functions."""
+
+    def test_sanitize_filename(self):
+        """Test _sanitize_filename function."""
+        from vllm_omni.entrypoints.openai.serving_speech import _sanitize_filename
+
+        # Test normal filenames
+        assert _sanitize_filename("test.wav") == "test.wav"
+        assert _sanitize_filename("test-file.mp3") == "test-file.mp3"
+        assert _sanitize_filename("test_file.flac") == "test_file.flac"
+
+        # Test path traversal attempts
+        assert _sanitize_filename("../../../etc/passwd") == "passwd"
+        assert _sanitize_filename("/absolute/path/file.wav") == "file.wav"
+
+        # Test special characters
+        assert _sanitize_filename("file with spaces.wav") == "file_with_spaces.wav"
+        assert _sanitize_filename("file&with&special&chars.wav") == "file_with_special_chars.wav"
+        assert _sanitize_filename("file@with#special$chars%.wav") == "file_with_special_chars_.wav"
+
+        # Test empty filename
+        assert _sanitize_filename("") == "file"
+
+        # Test very long filename
+        long_name = "a" * 300
+        sanitized = _sanitize_filename(long_name)
+        assert len(sanitized) == 255
+        assert sanitized.startswith("a")
+
+    def test_validate_path_within_directory(self, tmp_path):
+        """Test _validate_path_within_directory function."""
+        from vllm_omni.entrypoints.openai.serving_speech import _validate_path_within_directory
+
+        # Create test directory structure
+        base_dir = tmp_path / "uploads"
+        base_dir.mkdir()
+
+        # Valid paths within directory
+        valid_file = base_dir / "test.wav"
+        valid_subdir_file = base_dir / "subdir" / "test.wav"
+        valid_subdir_file.parent.mkdir()
+
+        assert _validate_path_within_directory(valid_file, base_dir) is True
+        assert _validate_path_within_directory(valid_subdir_file, base_dir) is True
+
+        # Invalid paths outside directory
+        outside_file = tmp_path / "outside.wav"
+        assert _validate_path_within_directory(outside_file, base_dir) is False
+
+        # Test with symlink (should fail)
+        if hasattr(os, "symlink"):
+            link_target = tmp_path / "target.wav"
+            link_target.touch()
+            symlink = base_dir / "link.wav"
+            os.symlink(link_target, symlink)
+            # Symlinks to outside should be rejected
+            assert _validate_path_within_directory(symlink, base_dir) is False
+
+        # Test with non-existent file (should still validate path)
+        non_existent = base_dir / "nonexistent.wav"
+        assert _validate_path_within_directory(non_existent, base_dir) is True
+
+
 class TestStreamingProtocolValidation:
     """Unit tests for the stream field validators in OpenAICreateSpeechRequest."""
 

tests/entrypoints/openai_api/test_video_server.py

@@ -5,6 +5,7 @@ Unit tests for OpenAI-compatible video generation endpoints.
 """
 
 import io
+from types import SimpleNamespace
 
 import pytest
 from fastapi import FastAPI
@@ -21,8 +22,12 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class MockVideoResult:
-    def __init__(self, videos):
+    def __init__(self, videos, audios=None, sample_rate=None):
         self.multimodal_output = {"video": videos}
+        if audios is not None:
+            self.multimodal_output["audio"] = audios
+        if sample_rate is not None:
+            self.multimodal_output["audio_sample_rate"] = sample_rate
 
 
 class FakeAsyncOmni:
@@ -91,6 +96,7 @@ def test_t2v_video_generation_form(test_client, mocker: MockerFixture):
     assert captured.height == 360
     assert captured.num_frames == 24
     assert captured.fps == 12
+    assert captured.frame_rate == 12.0
     assert fps_values == [12, 12]
 
 
@@ -121,6 +127,32 @@ def test_i2v_video_generation_form(test_client, mocker: MockerFixture):
     assert input_image.size == (48, 32)
 
 
+def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client, mocker: MockerFixture):
+    image_bytes = _make_test_image_bytes((48, 32))
+
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        return_value="Zg==",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "A bear playing with yarn.",
+            "width": "96",
+            "height": "64",
+        },
+        files={"input_reference": ("input.png", image_bytes, "image/png")},
+    )
+
+    assert response.status_code == 200
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    prompt = engine.captured_prompt
+    input_image = prompt["multi_modal_data"]["image"]
+    assert isinstance(input_image, Image.Image)
+    assert input_image.size == (96, 64)
+
+
 def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
     fps_values = []
 
@@ -168,6 +200,42 @@ def test_size_param_sets_width_height(test_client, mocker: MockerFixture):
     assert captured.height == 240
 
 
+def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFixture):
+    audio_sample_rates = []
+
+    def _fake_encode(video, fps, audio=None, audio_sample_rate=None):
+        audio_sample_rates.append(audio_sample_rate)
+        return "Zg=="
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    engine.model_config = SimpleNamespace(
+        hf_config=SimpleNamespace(
+            vocoder=SimpleNamespace(
+                config=SimpleNamespace(output_sampling_rate=16000),
+            ),
+        ),
+    )
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        engine.captured_prompt = prompt
+        engine.captured_sampling_params_list = sampling_params_list
+        yield MockVideoResult([object()], audios=[object()])
+
+    engine.generate = _generate
+
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        side_effect=_fake_encode,
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={"prompt": "video with audio"},
+    )
+
+    assert response.status_code == 200
+    assert audio_sample_rates == [16000]
+
+
 def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
     mocker.patch(
         "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -11,12 +11,18 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 MODEL = "riverclouds/qwen_image_random"
 
 
+def _noop_inline_engine(self, model, stage_config, kwargs):
+    self._inline_diffusion = False
+    self._inline_engine = None
+
+
 def test_default_stage_config_includes_cache_backend(monkeypatch):
     """Ensure cache_backend/cache_config are preserved in default diffusion stage."""
     monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
     monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
+    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
 
     omni = AsyncOmni(
         model=MODEL,
@@ -47,6 +53,7 @@ def test_default_cache_config_used_when_missing(monkeypatch):
     monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
+    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
 
     omni = AsyncOmni(
         model=MODEL,
@@ -65,6 +72,7 @@ def test_default_stage_devices_from_sequence_parallel(monkeypatch):
     monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
+    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
 
     omni = AsyncOmni(
         model=MODEL,

tests/entrypoints/test_cfg_companion_tracker.py

@@ -0,0 +1,114 @@
+import time
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.entrypoints.cfg_companion_tracker import CfgCompanionTracker
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def dummy_expand_func(prompt, sp0):
+    if prompt == "expand_me":
+        return [SimpleNamespace(prompt={"prompt": "neg"}, role="cfg_text", request_id_suffix="__cfg_text")]
+    return []
+
+
+@pytest.fixture
+def tracker():
+    sp0 = SimpleNamespace()
+    return CfgCompanionTracker(prompt_expand_func=dummy_expand_func, stage0_sampling_params=sp0, timeout_s=0.1)
+
+
+def test_companion_tracker_initialization(tracker):
+    assert not tracker.is_active
+    assert tracker.num_companions == 0
+
+
+def test_expand_prompts_registers_companions(tracker):
+    request_id_to_prompt = {"req1": "expand_me", "req2": "do_not_expand"}
+
+    pairs = tracker.expand_prompts(request_id_to_prompt)
+
+    assert len(pairs) == 1
+    companion_id, prompt = pairs[0]
+    assert companion_id == "req1__cfg_text"
+    assert prompt == {"prompt": "neg"}
+
+    assert tracker.is_active
+    assert tracker.num_companions == 1
+    assert tracker.is_companion("req1__cfg_text")
+    assert not tracker.is_companion("req2__cfg_text")
+    assert tracker.has_companions("req1")
+    assert not tracker.has_companions("req2")
+
+    comp_map = tracker.get_companion_request_ids("req1")
+    assert comp_map == {"cfg_text": "req1__cfg_text"}
+
+
+def test_companion_lifecycle_success(tracker):
+    request_id_to_prompt = {"req1": "expand_me"}
+    tracker.expand_prompts(request_id_to_prompt)
+
+    # Defer parent
+    engine_outputs = {"out": 123}
+    tracker.defer_parent("req1", engine_outputs, stage_id=0)
+
+    # Initially not done
+    assert not tracker.all_companions_done("req1")
+
+    # Companion completes
+    parent_id = tracker.on_companion_completed("req1__cfg_text")
+
+    # Parent should be returned since all companions are done and it is pending
+    assert parent_id == "req1"
+    assert tracker.all_companions_done("req1")
+
+    # Pop pending parent
+    popped = tracker.pop_pending_parent("req1")
+    assert popped is not None
+    assert popped["engine_outputs"] == engine_outputs
+    assert popped["stage_id"] == 0
+
+
+def test_companion_lifecycle_failure(tracker):
+    request_id_to_prompt = {"req1": "expand_me"}
+    tracker.expand_prompts(request_id_to_prompt)
+
+    tracker.defer_parent("req1", {"out": 123}, stage_id=0)
+
+    # Companion fails
+    parent_id, aborted = tracker.on_companion_error("req1__cfg_text")
+
+    assert parent_id == "req1"
+    assert aborted is True
+    assert tracker.is_parent_failed("req1")
+
+    # Parent should be removed from pending list
+    assert tracker.pop_pending_parent("req1") is None
+
+    # Consume failure
+    tracker.consume_parent_failure("req1")
+    assert not tracker.is_parent_failed("req1")
+
+
+def test_companion_lifecycle_timeout(tracker):
+    request_id_to_prompt = {"req1": "expand_me"}
+    tracker.expand_prompts(request_id_to_prompt)
+
+    tracker.defer_parent("req1", {"out": 123}, stage_id=0)
+
+    # Initially no timeouts
+    timeouts = tracker.check_timeouts()
+    assert len(timeouts) == 0
+
+    # Wait for timeout
+    time.sleep(0.15)
+
+    # Check timeouts again
+    timeouts = tracker.check_timeouts()
+    assert len(timeouts) == 1
+    assert timeouts[0] == "req1"
+
+    # Should be removed from pending
+    assert tracker.pop_pending_parent("req1") is None

tests/entrypoints/test_omni_diffusion.py

@@ -4,6 +4,7 @@ from queue import Empty, Queue
 from typing import Any
 
 import pytest
+import torch
 from pytest_mock import MockerFixture
 
 from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
@@ -1340,3 +1341,91 @@ def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: Mo
 
     # Verify stop_stage_worker was called (process should be set)
     assert omni.stage_list[0]._proc is not None
+
+
+# Tests below are for diffusion dtype normalization fix from the following:
+# https://github.com/vllm-project/vllm-omni/pull/1391
+# In the future we should ensure dtypes are parsed in less hacky way and make
+# these tests more atomic.
+@pytest.mark.parametrize("dtype", ["float16", torch.float16])
+def test_dtype_normalization_valid_types(
+    monkeypatch, dtype: str | torch.dtype, mocker: MockerFixture, fake_stage_config
+):
+    """Ensure Diffusion Config builder coerces valid types correctly."""
+
+    def _fake_loader(model: str, base_engine_args=None):
+        # Return not stage configs to fall back to the diffusion cfg builder
+        return []
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+
+    omni = Omni(model="any", init_timeout=1, dtype=dtype)
+
+    # Dtype parsing being checked is on the diffusion path
+    assert len(omni.stage_configs) == 1
+    assert omni.stage_configs[0]["stage_type"] == "diffusion"
+
+    # Regardless of whether a str / dtype is passed, it resolves correctly
+    engine_args = omni.stage_configs[0]["engine_args"]
+    assert "dtype" in engine_args
+    assert isinstance(engine_args["dtype"], str)
+    assert engine_args["dtype"] == "float16"
+
+
+def test_dtype_normalization_invalid_types(monkeypatch, mocker: MockerFixture, fake_stage_config):
+    """Ensure Diffusion Config builder correctly handles bad dtype overrides."""
+
+    def _fake_loader(model: str, base_engine_args=None):
+        # Return not stage configs to fall back to the diffusion cfg builder
+        return []
+
+    class NotATorchDtype:
+        pass
+
+    import sys
+
+    for module_name in [
+        "vllm_omni.entrypoints.utils",
+        "vllm_omni.entrypoints.omni",
+        "vllm_omni.entrypoints.omni_stage",
+    ]:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
+    _setup_ipc_mocks(monkeypatch)
+    _setup_log_mocks(monkeypatch)
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        _fake_loader,
+        raising=False,
+    )
+
+    # Raise TypeError if we get an unrecognized type
+    with pytest.raises(TypeError):
+        Omni(model="any", init_timeout=1, dtype=NotATorchDtype)

tests/entrypoints/test_omni_llm.py

@@ -1,3 +1,4 @@
+import inspect
 import uuid
 import warnings
 from queue import Empty, Queue
@@ -1194,3 +1195,72 @@ def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: Mo
 
     # Verify stop_stage_worker was called (process should be set)
     assert omni.stage_list[0]._proc is not None
+
+
+# ---------------------------------------------------------------------------
+# Signature compatibility tests — catch upstream API drift early
+# ---------------------------------------------------------------------------
+
+_INHERITANCE_PAIRS: list[tuple[str, str, str, str]] = [
+    (
+        "vllm.entrypoints.llm",
+        "LLM",
+        "vllm_omni.entrypoints.omni_llm",
+        "OmniLLM",
+    ),
+]
+
+
+def _import_class(module_path: str, class_name: str):
+    mod = __import__(module_path, fromlist=[class_name])
+    return getattr(mod, class_name)
+
+
+@pytest.mark.parametrize(
+    "up_mod,up_cls,omni_mod,omni_cls",
+    _INHERITANCE_PAIRS,
+    ids=[f"{pair[3]}({pair[1]})" for pair in _INHERITANCE_PAIRS],
+)
+def test_overridden_method_signatures_compatible(up_mod: str, up_cls: str, omni_mod: str, omni_cls: str):
+    """All params accepted by a base-class method must also be accepted by
+    the overriding method in the vllm-omni subclass.  This prevents
+    TypeError at runtime when upstream callers pass new arguments."""
+    BaseCls = _import_class(up_mod, up_cls)
+    OmniCls = _import_class(omni_mod, omni_cls)
+
+    failures: list[str] = []
+    for name in dir(OmniCls):
+        if name.startswith("__") and name != "__init__":
+            continue
+        omni_method = getattr(OmniCls, name, None)
+        base_method = getattr(BaseCls, name, None)
+        if not (callable(omni_method) and callable(base_method)):
+            continue
+        if omni_method is base_method:
+            continue
+        try:
+            base_sig = inspect.signature(base_method)
+            omni_sig = inspect.signature(omni_method)
+        except (ValueError, TypeError):
+            continue
+
+        omni_has_var_keyword = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in omni_sig.parameters.values())
+
+        base_params = base_sig.parameters
+        omni_params = set(omni_sig.parameters.keys())
+
+        missing = []
+        for pname, param in base_params.items():
+            if pname in omni_params:
+                continue
+            if omni_has_var_keyword and param.kind in (
+                inspect.Parameter.KEYWORD_ONLY,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            ):
+                continue
+            missing.append(pname)
+
+        if missing:
+            failures.append(f"{omni_cls}.{name}() missing params {sorted(missing)}; base={base_sig}, omni={omni_sig}")
+
+    assert not failures, f"Signature mismatches found ({len(failures)}):\n" + "\n".join(f"  - {f}" for f in failures)

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py

@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for CosyVoice3 components."""
+
+import pytest
+import torch
+import torch.nn as nn
+
+
+class TestPreLookaheadLayer:
+    """Tests for PreLookaheadLayer."""
+
+    @pytest.fixture
+    def layer(self):
+        from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer
+
+        return PreLookaheadLayer(in_channels=512, channels=512, pre_lookahead_len=3)
+
+    def test_forward_shape(self, layer):
+        """Test that output shape matches input shape."""
+        batch, seq_len, channels = 2, 10, 512
+        x = torch.randn(batch, seq_len, channels)
+
+        out = layer(x)
+
+        assert out.shape == x.shape
+
+    def test_forward_with_context(self, layer):
+        """Test forward with context for streaming."""
+        batch, seq_len, channels = 1, 10, 512
+        x = torch.randn(batch, seq_len, channels)
+        context = torch.randn(batch, 3, channels)  # pre_lookahead_len=3
+
+        layer.eval()
+        out = layer(x, context=context)
+
+        assert out.shape == x.shape
+
+    def test_residual_connection(self, layer):
+        """Test that residual connection is applied."""
+        batch, seq_len, channels = 1, 5, 512
+        x = torch.zeros(batch, seq_len, channels)
+
+        # With zero input, output should also be close to zero due to residual
+        out = layer(x)
+
+        # Output should be close to input (residual) plus conv output
+        assert out.shape == x.shape
+
+
+class TestDiTAttention:
+    """Tests for DiTAttention with diffusion backend."""
+
+    @pytest.fixture
+    def attention(self):
+        from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTAttention
+
+        return DiTAttention(dim=512, heads=8, dim_head=64, dropout=0.0)
+
+    def test_forward_shape(self, attention):
+        """Test attention output shape."""
+        batch, seq_len, dim = 2, 16, 512
+        x = torch.randn(batch, seq_len, dim)
+
+        out = attention(x)
+
+        assert out.shape == x.shape
+
+    def test_forward_with_mask(self, attention):
+        """Test attention with mask."""
+        batch, seq_len, dim = 2, 16, 512
+        x = torch.randn(batch, seq_len, dim)
+        mask = torch.ones(batch, seq_len, dtype=torch.bool)
+        mask[:, -3:] = False  # Mask last 3 positions
+
+        out = attention(x, mask=mask)
+
+        assert out.shape == x.shape
+        # Masked positions should be zero
+        assert torch.allclose(out[:, -3:], torch.zeros_like(out[:, -3:]))
+
+    def test_qkv_projections(self, attention):
+        """Test that Q/K/V projections exist and have correct dimensions."""
+        assert hasattr(attention, "to_q")
+        assert hasattr(attention, "to_k")
+        assert hasattr(attention, "to_v")
+        assert attention.to_q.out_features == 512  # heads * dim_head
+        assert attention.to_k.out_features == 512
+        assert attention.to_v.out_features == 512
+
+
+class TestDiTBlock:
+    """Tests for DiTBlock."""
+
+    @pytest.fixture
+    def block(self):
+        from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTBlock
+
+        return DiTBlock(dim=512, heads=8, dim_head=64, ff_mult=4, dropout=0.0)
+
+    def test_forward_shape(self, block):
+        """Test block output shape."""
+        batch, seq_len, dim = 2, 16, 512
+        x = torch.randn(batch, seq_len, dim)
+        t = torch.randn(batch, dim)  # Timestep embedding
+
+        out = block(x, t)
+
+        assert out.shape == x.shape
+
+    def test_adalayernorm_modulation(self, block):
+        """Test that AdaLayerNorm modulates based on timestep."""
+        batch, seq_len, dim = 1, 8, 512
+        x = torch.randn(batch, seq_len, dim)
+        t1 = torch.zeros(batch, dim)
+        t2 = torch.ones(batch, dim)
+
+        out1 = block(x, t1)
+        out2 = block(x, t2)
+
+        # Different timesteps should produce different outputs
+        assert not torch.allclose(out1, out2)
+
+
+class TestDiT:
+    """Tests for the full DiT model."""
+
+    @pytest.fixture
+    def dit(self):
+        from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT
+
+        return DiT(
+            dim=256,
+            depth=2,
+            heads=4,
+            dim_head=64,
+            dropout=0.0,
+            ff_mult=2,
+            mel_dim=80,
+            mu_dim=80,
+            spk_dim=80,
+            long_skip_connection=True,
+        )
+
+    def test_forward_shape(self, dit):
+        """Test DiT forward output shape."""
+        batch, mel_dim, seq_len = 1, 80, 32
+        x = torch.randn(batch, mel_dim, seq_len)
+        mask = torch.ones(batch, 1, seq_len)
+        mu = torch.randn(batch, mel_dim, seq_len)
+        t = torch.tensor([0.5])
+        spks = torch.randn(batch, 80)
+        cond = torch.randn(batch, mel_dim, seq_len)
+
+        out = dit(x, mask, mu, t, spks=spks, cond=cond)
+
+        assert out.shape == (batch, mel_dim, seq_len)
+
+    def test_timestep_embedding(self, dit):
+        """Test that different timesteps produce different outputs."""
+        batch, mel_dim, seq_len = 1, 80, 16
+        x = torch.randn(batch, mel_dim, seq_len)
+        mask = torch.ones(batch, 1, seq_len)
+        mu = torch.randn(batch, mel_dim, seq_len)
+        spks = torch.randn(batch, 80)
+        cond = torch.randn(batch, mel_dim, seq_len)
+
+        out1 = dit(x, mask, mu, torch.tensor([0.0]), spks=spks, cond=cond)
+        out2 = dit(x, mask, mu, torch.tensor([1.0]), spks=spks, cond=cond)
+
+        assert not torch.allclose(out1, out2)
+
+
+class TestCFM:
+    """Tests for Conditional Flow Matching classes."""
+
+    @pytest.fixture
+    def dummy_estimator(self):
+        """Create a dummy estimator for testing."""
+
+        class DummyEstimator(nn.Module):
+            def __init__(self, mel_dim=80):
+                super().__init__()
+                self.mel_dim = mel_dim
+
+            def forward(self, x, mask, mu, t, spks=None, cond=None):
+                return torch.zeros_like(x)
+
+        return DummyEstimator()
+
+    def test_causal_conditional_cfm_forward(self, dummy_estimator):
+        """Test CausalConditionalCFM forward pass."""
+        from omegaconf import DictConfig
+
+        from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import CausalConditionalCFM
+
+        cfm_params = DictConfig(
+            {
+                "sigma_min": 1e-6,
+                "solver": "euler",
+                "t_scheduler": "cosine",
+                "training_cfg_rate": 0.2,
+                "inference_cfg_rate": 0.7,
+            }
+        )
+
+        cfm = CausalConditionalCFM(
+            in_channels=80,
+            cfm_params=cfm_params,
+            n_spks=1,
+            spk_emb_dim=80,
+            estimator=dummy_estimator,
+        )
+
+        batch, mel_dim, seq_len = 1, 80, 32
+        mu = torch.randn(batch, mel_dim, seq_len)
+        mask = torch.ones(batch, 1, seq_len)
+        spks = torch.randn(batch, 80)
+        cond = torch.randn(batch, mel_dim, seq_len)
+
+        out, _ = cfm(mu, mask, n_timesteps=2, spks=spks, cond=cond)
+
+        assert out.shape == mu.shape
+
+
+class TestSDPAFallback:
+    """Test SDPA fallback for float32 inputs."""
+
+    def test_float32_uses_sdpa(self):
+        """Test that float32 inputs use SDPA fallback."""
+        from vllm_omni.diffusion.attention.layer import Attention
+
+        attn = Attention(
+            num_heads=8,
+            head_size=64,
+            causal=False,
+            softmax_scale=1.0 / 8.0,
+        )
+
+        batch, seq_len, heads, dim = 1, 16, 8, 64
+        q = torch.randn(batch, seq_len, heads, dim, dtype=torch.float32)
+        k = torch.randn(batch, seq_len, heads, dim, dtype=torch.float32)
+        v = torch.randn(batch, seq_len, heads, dim, dtype=torch.float32)
+
+        # Should not raise error - SDPA fallback handles float32
+        out = attn(q, k, v)
+
+        assert out.shape == (batch, seq_len, heads, dim)
+        assert out.dtype == torch.float32

tests/model_executor/models/cosyvoice3/test_cosyvoice3_utils.py

@@ -0,0 +1,230 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for CosyVoice3 utility functions."""
+
+import pytest
+import torch
+
+
+class TestMakePadMask:
+    """Tests for make_pad_mask utility."""
+
+    def test_basic_mask_creation(self):
+        """Test basic padding mask creation."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask
+
+        lengths = torch.tensor([5, 3, 2])
+        mask = make_pad_mask(lengths)
+
+        # Expected: [[F,F,F,F,F], [F,F,F,T,T], [F,F,T,T,T]]
+        expected = torch.tensor(
+            [
+                [False, False, False, False, False],
+                [False, False, False, True, True],
+                [False, False, True, True, True],
+            ]
+        )
+        assert torch.equal(mask, expected)
+
+    def test_single_sequence(self):
+        """Test mask for single sequence."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask
+
+        lengths = torch.tensor([3])
+        mask = make_pad_mask(lengths)
+
+        expected = torch.tensor([[False, False, False]])
+        assert torch.equal(mask, expected)
+
+    def test_explicit_max_len(self):
+        """Test with explicit max_len parameter."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask
+
+        lengths = torch.tensor([2, 3])
+        mask = make_pad_mask(lengths, max_len=5)
+
+        expected = torch.tensor(
+            [
+                [False, False, True, True, True],
+                [False, False, False, True, True],
+            ]
+        )
+        assert torch.equal(mask, expected)
+
+    def test_zero_max_len_uses_max_length(self):
+        """Test that max_len=0 falls back to max(lengths)."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask
+
+        lengths = torch.tensor([2, 4])
+        mask = make_pad_mask(lengths, max_len=0)
+
+        # Should use max(lengths) = 4
+        assert mask.shape == (2, 4)
+
+    def test_full_length_sequence(self):
+        """Test sequence that fills the entire length (no padding)."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask
+
+        lengths = torch.tensor([5])
+        mask = make_pad_mask(lengths, max_len=5)
+
+        # No positions should be masked
+        assert not mask.any()
+
+    def test_empty_sequence(self):
+        """Test sequence with zero length (all padding)."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask
+
+        lengths = torch.tensor([0])
+        mask = make_pad_mask(lengths, max_len=3)
+
+        # All positions should be masked
+        assert mask.all()
+
+    def test_device_preservation(self):
+        """Test that mask is created on the same device as input."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask
+
+        lengths = torch.tensor([3, 2])
+        mask = make_pad_mask(lengths)
+
+        assert mask.device == lengths.device
+
+
+class TestDynamicRangeCompression:
+    """Tests for dynamic_range_compression_torch utility."""
+
+    def test_basic_compression(self):
+        """Test basic compression with default parameters."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import dynamic_range_compression_torch
+
+        x = torch.tensor([1.0, 10.0, 100.0])
+        result = dynamic_range_compression_torch(x)
+
+        # Result should be log(x * 1) = log(x)
+        expected = torch.log(x)
+        assert torch.allclose(result, expected)
+
+    def test_clipping_small_values(self):
+        """Test that small values are clipped."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import dynamic_range_compression_torch
+
+        x = torch.tensor([0.0, 1e-10, 1e-6])
+        clip_val = 1e-5
+        result = dynamic_range_compression_torch(x, clip_val=clip_val)
+
+        # All values below clip_val should be clipped
+        expected = torch.log(torch.tensor([clip_val, clip_val, clip_val]))
+        assert torch.allclose(result, expected)
+
+    def test_scaling_factor(self):
+        """Test with different scaling factor c."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import dynamic_range_compression_torch
+
+        x = torch.tensor([1.0, 10.0])
+        c = 2.0
+        result = dynamic_range_compression_torch(x, c=c)
+
+        expected = torch.log(x * c)
+        assert torch.allclose(result, expected)
+
+
+class TestSpectralNormalize:
+    """Tests for spectral_normalize_torch utility."""
+
+    def test_is_wrapper_for_compression(self):
+        """Test that spectral_normalize is a wrapper for dynamic_range_compression."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import (
+            dynamic_range_compression_torch,
+            spectral_normalize_torch,
+        )
+
+        magnitudes = torch.tensor([1.0, 10.0, 100.0])
+        result = spectral_normalize_torch(magnitudes)
+        expected = dynamic_range_compression_torch(magnitudes)
+
+        assert torch.allclose(result, expected)
+
+
+class TestExactDiv:
+    """Tests for exact_div utility."""
+
+    def test_exact_division(self):
+        """Test exact division with valid inputs."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import exact_div
+
+        assert exact_div(10, 2) == 5
+        assert exact_div(100, 25) == 4
+        assert exact_div(0, 5) == 0
+
+    def test_non_exact_division_raises(self):
+        """Test that non-exact division raises AssertionError."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import exact_div
+
+        with pytest.raises(AssertionError):
+            exact_div(10, 3)
+
+        with pytest.raises(AssertionError):
+            exact_div(7, 2)
+
+
+class TestConcatTextWithPromptIds:
+    """Tests for concat_text_with_prompt_ids utility."""
+
+    def test_basic_concat(self):
+        """Test basic concatenation."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids
+
+        text = torch.tensor([[5, 6, 7]], dtype=torch.int64)
+        text_len = torch.tensor([3], dtype=torch.int32)
+        prompt_text = torch.tensor([[1, 2]], dtype=torch.int64)
+        prompt_text_len = torch.tensor([2], dtype=torch.int32)
+
+        result, result_len = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len)
+
+        assert result.tolist() == [[1, 2, 5, 6, 7]]
+        assert result_len.item() == 5
+
+    def test_empty_prompt(self):
+        """Test with empty prompt."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids
+
+        text = torch.tensor([[1, 2, 3]], dtype=torch.int64)
+        text_len = torch.tensor([3], dtype=torch.int32)
+        prompt_text = torch.tensor([[]], dtype=torch.int64).reshape(1, 0)
+        prompt_text_len = torch.tensor([0], dtype=torch.int32)
+
+        result, result_len = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len)
+
+        assert result.tolist() == [[1, 2, 3]]
+        assert result_len.item() == 3
+
+    def test_batch_concat(self):
+        """Test batched concatenation."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids
+
+        text = torch.tensor([[5, 6], [7, 8]], dtype=torch.int64)
+        text_len = torch.tensor([2, 2], dtype=torch.int32)
+        prompt_text = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
+        prompt_text_len = torch.tensor([2, 2], dtype=torch.int32)
+
+        result, result_len = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len)
+
+        assert result.tolist() == [[1, 2, 5, 6], [3, 4, 7, 8]]
+        assert result_len.tolist() == [4, 4]
+
+    def test_prompt_comes_first(self):
+        """Test that prompt tokens come before text tokens."""
+        from vllm_omni.model_executor.models.cosyvoice3.utils import concat_text_with_prompt_ids
+
+        text = torch.tensor([[100, 200]], dtype=torch.int64)
+        text_len = torch.tensor([2], dtype=torch.int32)
+        prompt_text = torch.tensor([[1, 2, 3]], dtype=torch.int64)
+        prompt_text_len = torch.tensor([3], dtype=torch.int32)
+
+        result, _ = concat_text_with_prompt_ids(text, text_len, prompt_text, prompt_text_len)
+
+        # First 3 tokens should be prompt
+        assert result[0, :3].tolist() == [1, 2, 3]
+        # Last 2 tokens should be text
+        assert result[0, 3:].tolist() == [100, 200]

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -4,27 +4,60 @@
 from collections import defaultdict
 from types import SimpleNamespace
 
+import pytest
 import torch
 
 from vllm_omni.model_executor.stage_input_processors.qwen3_tts import talker2code2wav_async_chunk
 
+_FRAME = [1, 2, 3, 4]  # 4-codebook frame
+_Q = len(_FRAME)  # num quantizers
 
-def _req(external_req_id: str, *, finished: bool):
+
+def _req(rid: str, *, finished: bool, initial_codec_chunk_frames: int | None = None):
+    ai = None
+    if initial_codec_chunk_frames is not None:
+        entry = SimpleNamespace(list_data=[initial_codec_chunk_frames])
+        ai = SimpleNamespace(entries={"initial_codec_chunk_frames": entry})
     return SimpleNamespace(
-        external_req_id=external_req_id,
+        external_req_id=rid,
         is_finished=lambda: finished,
+        additional_information=ai,
     )
 
 
-def test_talker2code2wav_async_chunk_does_not_emit_empty_chunk_when_not_finished():
-    transfer_manager = SimpleNamespace(
+def _tm(*, chunk_frames=25, left_context=25, initial_chunk=0):
+    return SimpleNamespace(
         code_prompt_token_ids=defaultdict(list),
-        connector=SimpleNamespace(config={"extra": {"codec_chunk_frames": 25, "codec_left_context_frames": 25}}),
+        put_req_chunk=defaultdict(int),
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "codec_chunk_frames": chunk_frames,
+                    "codec_left_context_frames": left_context,
+                    "initial_codec_chunk_frames": initial_chunk,
+                }
+            }
+        ),
     )
+
+
+def _call(tm, rid, *, n_frames, put_req=0, finished=False, req_ic=None):
+    tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)]
+    tm.put_req_chunk[rid] = put_req
+    return talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        pooling_output={"audio_codes": torch.zeros((0,))},
+        request=_req(rid, finished=finished, initial_codec_chunk_frames=req_ic),
+        is_finished=finished,
+    )
+
+
+def test_does_not_emit_empty_chunk_when_not_finished():
+    tm = _tm()
     request = _req("rid-empty", finished=False)
 
     payload = talker2code2wav_async_chunk(
-        transfer_manager=transfer_manager,
+        transfer_manager=tm,
         pooling_output={"audio_codes": torch.zeros((0,))},
         request=request,
     )
@@ -32,36 +65,29 @@ def test_talker2code2wav_async_chunk_does_not_emit_empty_chunk_when_not_finished
     assert payload is None
 
 
-def test_talker2code2wav_async_chunk_flushes_tail_when_finished_without_pooler_output():
-    transfer_manager = SimpleNamespace(
-        code_prompt_token_ids=defaultdict(list),
-        connector=SimpleNamespace(config={"extra": {"codec_chunk_frames": 25, "codec_left_context_frames": 25}}),
-    )
-    request_id = "rid-tail"
-    transfer_manager.code_prompt_token_ids[request_id] = [[1, 2, 3, 4] for _ in range(24)]
-    request = _req(request_id, finished=True)
+def test_flushes_tail_when_finished_without_pooler_output():
+    tm = _tm()
+    rid = "rid-tail"
+    tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(24)]
+    request = _req(rid, finished=True)
 
     payload = talker2code2wav_async_chunk(
-        transfer_manager=transfer_manager,
-        pooling_output=None,  # e.g. EOS step with no audio_codes
+        transfer_manager=tm,
+        pooling_output=None,
         request=request,
     )
 
     assert payload is not None
     assert payload["finished"].item() is True
-    # ctx_frames header + flat codes
-    assert len(payload["code_predictor_codes"]) == 1 + 4 * 24
+    assert len(payload["code_predictor_codes"]) == _Q * 24
 
 
-def test_talker2code2wav_async_chunk_emits_eof_marker_when_finished_with_no_frames():
-    transfer_manager = SimpleNamespace(
-        code_prompt_token_ids=defaultdict(list),
-        connector=SimpleNamespace(config={"extra": {"codec_chunk_frames": 25, "codec_left_context_frames": 25}}),
-    )
+def test_emits_eof_marker_when_finished_with_no_frames():
+    tm = _tm()
     request = _req("rid-eof", finished=True)
 
     payload = talker2code2wav_async_chunk(
-        transfer_manager=transfer_manager,
+        transfer_manager=tm,
         pooling_output=None,
         request=request,
     )
@@ -70,3 +96,59 @@ def test_talker2code2wav_async_chunk_emits_eof_marker_when_finished_with_no_fram
         "code_predictor_codes": [],
         "finished": torch.tensor(True, dtype=torch.bool),
     }
+
+
+_CASES = [
+    # Normal path (initial=0): emit at chunk_size boundaries
+    ((25, 25, 0), (24, 0, False), None),
+    ((25, 25, 0), (25, 0, False), (0, 25)),
+    # Initial-chunk phase: hold, first emit, second emit
+    ((25, 25, 10), (9, 0, False), None),
+    ((25, 25, 10), (10, 0, False), (0, 10)),
+    ((25, 25, 10), (20, 1, False), (10, 20)),
+    # Non-divisible: holds at chunk boundary
+    ((25, 25, 12), (25, 2, False), None),
+    # Normal phase: offset by initial_coverage (chunk//initial * initial)
+    ((25, 25, 10), (45, 2, False), (20, 45)),
+    # Second normal emit (offset must stay stable)
+    ((25, 25, 10), (70, 3, False), (25, 50)),
+    # initial >= chunk clamps to chunk_size (behaves as normal)
+    ((25, 25, 30), (25, 0, False), (0, 25)),
+    # finished=True flushes IC tail
+    ((25, 25, 10), (5, 0, True), (0, 5)),
+    # finished=True flushes non-divisible IC residual
+    ((25, 25, 12), (25, 2, True), (24, 25)),
+    # finished=True flushes normal phase tail
+    ((25, 25, 10), (30, 2, True), (20, 30)),
+]
+
+
+@pytest.mark.parametrize("config, state, expected", _CASES)
+def test_streaming_decoding_with_variable_initial(config, state, expected):
+    chunk_frames, left_context, initial_chunk = config
+    n_frames, put_req, finished = state
+
+    tm = _tm(chunk_frames=chunk_frames, left_context=left_context, initial_chunk=initial_chunk)
+    payload = _call(tm, "r", n_frames=n_frames, put_req=put_req, finished=finished)
+
+    if expected is None:
+        assert payload is None
+    else:
+        exp_ctx, exp_window = expected
+        assert payload is not None
+        assert payload["left_context_size"] == exp_ctx
+        assert len(payload["code_predictor_codes"]) == _Q * exp_window
+
+
+def test_per_request_override_activates_initial_phase():
+    tm = _tm(initial_chunk=0)
+    payload = _call(tm, "r-override", n_frames=10, req_ic=10)
+    assert payload is not None
+    assert payload["left_context_size"] == 0
+    assert len(payload["code_predictor_codes"]) == _Q * 10
+
+
+def test_per_request_override_wins_over_stage_config():
+    tm = _tm(initial_chunk=5)
+    payload = _call(tm, "r-override2", n_frames=10, put_req=0, req_ic=15)
+    assert payload is None

tests/perf/scripts/run_benchmark.py

@@ -1,9 +1,5 @@
-import os
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
-
 import json
+import os
 import subprocess
 import threading
 from datetime import datetime
@@ -14,6 +10,9 @@ import pytest
 
 from tests.conftest import OmniServer, modify_stage_config
 
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
 
 def load_configs(config_path: str) -> list[dict[str, Any]]:
     try:
@@ -123,6 +122,8 @@ def run_benchmark(args: list, test_name: str, flow, dataset_name: str, num_promp
             "--endpoint",
             "/v1/chat/completions",
             "--save-result",
+            "--result-dir",
+            os.environ.get("BENCHMARK_DIR", "tests"),
             "--result-filename",
             result_filename,
         ]
@@ -137,9 +138,9 @@ def run_benchmark(args: list, test_name: str, flow, dataset_name: str, num_promp
     for line in iter(process.stderr.readline, ""):
         print(line, end=" ")
 
-    if "--result-dir" in args:
-        index = args.index("--result-dir")
-        result_dir = args[index + 1]
+    if "--result-dir" in command:
+        index = command.index("--result-dir")
+        result_dir = command[index + 1]
     else:
         result_dir = "./"
 
@@ -187,6 +188,10 @@ def benchmark_params(request, omni_server):
     if param_index >= len(all_params):
         raise ValueError(f"No benchmark parameters found for index {param_index} in test: {test_name}")
 
+    if all_params[param_index]["dataset_name"] == "random-mm":
+        # TODO: Due to known issues, skip the random-mm dataset.
+        pytest.skip("Skipping parameter for random-mm dataset.")
+
     current = param_index + 1
     total = len(all_params)
     print(f"\n  Running benchmark {current}/{total} for {test_name}")

tests/perf/stage_configs/qwen3_omni.yaml

@@ -85,7 +85,7 @@ stage_args:
       engine_output_type: audio  # Final output: audio waveform
       gpu_memory_utilization: 0.1
       distributed_executor_backend: "mp"
-      max_num_batched_tokens: 1000000
+      max_num_batched_tokens: 100000
       hf_config_name: thinker_config
     engine_input_source: [1]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav

tests/perf/tests/test.json

@@ -10,16 +10,14 @@
                 "dataset_name": "random",
                 "num_prompts": [
                     10,
-                    40,
-                    100
+                    40
                 ],
                 "max_concurrency": [
                     1,
-                    4,
-                    10
+                    4
                 ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
+                "random_input_len": 100,
+                "random_output_len": 100,
                 "ignore_eos": true,
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
@@ -32,13 +30,11 @@
                 "dataset_name": "random-mm",
                 "num_prompts": [
                     10,
-                    40,
-                    100
+                    40
                 ],
                 "max_concurrency": [
                     1,
-                    4,
-                    10
+                    4
                 ],
                 "random_input_len": 100,
                 "random_output_len": 100,
@@ -94,16 +90,14 @@
                 "dataset_name": "random",
                 "num_prompts": [
                     10,
-                    40,
-                    100
+                    40
                 ],
                 "max_concurrency": [
                     1,
-                    4,
-                    10
+                    4
                 ],
-                "random_input_len": 2500,
-                "random_output_len": 900,
+                "random_input_len": 100,
+                "random_output_len": 100,
                 "ignore_eos": true,
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
@@ -116,13 +110,11 @@
                 "dataset_name": "random-mm",
                 "num_prompts": [
                     10,
-                    40,
-                    100
+                    40
                 ],
                 "max_concurrency": [
                     1,
-                    4,
-                    10
+                    4
                 ],
                 "random_input_len": 100,
                 "random_output_len": 100,

tests/worker/test_gpu_generation_model_runner.py

@@ -34,6 +34,7 @@ def _make_runner(multimodal_outputs):
     runner.use_async_scheduling = False
     runner.device = torch.device("cpu")
     runner.supports_mm_inputs = False
+    runner.speculative_config = None
     return runner
 
 

tests/worker/test_omni_gpu_model_runner.py

@@ -64,6 +64,7 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
     # Minimal attributes used by OmniGPUModelRunner._talker_mtp_forward
     runner.input_batch = DummyInputBatch(list(req_ids))
     runner.requests = {rid: DummyReqState() for rid in req_ids}
+    runner.model_intermediate_buffer = {}
 
     # query_start_loc.cpu[req_index] is used to locate the token position
     # in the flattened `inputs_embeds`.
@@ -167,6 +168,61 @@ def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):
     assert torch.allclose(inputs_embeds, before)
 
 
+def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch):
+    """Validate that _update_intermediate_buffer writes to model_intermediate_buffer
+    (forward path) and mirrors to additional_information_cpu setattr (backward compat)."""
+    import vllm_omni.worker.gpu_model_runner as mod
+
+    monkeypatch.setattr(mod, "set_forward_context", _noop_forward_context)
+
+    runner = _make_runner(req_ids=("r1",), hidden_size=4)
+
+    update = {"my_tensor": torch.tensor([1.0, 2.0]), "my_list": [3, 4]}
+    OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", update)
+
+    # Forward: buffer is populated
+    assert "r1" in runner.model_intermediate_buffer
+    buf = runner.model_intermediate_buffer["r1"]
+    assert torch.allclose(buf["my_tensor"], torch.tensor([1.0, 2.0]))
+    assert buf["my_list"] == [3, 4]
+
+    # Backward compat: setattr is also populated
+    info_cpu = runner.requests["r1"].additional_information_cpu
+    assert torch.allclose(info_cpu["my_tensor"], torch.tensor([1.0, 2.0]))
+    assert info_cpu["my_list"] == [3, 4]
+
+
+def test_update_intermediate_buffer_accumulates():
+    """Validate that successive merges accumulate keys in the buffer."""
+    runner = _make_runner(req_ids=("r1",), hidden_size=4)
+
+    OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {"a": torch.tensor([1.0])})
+    OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {"b": torch.tensor([2.0])})
+
+    buf = runner.model_intermediate_buffer["r1"]
+    assert "a" in buf and "b" in buf
+    assert torch.allclose(buf["a"], torch.tensor([1.0]))
+    assert torch.allclose(buf["b"], torch.tensor([2.0]))
+
+
+def test_update_intermediate_buffer_skips_empty_update():
+    """Validate that an empty update dict is a no-op."""
+    runner = _make_runner(req_ids=("r1",), hidden_size=4)
+
+    OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {})
+
+    assert "r1" not in runner.model_intermediate_buffer
+
+
+def test_update_intermediate_buffer_skips_unknown_req_id():
+    """Validate that merge is a no-op when req_id is not in self.requests."""
+    runner = _make_runner(req_ids=("r1",), hidden_size=4)
+
+    OmniGPUModelRunner._update_intermediate_buffer(runner, "unknown_req", {"key": torch.tensor([1.0])})
+
+    assert "unknown_req" not in runner.model_intermediate_buffer
+
+
 def test_maybe_attach_mimo_audio_req_infos_enriches_dict():
     runner = _make_runner_for_mimo()
     req_id = "r_mimo"

tools/nightly/generate_nightly_perf_excel.py

@@ -0,0 +1,437 @@
+#!/usr/bin/env python3
+"""
+Generate a nightly Excel performance report from JSON results.
+
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+from collections.abc import Iterable, Sequence
+from datetime import datetime, timezone
+from typing import Any
+
+from openpyxl import Workbook
+from openpyxl.styles import PatternFill
+from openpyxl.utils import get_column_letter
+
+LOGGER = logging.getLogger(__name__)
+
+GREY_BLOCK_FILL = PatternFill(start_color="D3D3D3", fill_type="solid")
+
+# Benchmark metric columns: grey the latest row's cell when value changed vs previous date.
+BENCHMARK_COLUMNS: tuple[str, ...] = (
+    "num_prompts",
+    "request_rate",
+    "burstiness",
+    "max_concurrency",
+    "duration",
+    "completed",
+    "failed",
+    "request_throughput",
+    "output_throughput",
+    "total_token_throughput",
+    "mean_ttft_ms",
+    "p99_ttft_ms",
+    "mean_tpot_ms",
+    "p99_tpot_ms",
+    "mean_itl_ms",
+    "p99_itl_ms",
+    "mean_e2el_ms",
+    "p99_e2el_ms",
+    "mean_audio_rtf",
+    "p99_audio_rtf",
+    "mean_audio_duration_s",
+    "p99_audio_duration_s",
+)
+# Columns that get float coercion and number format in Excel. Excludes request_rate ("inf" str)
+# and max_concurrency (null); leave those as-is. If they become float in the future, they are
+# still written correctly from JSON without coercion here.
+NUMERIC_FORMAT_COLUMNS: tuple[str, ...] = tuple(
+    c for c in BENCHMARK_COLUMNS if c not in ("request_rate", "max_concurrency")
+)
+
+_COLUMNS_FILENAME = "nightly_perf_summary_columns.txt"
+DEFAULT_INPUT_DIR = os.getenv("DEFAULT_INPUT_DIR") if os.getenv("DEFAULT_INPUT_DIR") else "tests"
+DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR") if os.getenv("DEFAULT_OUTPUT_DIR") else "tests"
+
+
+def _load_summary_columns(script_dir: str) -> list[str]:
+    """Load summary column names from a file next to this script; fallback to default if missing."""
+    path = os.path.join(script_dir, _COLUMNS_FILENAME)
+    default = [
+        "date",
+        "endpoint_type",
+        "backend",
+        "model_id",
+        "tokenizer_id",
+        "num_prompts",
+        "request_rate",
+        "burstiness",
+        "max_concurrency",
+        "duration",
+        "completed",
+        "failed",
+        "request_throughput",
+        "output_throughput",
+        "total_token_throughput",
+        "mean_ttft_ms",
+        "p99_ttft_ms",
+        "mean_tpot_ms",
+        "p99_tpot_ms",
+        "mean_itl_ms",
+        "p99_itl_ms",
+        "mean_e2el_ms",
+        "p99_e2el_ms",
+        "mean_audio_rtf",
+        "p99_audio_rtf",
+        "mean_audio_duration_s",
+        "p99_audio_duration_s",
+        "commit_sha",
+        "build_id",
+        "build_url",
+        "source_file",
+    ]
+    if not os.path.isfile(path):
+        return default
+    columns: list[str] = []
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            s = line.strip()
+            if s and not s.startswith("#"):
+                columns.append(s)
+    return columns if columns else default
+
+
+def _vllm_omni_root() -> str:
+    """Resolve vllm-omni repo root: directory that contains a 'tests' subdir (and usually 'tools')."""
+    path = os.path.dirname(os.path.abspath(__file__))
+    while path and path != os.path.dirname(path):
+        if os.path.isdir(os.path.join(path, "tests")):
+            return path
+        path = os.path.dirname(path)
+    return os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))
+
+
+def _default_input_dir() -> str:
+    """Default: vllm-omni root / DEFAULT_INPUT_DIR (where performance JSON files live)."""
+    root = _vllm_omni_root()
+    return os.path.join(root, DEFAULT_INPUT_DIR)
+
+
+def _default_output_file() -> str:
+    """Default: vllm-omni root / DEFAULT_OUTPUT_DIR / nightly_perf_<timestamp>.xlsx."""
+    ts = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
+    return os.path.join(_vllm_omni_root(), DEFAULT_OUTPUT_DIR, f"nightly_perf_{ts}.xlsx")
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Read performance JSON files from vllm-omni/tests/ and generate an Excel report."
+    )
+    parser.add_argument(
+        "--input-dir",
+        type=str,
+        default=_default_input_dir(),
+        help="Directory containing performance JSON files; default is <vllm-omni-root>/DEFAULT_INPUT_DIR.",
+    )
+    parser.add_argument(
+        "--output-file",
+        type=str,
+        default=_default_output_file(),
+        help="Output path of the Excel report; \
+            default is <vllm-omni-root>/DEFAULT_OUTPUT_DIR/nightly_perf_<timestamp>.xlsx.",
+    )
+    parser.add_argument(
+        "--commit-sha",
+        type=str,
+        default=None,
+        help="Optional commit SHA; defaults to environment variable BUILDKITE_COMMIT if unset.",
+    )
+    parser.add_argument(
+        "--build-id",
+        type=str,
+        default=None,
+        help="Optional build ID; defaults to environment variable BUILDKITE_BUILD_ID if unset.",
+    )
+    parser.add_argument(
+        "--build-url",
+        type=str,
+        default=None,
+        help="Optional build URL; defaults to environment variable BUILDKITE_BUILD_URL if unset.",
+    )
+    return parser.parse_args()
+
+
+def _load_json_file(path: str) -> dict[str, Any] | None:
+    """Safely load a single JSON file; return None and log a warning on failure."""
+    try:
+        with open(path, encoding="utf-8") as f:
+            data = json.load(f)
+    except (OSError, json.JSONDecodeError) as exc:
+        LOGGER.warning("failed to load json '%s': %s", path, exc)
+        return None
+
+    if not isinstance(data, dict):
+        LOGGER.warning("json root in '%s' is not an object, skip", path)
+        return None
+
+    return data
+
+
+def _iter_json_records(input_dir: str) -> Iterable[dict[str, Any]]:
+    """Iterate over JSON files in the input directory and yield normalized records.
+    commit_sha/build_id/build_url are not set here; they are applied later only to
+    rows with the latest date (see _apply_build_metadata_to_latest_only).
+    """
+    if not os.path.isdir(input_dir):
+        LOGGER.warning("input dir '%s' does not exist or is not a directory", input_dir)
+        return
+
+    for entry in sorted(os.listdir(input_dir)):
+        if not entry.endswith(".json"):
+            continue
+        full_path = os.path.join(input_dir, entry)
+        if not os.path.isfile(full_path):
+            continue
+
+        data = _load_json_file(full_path)
+        if data is None:
+            continue
+
+        record: dict[str, Any] = dict(data)
+        record.setdefault("date", datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S"))
+        record["source_file"] = os.path.basename(full_path)
+        yield record
+
+
+def _collect_records(input_dir: str) -> list[dict[str, Any]]:
+    """Collect all JSON records into a list."""
+    records: list[dict[str, Any]] = []
+    for record in _iter_json_records(input_dir):
+        records.append(record)
+    return records
+
+
+def _apply_build_metadata_to_latest_only(
+    records: Sequence[dict[str, Any]],
+    commit_sha: str | None,
+    build_id: str | None,
+    build_url: str | None,
+) -> None:
+    """Set commit_sha, build_id, build_url only on rows with the latest date.
+    Other rows get None so that build info is not duplicated for older benchmark data.
+    """
+    if not records:
+        return
+    max_date = max((r.get("date") or "") for r in records)
+    for r in records:
+        if (r.get("date") or "") == max_date:
+            r["commit_sha"] = commit_sha
+            r["build_id"] = build_id
+            r["build_url"] = build_url
+        else:
+            r["commit_sha"] = None
+            r["build_id"] = None
+            r["build_url"] = None
+
+
+def _sort_records_for_summary(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Sort so that same model_id is grouped, newest date first within each group."""
+    by_date_desc = sorted(records, key=lambda r: (r.get("date") or ""), reverse=True)
+    return sorted(by_date_desc, key=lambda r: (r.get("model_id") or ""))
+
+
+def _values_differ(a: Any, b: Any) -> bool:
+    """Return True if two values are considered different; avoid direct float equality."""
+    if a is None and b is None:
+        return False
+    if a is None or b is None:
+        return True
+    if isinstance(a, float) and isinstance(b, float):
+        if a != a and b != b:
+            return False
+        if a != a or b != b:
+            return True
+        return abs(a - b) > 1e-9
+    return a != b
+
+
+def _apply_benchmark_change_highlight(
+    ws,
+    summary_columns: Sequence[str],
+    records: Sequence[dict[str, Any]],
+) -> None:
+    """Grey cells in the latest row of each model when a benchmark metric changed vs previous date."""
+    if not records:
+        return
+    col_to_index = {c: i + 1 for i, c in enumerate(summary_columns)}
+    # Walk by model_id blocks (records already sorted by model_id, date desc).
+    i = 0
+    while i < len(records):
+        model_id = records[i].get("model_id")
+        block_start = i
+        while i < len(records) and records[i].get("model_id") == model_id:
+            i += 1
+        block_end = i
+        newest_idx = block_start
+        prev_idx = block_start + 1 if block_start + 1 < block_end else None
+        if prev_idx is None:
+            continue
+        excel_row = newest_idx + 2
+        for col in BENCHMARK_COLUMNS:
+            if col not in col_to_index:
+                continue
+            cur_val = records[newest_idx].get(col)
+            prev_val = records[prev_idx].get(col)
+            if _values_differ(cur_val, prev_val):
+                ws.cell(row=excel_row, column=col_to_index[col]).fill = GREY_BLOCK_FILL
+
+
+def _build_raw_columns(
+    records: Sequence[dict[str, Any]],
+    summary_columns: Sequence[str],
+) -> list[str]:
+    """Infer the column set for the raw sheet based on all records."""
+    keys: set[str] = set()
+    for record in records:
+        keys.update(record.keys())
+    # Ensure summary columns appear first; remaining columns sorted alphabetically.
+    ordered_keys: list[str] = []
+    for key in summary_columns:
+        if key in keys:
+            ordered_keys.append(key)
+            keys.discard(key)
+    ordered_keys.extend(sorted(keys))
+    return ordered_keys
+
+
+def _to_float_if_numeric(value: Any) -> Any:
+    """Coerce to float when possible so Excel treats as number; avoid #### from narrow columns."""
+    if value is None:
+        return value
+    if isinstance(value, (int, float)):
+        return float(value) if isinstance(value, int) else value
+    if isinstance(value, str):
+        try:
+            return float(value)
+        except (ValueError, TypeError):
+            return value
+    return value
+
+
+def _write_sheet(
+    ws,
+    columns: Sequence[str],
+    rows: Iterable[dict[str, Any]],
+    numeric_columns: Sequence[str] = (),
+) -> None:
+    """Write column names and row data into the given worksheet."""
+    numeric_set = set(numeric_columns)
+    ws.append(list(columns))
+    for record in rows:
+        row_values = []
+        for col in columns:
+            v = record.get(col)
+            if col in numeric_set:
+                v = _to_float_if_numeric(v)
+            row_values.append(v)
+        ws.append(row_values)
+
+
+def _format_benchmark_columns(
+    ws,
+    columns: Sequence[str],
+    num_data_rows: int,
+) -> None:
+    """Set number format and column width for numeric benchmark columns so values display (no ####)."""
+    numeric_set = set(NUMERIC_FORMAT_COLUMNS)
+    for c, col_name in enumerate(columns):
+        if col_name not in numeric_set:
+            continue
+        col_letter = get_column_letter(c + 1)
+        ws.column_dimensions[col_letter].width = 14
+        for r in range(2, 2 + num_data_rows):
+            cell = ws.cell(row=r, column=c + 1)
+            if cell.value is not None and isinstance(cell.value, (int, float)):
+                cell.number_format = "0.0000"
+            elif isinstance(cell.value, str):
+                try:
+                    cell.value = float(cell.value)
+                    cell.number_format = "0.0000"
+                except (ValueError, TypeError):
+                    pass
+
+
+def _ensure_parent_dir(path: str) -> None:
+    """Ensure that the parent directory of the output file exists."""
+    parent = os.path.dirname(os.path.abspath(path))
+    if not parent:
+        return
+    os.makedirs(parent, exist_ok=True)
+
+
+def generate_excel_report(
+    input_dir: str,
+    output_file: str,
+    commit_sha: str | None,
+    build_id: str | None,
+    build_url: str | None,
+) -> None:
+    """Main logic: load JSON records and generate an Excel report."""
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    summary_columns = _load_summary_columns(script_dir)
+
+    records = _collect_records(input_dir)
+    if not records:
+        LOGGER.warning("no valid json records found under '%s'", input_dir)
+
+    sorted_records = _sort_records_for_summary(records)
+    _apply_build_metadata_to_latest_only(sorted_records, commit_sha, build_id, build_url)
+
+    wb = Workbook()
+    ws_summary = wb.active
+    ws_summary.title = "summary"
+
+    _write_sheet(ws_summary, summary_columns, sorted_records, numeric_columns=NUMERIC_FORMAT_COLUMNS)
+    _format_benchmark_columns(ws_summary, summary_columns, len(sorted_records))
+    _apply_benchmark_change_highlight(ws_summary, summary_columns, sorted_records)
+
+    if sorted_records:
+        raw_columns = _build_raw_columns(sorted_records, summary_columns)
+        ws_raw = wb.create_sheet(title="raw")
+        _write_sheet(ws_raw, raw_columns, sorted_records)
+
+    _ensure_parent_dir(output_file)
+    wb.save(output_file)
+    LOGGER.info("excel report saved to '%s'", output_file)
+
+
+def main() -> None:
+    """Command-line entrypoint."""
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    )
+
+    args = parse_args()
+
+    commit_sha = args.commit_sha or os.getenv("BUILDKITE_COMMIT")
+    build_id = args.build_id or os.getenv("BUILDKITE_BUILD_ID")
+    build_url = args.build_url or os.getenv("BUILDKITE_BUILD_URL")
+
+    generate_excel_report(
+        input_dir=args.input_dir,
+        output_file=args.output_file,
+        commit_sha=commit_sha,
+        build_id=build_id,
+        build_url=build_url,
+    )
+
+
+if __name__ == "__main__":
+    main()

tools/nightly/send_nightly_perf_email.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+"""
+Send the nightly performance Excel report by email to DAILY_EMAIL_LIST.
+
+Reads SMTP and recipient config from environment variables. Use --dry-run to
+validate config and print subject/body without sending.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+import smtplib
+import sys
+from email.header import Header
+from email.mime.application import MIMEApplication
+from email.mime.multipart import MIMEMultipart
+from email.mime.text import MIMEText
+
+LOGGER = logging.getLogger(__name__)
+
+# Env keys for SMTP and recipients.
+ENV_SMTP_HOST = "SMTP_HOST"
+ENV_SMTP_PORT = "SMTP_PORT"
+ENV_SMTP_USERNAME = "SMTP_USERNAME"
+ENV_SMTP_PASSWORD = "SMTP_PASSWORD"
+ENV_DAILY_EMAIL_LIST = "DAILY_EMAIL_LIST"
+ENV_EMAIL_SENDER = "EMAIL_SENDER"
+ENV_EMAIL_SUBJECT_PREFIX = "EMAIL_SUBJECT_PREFIX"
+ENV_BUILD_URL = "BUILDKITE_BUILD_URL"
+ENV_COMMIT = "BUILDKITE_COMMIT"
+
+DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR")
+
+# Do not attach Excel if size >= this (bytes); body will suggest downloading from build URL.
+MAX_ATTACHMENT_BYTES = 20 * 1024 * 1024
+
+SMTP_RETRIES = 3
+SMTP_RETRY_DELAY_SEC = 5
+
+
+def _get_required_env() -> dict[str, str]:
+    """Read required env vars; raise SystemExit with clear message if any missing."""
+    required = {
+        ENV_SMTP_HOST: os.environ.get(ENV_SMTP_HOST),
+        ENV_SMTP_PORT: os.environ.get(ENV_SMTP_PORT),
+        ENV_SMTP_USERNAME: os.environ.get(ENV_SMTP_USERNAME),
+        ENV_SMTP_PASSWORD: os.environ.get(ENV_SMTP_PASSWORD),
+        ENV_DAILY_EMAIL_LIST: os.environ.get(ENV_DAILY_EMAIL_LIST),
+    }
+    missing = [k for k, v in required.items() if not (v and str(v).strip())]
+    if missing:
+        raise SystemExit(f"Missing required env vars: {', '.join(missing)}. Set them (e.g. in Buildkite secrets).")
+    return {k: str(v).strip() for k, v in required.items()}
+
+
+def _get_latest_file(folder_path: str) -> str:
+    """Get the latest modified file from the folder path."""
+    files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(".xlsx")]
+    if not files:
+        raise SystemExit(f"No Excel files found in {folder_path}")
+    return max(files, key=os.path.getmtime)
+
+
+def _recipients_list(comma_separated: str) -> list[str]:
+    """Parse DAILY_EMAIL_LIST into a list of addresses."""
+    return [a.strip() for a in comma_separated.split(",") if a.strip()]
+
+
+def _build_body(
+    date_str: str,
+    commit_sha: str | None,
+    build_url: str | None,
+    attachment_skipped: bool,
+) -> str:
+    """Plain-text body with build metadata; note when attachment was skipped."""
+    lines = [
+        f"Nightly performance report date: {date_str}",
+        "",
+        f"Commit: {commit_sha or 'N/A'}",
+        f"Build: {build_url or 'N/A'}",
+        "",
+    ]
+    if attachment_skipped:
+        lines.append(
+            "Report file was too large to attach. Please download the Excel from the build artifacts (Build URL above)."
+        )
+        lines.append("")
+    return "\n".join(lines)
+
+
+def _build_subject(prefix: str | None, date_str: str) -> str:
+    """Subject line: optional prefix + date."""
+    base = f"Nightly Perf {date_str}"
+    if prefix and prefix.strip():
+        return f"{prefix.strip()} {base}"
+    return base
+
+
+def _send_mail(
+    report_file: str,
+    date_str: str,
+    dry_run: bool,
+) -> None:
+    """Load config, build message, and send (or dry-run)."""
+    cfg = _get_required_env()
+    recipients = _recipients_list(cfg[ENV_DAILY_EMAIL_LIST])
+    if not recipients:
+        raise SystemExit("DAILY_EMAIL_LIST is empty after parsing.")
+
+    commit_sha = os.environ.get(ENV_COMMIT)
+    build_url = os.environ.get(ENV_BUILD_URL)
+    sender = os.environ.get(ENV_EMAIL_SENDER) or cfg[ENV_SMTP_USERNAME]
+    prefix = os.environ.get(ENV_EMAIL_SUBJECT_PREFIX)
+
+    size = os.path.getsize(report_file) if os.path.isfile(report_file) else 0
+    attach_excel = size < MAX_ATTACHMENT_BYTES and size > 0
+    attachment_skipped = os.path.isfile(report_file) and not attach_excel
+
+    body = _build_body(
+        commit_sha=commit_sha, build_url=build_url, date_str=date_str, attachment_skipped=attachment_skipped
+    )
+    subject = _build_subject(prefix=prefix, date_str=date_str)
+
+    msg = MIMEMultipart()
+    msg["Subject"] = Header(subject, "utf-8")
+    msg["From"] = sender
+    msg["To"] = ", ".join(recipients)
+    msg.attach(MIMEText(body, "plain", "utf-8"))
+
+    if attach_excel:
+        with open(report_file, "rb") as f:
+            part = MIMEApplication(f.read(), _subtype="vnd.openxmlformats-officedocument.spreadsheetml.sheet")
+        part.add_header("Content-Disposition", "attachment", filename=os.path.basename(report_file))
+        msg.attach(part)
+
+    if dry_run:
+        LOGGER.info("dry-run: not sending mail")
+        print("To:", recipients, file=sys.stderr)
+        print("Subject:", subject, file=sys.stderr)
+        print("Attachment:", "yes" if attach_excel else "no (size limit)", file=sys.stderr)
+        print("Body preview:", body[:300] + ("..." if len(body) > 300 else ""), file=sys.stderr)
+        return
+
+    port = int(cfg[ENV_SMTP_PORT], 10)
+    last_err: Exception | None = None
+    for attempt in range(SMTP_RETRIES):
+        try:
+            with smtplib.SMTP(cfg[ENV_SMTP_HOST], port=port, timeout=30) as smtp:
+                smtp.starttls()
+                smtp.login(cfg[ENV_SMTP_USERNAME], cfg[ENV_SMTP_PASSWORD])
+                smtp.sendmail(sender, recipients, msg.as_string())
+            LOGGER.info("sent nightly perf email to %d recipient(s)", len(recipients))
+            return
+        except Exception as e:
+            last_err = e
+            LOGGER.warning("SMTP attempt %d/%d failed: %s", attempt + 1, SMTP_RETRIES, e)
+            if attempt < SMTP_RETRIES - 1:
+                import time
+
+                time.sleep(SMTP_RETRY_DELAY_SEC)
+    raise SystemExit(f"Failed to send email after {SMTP_RETRIES} attempts.") from last_err
+
+
+def _date_from_filename(path: str) -> str:
+    """Try to derive a date string from report filename (e.g. nightly_perf_20260211-020704.xlsx -> 2026-02-11)."""
+    base = os.path.splitext(os.path.basename(path))[0]
+    if base.startswith("nightly_perf_") and len(base) > 13:
+        raw = base.replace("nightly_perf_", "")[:8]
+        if len(raw) == 8 and raw.isdigit():
+            return f"{raw[:4]}-{raw[4:6]}-{raw[6:8]}"
+    return base or "unknown"
+
+
+def _vllm_omni_root() -> str:
+    """Resolve vllm-omni repo root: directory that contains a 'tests' subdir (and usually 'tools')."""
+    path = os.path.dirname(os.path.abspath(__file__))
+    while path and path != os.path.dirname(path):
+        if os.path.isdir(os.path.join(path, "tests")):
+            return path
+        path = os.path.dirname(path)
+    return os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))
+
+
+def _default_output_dir() -> str:
+    """Default: vllm-omni root / DEFAULT_OUTPUT_DIR (where performance .xlsx files live)."""
+    root = _vllm_omni_root()
+    return os.path.join(root, DEFAULT_OUTPUT_DIR)
+
+
+def parse_args() -> argparse.Namespace:
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Send nightly performance Excel report by email (config from env).",
+    )
+    parser.add_argument(
+        "--report-file",
+        type=str,
+        default=_default_output_dir(),
+        help="Folder/file path to the nightly_perf_*.xlsx file; default is DEFAULT_OUTPUT_DIR.",
+    )
+    parser.add_argument(
+        "--date",
+        type=str,
+        default=None,
+        help="Date string for subject/body (default: inferred from report filename).",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Print recipient, subject, and body; do not send.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    """Entrypoint."""
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    )
+    args = parse_args()
+
+    report_file = _get_latest_file(args.report_file) if os.path.isdir(args.report_file) else args.report_file
+    if not os.path.isfile(report_file):
+        raise SystemExit(f"Report file not found: {report_file}")
+
+    date_str = args.date or _date_from_filename(report_file)
+    _send_mail(report_file=report_file, date_str=date_str, dry_run=args.dry_run)
+
+
+if __name__ == "__main__":
+    main()