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

2026-02-05 01:44:34 +00:00
commit 0144a424ed
@@ -9,14 +9,16 @@ 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://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> |
 </p>

+
 ---

 *Latest News* 🔥
- [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.
- [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), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true).
+
+- [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.

 ---
@@ -67,6 +69,19 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
 We welcome and value any contributions and collaborations.
 Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en/latest/contributing/) for how to get involved.

+## Citation
+
+If you use vLLM-Omni for your research, please cite our [paper](https://arxiv.org/abs/2602.02204):
+
+```bibtex
+@article{yin2026vllmomni,
+  title={vLLM-Omni: Fully Disaggregated Serving for Any-to-Any Multimodal Models},
+  author={Peiqi Yin, Jiangyun Zhu, Han Gao, Chenguang Zheng, Yongxiang Huang, Taichang Zhou, Ruirui Yang, Weizhi Liu, Weiqing Chen, Canlin Guo, Didan Deng, Zifeng Mo, Cong Wang, James Cheng, Roger Wang, Hongsheng Liu},
+  journal={arXiv preprint arXiv:2602.02204},
+  year={2026}
+}
+```
+
 ## Join the Community
 Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).

@@ -3,19 +3,18 @@ ARG VLLM_BASE_TAG=v0.15.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
+
 COPY . .
+
 # Install system dependencies
 RUN apt-get update && \
-  apt-get install -y ffmpeg && \
-  apt-get clean && \
-  rm -rf /var/lib/apt/lists/*
-
-# Install vLLM wheel for the specific commit
-# URL format: https://wheels.vllm.ai/<commit_hash>/
-# Current vLLM commit: f176443446f659dbab5315e056e605d8984fd976 (v0.15.0)
-RUN uv pip install --system --upgrade --force-reinstall vllm --extra-index-url https://wheels.vllm.ai/f176443446f659dbab5315e056e605d8984fd976
+    apt-get install -y ffmpeg && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*

 # Install vllm-omni into the same uv-managed Python environment used by the base image.
-RUN uv pip install --system --no-cache-dir ".[dev]"
+RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
+
 RUN ln -sf /usr/bin/python3 /usr/bin/python
+
 ENTRYPOINT []
@@ -1,15 +0,0 @@
-ARG VLLM_ASCEND_IMAGE=quay.nju.edu.cn/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.11.0rc2
-FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
-
-ARG APP_DIR=/vllm-workspace/vllm-omni
-WORKDIR ${APP_DIR}
-
-COPY . .
-
-# Install vllm-omni with dev dependencies
-RUN pip install --no-cache-dir -e ".[dev]"
-
-ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
-
-ENTRYPOINT []
@@ -0,0 +1,19 @@
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.14.0rc1
+FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
+
+ARG APP_DIR=/vllm-workspace/vllm-omni
+WORKDIR ${APP_DIR}
+
+COPY . .
+
+# Remove this replace when the dispatch of requirements is ready
+RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
+ && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
+
+# Install vllm-omni with dev dependencies
+RUN pip install --no-cache-dir -e .
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT []
@@ -0,0 +1,19 @@
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
+FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
+
+ARG APP_DIR=/vllm-workspace/vllm-omni
+WORKDIR ${APP_DIR}
+
+COPY . .
+
+# Remove this replace when the dispatch of requirements is ready
+RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
+ && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
+
+# Install vllm-omni with dev dependencies
+RUN pip install --no-cache-dir -e .
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT []
@@ -11,17 +11,23 @@ nav:
  - Examples:
    - examples/README.md
    - Offline Inference:
+      - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
      - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
      - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
+      - LoRA Inference Examples: user_guide/examples/offline_inference/lora_inference.md
      - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
      - Qwen3-TTS Offline Inference: user_guide/examples/offline_inference/qwen3_tts.md
+      - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
      - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
      - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
    - Online Serving:
+      - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
      - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
+      - Online LoRA Inference (Diffusion): user_guide/examples/online_serving/lora_inference.md
      - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
+      - Qwen3-TTS Online Serving: user_guide/examples/online_serving/qwen3_tts.md
      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
  - General:
    - usage/*
@@ -54,6 +60,7 @@ nav:
    - Feature Design:
      - design/feature/disaggregated_inference.md
      - design/feature/ray_based_execution.md
+      - design/feature/omni_connectors/
    - Module Design:
      - design/module/ar_module.md
      - design/module/dit_module.md
@@ -10,6 +10,9 @@ Main entry points for vLLM-Omni inference and serving.
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
 - [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
+- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
+- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
+- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
 - [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
 - [vllm_omni.entrypoints.log_utils.OrchestratorMetrics][]
@@ -26,7 +29,9 @@ Main entry points for vLLM-Omni inference and serving.

 Input data structures for multi-modal inputs.

+- [vllm_omni.inputs.data.OmniDiffusionSamplingParams][]
 - [vllm_omni.inputs.data.OmniEmbedsPrompt][]
+- [vllm_omni.inputs.data.OmniTextPrompt][]
 - [vllm_omni.inputs.data.OmniTokenInputs][]
 - [vllm_omni.inputs.data.OmniTokensPrompt][]
 - [vllm_omni.inputs.parse.parse_singleton_prompt_omni][]
@@ -58,6 +63,7 @@ Core scheduling and caching components.
 - [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
 - [vllm_omni.core.sched.output.OmniCachedRequestData][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
+- [vllm_omni.core.sched.output.OmniSchedulerOutput][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][]
@@ -88,20 +94,23 @@ Configuration classes.

 Worker classes and model runners for distributed inference.

- [vllm_omni.diffusion.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]
- [vllm_omni.diffusion.worker.gpu_diffusion_worker.GPUDiffusionWorker][]
- [vllm_omni.diffusion.worker.gpu_diffusion_worker.WorkerProc][]
- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][]
- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][]
+- [vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner][]
+- [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
+- [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
+- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
+- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
+- [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]
+- [vllm_omni.platforms.npu.worker.npu_generation_model_runner.NPUGenerationModelRunner][]
+- [vllm_omni.platforms.npu.worker.npu_generation_worker.NPUGenerationWorker][]
+- [vllm_omni.platforms.npu.worker.npu_model_runner.OmniNPUModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_ar_model_runner.XPUARModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker][]
+- [vllm_omni.platforms.xpu.worker.xpu_generation_model_runner.XPUGenerationModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_generation_worker.XPUGenerationWorker][]
 - [vllm_omni.worker.gpu_ar_model_runner.ExecuteModelState][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
 - [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
 - [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
 - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
- [vllm_omni.worker.npu.npu_ar_model_runner.ExecuteModelState][]
- [vllm_omni.worker.npu.npu_ar_model_runner.NPUARModelRunner][]
- [vllm_omni.worker.npu.npu_ar_worker.NPUARWorker][]
- [vllm_omni.worker.npu.npu_generation_model_runner.NPUGenerationModelRunner][]
- [vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker][]
- [vllm_omni.worker.npu.npu_model_runner.OmniNPUModelRunner][]
+- [vllm_omni.worker.mixins.OmniWorkerMixin][]
@@ -52,7 +52,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
 - `--percentile-metrics`  
        "Comma-separated list of selected metrics to report percentiles."
                    "This argument specifies the metrics to report percentiles."
-                    'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf". '
+                    'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '

 - `--save-result`  
 Specify to save benchmark results to a json file
@@ -220,48 +220,48 @@ If successful, you will see the following output:
 ============ Serving Benchmark Result ============
 Successful requests:                     2
 Failed requests:                         0
-Benchmark duration (s):                  24.35
-Request throughput (req/s):              0.08
+Benchmark duration (s):                  3.89
+Request throughput (req/s):              0.51
 Peak concurrent requests:                2.00
 ----------------End-to-end Latency----------------
-Mean E2EL (ms):                          22576.23
-Median E2EL (ms):                        22576.23
-P99 E2EL (ms):                           24205.72
+Mean E2EL (ms):                          3824.76
+Median E2EL (ms):                        3824.76
+P99 E2EL (ms):                           3888.54
 ================== Text Result ===================
 Total input tokens:                      30
-Total generated tokens:                  8973
-Output token throughput (tok/s):         368.52
-Peak output token throughput (tok/s):    81.00
+Total generated tokens:                  10101
+Output token throughput (tok/s):         2595.57
+Peak output token throughput (tok/s):    111.00
 Peak concurrent requests:                2.00
-Total Token throughput (tok/s):          369.76
+Total Token throughput (tok/s):          2603.28
 ---------------Time to First Token----------------
-Mean TTFT (ms):                          125.16
-Median TTFT (ms):                        125.16
-P99 TTFT (ms):                           155.88
+Mean TTFT (ms):                          117.15
+Median TTFT (ms):                        117.15
+P99 TTFT (ms):                           142.69
 -----Time per Output Token (excl. 1st token)------
-Mean TPOT (ms):                          5.01
-Median TPOT (ms):                        5.01
-P99 TPOT (ms):                           5.42
+Mean TPOT (ms):                          0.73
+Median TPOT (ms):                        0.73
+P99 TPOT (ms):                           0.74
 ---------------Inter-token Latency----------------
-Mean ITL (ms):                           34.15
-Median ITL (ms):                         0.01
-P99 ITL (ms):                            376.19
+Mean ITL (ms):                           16.47
+Median ITL (ms):                         16.19
+P99 ITL (ms):                            52.55
 ================== Audio Result ==================
-Total audio duration generated(s):       3.95
-Total audio frames generated:            94890
-Audio throughput(audio duration/s):      0.16
+Total audio duration generated(s):       15.79
+Total audio frames generated:            379050
+Audio throughput(audio duration/s):      4.06
 ---------------Time to First Packet---------------
-Mean AUDIO_TTFP (ms):                    11756.89
-Median AUDIO_TTFP (ms):                  11756.89
-P99 AUDIO_TTFP (ms):                     20854.25
+Mean AUDIO_TTFP (ms):                    3701.37
+Median AUDIO_TTFP (ms):                  3701.37
+P99 AUDIO_TTFP (ms):                     3762.25
 -----------------Real Time Factor-----------------
-Mean AUDIO_RTF:                          3.75
-Median AUDIO_RTF:                        3.75
-P99 AUDIO_RTF:                           7.39
+Mean AUDIO_RTF:                          0.47
+Median AUDIO_RTF:                        0.47
+P99 AUDIO_RTF:                           0.48
 ==================================================
 ```
 Notes:
-We use (audio generation time - first packet latency) / audio duration to calculate RTF.
+We use audio generation time / audio duration to calculate RTF.

 </details>

@@ -22,7 +22,7 @@ Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.1

 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.14.0 --torch-backend=auto
+uv pip install vllm==0.15.0 --torch-backend=auto
 ```

 #### Installation of vLLM-Omni
@@ -36,10 +36,10 @@ uv pip install vllm-omni
 # --8<-- [start:build-wheel-from-source]

 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.14.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.15.0 release version of the library

 ```bash
-uv pip install vllm==0.14.0 --torch-backend=auto
+uv pip install vllm==0.15.0 --torch-backend=auto
 ```

 The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
@@ -58,11 +58,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.14.0
+git checkout v0.15.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.14.0/vllm-0.14.0-cp38-abi3-manylinux_2_31_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.15.0/vllm-0.15.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -16,7 +16,13 @@ vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a c

 ## Installation

-### Recommended
+### Set up using Docker
+
+=== "NPU"
+
+    --8<-- "docs/getting_started/installation/npu/npu.inc.md:pre-built-images"
+
+### Build wheel from source

 === "NPU"

@@ -13,10 +13,10 @@ export DEVICE0=/dev/davinci0
 export DEVICE1=/dev/davinci1
 # Update the vllm-ascend image
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
 docker run --rm \
    --name vllm-omni-npu \
    --shm-size=1g \
@@ -34,26 +34,82 @@ docker run --rm \
    -p 8000:8000 \
    -it $IMAGE bash

-# Install the missing dependency of mooncake in the origin image.
-apt update
-apt install libjemalloc2
-echo "export LD_PRELOAD=/usr/lib/$(uname -m)-linux-gnu/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
-source ~/.bashrc
-
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
 git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
+
+# Remove this replace when the dispatch of requirements is ready
+RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
+ && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
+
 cd vllm-omni
 pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
-
-# (Optional) Disable mooncake for stable capability
-mv /usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake \
-   /usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake.disabled
 ```

 The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code placed in `/vllm-workspace` installed in development mode.

 For other installation methods (pip installation, building from source, custom Docker builds), please refer to the [vllm-ascend installation guide](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).

+We are keeping [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
+
 # --8<-- [end:installation]
+
+# --8<-- [start:pre-built-images]
+
+`vllm-ascend` offers Docker images for deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
+
+Supported images as following.
+
+| image name | Hardware | OS |
+|-|-|-|
+| image-tag | Atlas A2 | Ubuntu |
+| image-tag-a3 | Atlas A3 | Ubuntu |
+
+vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
+
+Here's an example deployment command that has been verified on 6 x NPUs:
+
+```bash
+# Atlas A2:
+# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+# Atlas A3:
+# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0-a3
+export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+docker run --rm \
+    --name vllm-omni-npu \
+    --shm-size=1g \
+    --device /dev/davinci0 \
+    --device /dev/davinci1 \
+    --device /dev/davinci2 \
+    --device /dev/davinci3 \
+    --device /dev/davinci4 \
+    --device /dev/davinci5 \
+    --device /dev/davinci_manager \
+    --device /dev/devmm_svm \
+    --device /dev/hisi_hdc \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
+    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
+    -v /etc/ascend_install.info:/etc/ascend_install.info \
+    -v /root/.cache:/root/.cache \
+    -p 8000:8000 \
+    -it $IMAGE bash
+```
+
+!!! tip
+    You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
+
+Or build IMAGE from **source code**:
+
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm-omni
+# A2
+docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu .
+# A3
+# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 .
+```
+
+# --8<-- [end:pre-built-images]
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate

 # On CUDA
-uv pip install vllm==0.14.0 --torch-backend=auto
+uv pip install vllm==0.15.0 --torch-backend=auto

 # On ROCm
-uv pip install vllm==0.14.0 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
+uv pip install vllm==0.15.0 --extra-index-url https://wheels.vllm.ai/rocm/0.15.0/rocm700

 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -25,6 +25,7 @@ th {
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
+| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` |
 |`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` |
@@ -51,6 +52,7 @@ th {

 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
+| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
@@ -59,3 +61,8 @@ th {
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
+|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
+|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
@@ -50,6 +50,7 @@ The following table shows which models are currently supported by each accelerat
 | **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ |
 | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ |
 | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ |
+| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ❌ |

 ### VideoGen

@@ -2,6 +2,7 @@

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

+
 ## Set up

 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.
@@ -99,7 +100,7 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \

 BAGEL-7B-MoT supports **multiple modality modes** for different use cases.

-The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
+The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)

 #### 📌 Command Line Arguments (end2end.py)

@@ -177,3 +178,10 @@ sudo apt install ffmpeg
 | Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
 | Stage-1 (DiT)       | **26.50 GiB**                |
 | Total               | **~42 GiB + KV Cache**       |
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/bagel/end2end.py"
+    ``````
@@ -47,10 +47,15 @@ Key arguments:
 - `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images.
 - `--prompt` / `--negative_prompt`: text description (string).
 - `--cfg_scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.
- `--cfg_parallel_size`: the number of devices to run CFG Parallel. CFG Parallel is valid only if classifier-free guidance is enabled and `cfg_parallel_size` is set to 2.
 - `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--output`: path to save the generated PNG.
+- `--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`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.

 ## Example materials

@@ -52,12 +52,17 @@ Key arguments:
 - `--num_frames`: Number of frames (default 81).
 - `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high-noise stages for MoE).
 - `--negative_prompt`: Optional list of artifacts to suppress.
- `--cfg_parallel_size`: the number of devices to run CFG Parallel. CFG Parallel is valid only if classifier-free guidance is enabled and `cfg_parallel_size` is set to 2.
 - `--boundary_ratio`: Boundary split ratio for two-stage MoE models.
 - `--flow_shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
 - `--num_inference_steps`: Number of denoising steps (default 50).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: Path to save the generated video.
+- `--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`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.

 ## Example materials

@@ -1,8 +1,9 @@
-# LoRA-Inference
+# LoRA Inference(Diffusion)

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

-This contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
+
+This directory contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
 The example uses the  `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.

 ## Overview
@@ -1,4 +1,4 @@
-# Qwen3-TTS Offline Inference
+# Qwen3-TTS

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

@@ -16,6 +16,15 @@ Qwen3 TTS provides multiple task variants for speech generation:
 ## 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.

+### ROCm Dependencies
+
+You will need to install these two dependencies `onnxruntime-rocm` and `sox`.
+
+```
+pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
+pip install onnxruntime-rocm sox
+```
+
 ## Quick Start

 Run a single sample for a task:
@@ -0,0 +1,47 @@
+# Text-To-Audio
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_audio>.
+
+
+The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts.
+
+## Prerequisites
+
+If you use a gated model (e.g., `stabilityai/stable-audio-open-1.0`), ensure you have access:
+
+1. **Accept Model License**: Visit the model page on Hugging Face (e.g., [stabilityai/stable-audio-open-1.0]) and accept the user agreement.
+2. **Authenticate**: Log in to Hugging Face locally to access the gated model.
+   ```bash
+   huggingface-cli login
+   ```
+
+## Local CLI Usage
+
+```bash
+python text_to_audio.py \
+  --model stabilityai/stable-audio-open-1.0 \
+  --prompt "The sound of a hammer hitting a wooden surface" \
+  --negative_prompt "Low quality" \
+  --seed 42 \
+  --guidance_scale 7.0 \
+  --audio_length 10.0 \
+  --num_inference_steps 100 \
+  --output stable_audio_output.wav
+```
+
+Key arguments:
+
+- `--prompt`: text description (string).
+- `--negative_prompt`: negative prompt for classifier-free guidance.
+- `--seed`: integer seed for deterministic generation.
+- `--guidance_scale`: classifier-free guidance scale.
+- `--audio_length`: audio duration in seconds.
+- `--num_inference_steps`: diffusion sampling steps.(more steps = higher quality, slower).
+- `--output`: path to save the generated WAV file.
+
+## Example materials
+
+??? abstract "text_to_audio.py"
+    ``````py
+    --8<-- "examples/offline_inference/text_to_audio/text_to_audio.py"
+    ``````
@@ -51,7 +51,7 @@ if __name__ == "__main__":

    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
-    batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+    batched inputs, you can [modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a longer batch of prompts.

 Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
 This is useful for models that support negative prompts.
@@ -95,11 +95,16 @@ Key arguments:
 - `--prompt`: text description (string).
 - `--seed`: integer seed for deterministic sampling.
 - `--cfg_scale`: true CFG scale (model-specific guidance strength).
- `--cfg_parallel_size`: the number of devices to run CFG Parallel. CFG Parallel is valid only if classifier-free guidance is enabled and `cfg_parallel_size` is set to 2.
 - `--num_images_per_prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--height/--width`: output resolution (defaults 1024x1024).
 - `--output`: path to save the generated PNG.
+- `--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`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.

 > ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.

@@ -12,10 +12,11 @@ python text_to_video.py \
  --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
  --negative_prompt "<optional quality filter>" \
  --height 480 \
-  --width 640 \
-  --num_frames 32 \
+  --width 832 \
+  --num_frames 33 \
  --guidance_scale 4.0 \
  --guidance_scale_high 3.0 \
+  --flow_shift 12.0 \
  --num_inference_steps 40 \
  --fps 16 \
  --output t2v_out.mp4
@@ -24,14 +25,19 @@ python text_to_video.py \
 Key arguments:

 - `--prompt`: text description (string).
- `--height/--width`: output resolution (defaults 720x1280). Dimensions should align with Wan VAE downsampling (multiples of 8).
+- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
 - `--num_frames`: Number of frames (Wan default is 81).
- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high)..
+- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high).
 - `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
- `--cfg_parallel_size`: the number of devices to run CFG Parallel. CFG Parallel is valid only if classifier-free guidance is enabled and `cfg_parallel_size` is set to 2.
- `--boundary_ratio`: Boundary split ratio for low/high DiT.
+- `--boundary_ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
+- `--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`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.

 ## Example materials

@@ -2,9 +2,10 @@

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

+
 ## 🛠️ Installation

-Please refer to [README.md](../../../README.md)
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)

 ## Run examples (BAGEL-7B-MoT)

@@ -74,7 +75,7 @@ python openai_chat_client.py \

 BAGEL-7B-MoT supports **multiple modality modes** for different use cases.

-The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
+The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)

 | Modality    | Input        | Output | Description                            |
 | ----------- | ------------ | ------ | -------------------------------------- |
@@ -230,3 +231,14 @@ sudo apt install ffmpeg
 | Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
 | Stage-1 (DiT)       | **26.50 GiB**                |
 | Total               | **~42 GiB + KV Cache**       |
+
+## Example materials
+
+??? abstract "openai_chat_client.py"
+    ``````py
+    --8<-- "examples/online_serving/bagel/openai_chat_client.py"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/bagel/run_server.sh"
+    ``````
@@ -1,7 +1,8 @@
-# LoRA-Inference
+# LoRA Inference(Diffusion)

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

+
 This example shows how to use **per-request LoRA** with vLLM-Omni diffusion models via the OpenAI-compatible Chat Completions API.

 > Note: The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
@@ -55,15 +56,17 @@ lora_adapter/
 └── adapter_model.safetensors
 ```

+## Example materials
+
 ??? abstract "openai_chat_client.py"
    ``````py
    --8<-- "examples/online_serving/lora_inference/openai_chat_client.py"
    ``````
 ??? abstract "run_curl_lora_inference.sh"
-    ``````py
+    ``````sh
    --8<-- "examples/online_serving/lora_inference/run_curl_lora_inference.sh"
    ``````
 ??? abstract "run_server.sh"
-    ``````py
+    ``````sh
    --8<-- "examples/online_serving/lora_inference/run_server.sh"
    ``````
@@ -0,0 +1,189 @@
+# Qwen3-TTS Online Serving
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_tts>.
+
+
+This directory contains examples for running Qwen3-TTS models with vLLM-Omni's online serving API.
+
+## Supported Models
+
+| Model | Task Type | Description |
+|-------|-----------|-------------|
+| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
+
+## Quick Start
+
+### 1. Start the Server
+
+```bash
+# CustomVoice model (default)
+./run_server.sh
+
+# Or specify task type
+./run_server.sh CustomVoice
+./run_server.sh VoiceDesign
+./run_server.sh Base
+```
+
+### 2. Run the Client
+
+```bash
+# CustomVoice: Use predefined speaker
+python openai_speech_client.py \
+    --text "你好，我是通义千问" \
+    --voice Vivian \
+    --language Chinese
+
+# CustomVoice with style instruction
+python openai_speech_client.py \
+    --text "今天天气真好" \
+    --voice Ryan \
+    --instructions "用开心的语气说"
+
+# VoiceDesign: Describe the voice style
+python openai_speech_client.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
+    --task-type VoiceDesign \
+    --text "哥哥，你回来啦" \
+    --instructions "体现撒娇稚嫩的萝莉女声，音调偏高"
+
+# Base: Voice cloning
+python openai_speech_client.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --task-type Base \
+    --text "Hello, this is a cloned voice" \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Original transcript of the reference audio"
+```
+
+### 3. Using curl
+
+```bash
+# Simple TTS request
+curl -X POST http://localhost:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "Vivian",
+        "language": "English"
+    }' --output output.wav
+
+# With style instruction
+curl -X POST http://localhost:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "I am so excited!",
+        "voice": "Vivian",
+        "instructions": "Speak with great enthusiasm"
+    }' --output excited.wav
+
+# List available voices in CustomVoice models
+curl http://localhost:8000/v1/audio/voices
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/speech
+```
+
+This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
+
+### Request Body
+
+```json
+{
+    "input": "Text to synthesize",
+    "voice": "Vivian",
+    "response_format": "wav",
+    "task_type": "CustomVoice",
+    "language": "Auto",
+    "instructions": "Optional style instructions",
+    "ref_audio": "URL or base64 for voice cloning",
+    "ref_text": "Reference audio transcript",
+    "x_vector_only_mode": false,
+    "max_new_tokens": 2048
+}
+```
+
+> **Note:** The `model` field is optional when serving a single model, as the server already knows which model is loaded.
+
+### Response
+
+Returns audio data in the requested format (default: WAV).
+
+## Parameters
+
+### Standard OpenAI Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | required | Text to synthesize |
+| `voice` | string | "Vivian" | Speaker/voice name |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
+| `model` | string | optional | Model name (optional when serving single model) |
+
+### Qwen3-TTS Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base |
+| `language` | string | "Auto" | Language: Auto, Chinese, English, Japanese, Korean |
+| `instructions` | string | "" | Voice style/emotion instructions |
+| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
+
+### Voice Clone Parameters (Base task)
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `ref_audio` | string | Yes* | Reference audio (file path, URL, or base64) |
+| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
+| `x_vector_only_mode` | bool | false | Use speaker embedding only (no ICL) |
+
+## Python Usage
+
+```python
+import httpx
+
+# Simple request
+response = httpx.post(
+    "http://localhost:8000/v1/audio/speech",
+    json={
+        "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+        "input": "Hello world",
+        "voice": "Vivian",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Limitations
+
+- **No streaming**: Audio is generated completely before being returned. Streaming will be supported after the pipeline is disaggregated (see RFC #938).
+- **Single request**: Batch processing is not yet optimized for online serving.
+
+## Troubleshooting
+
+1. **Connection refused**: Make sure the server is running on the correct port
+2. **Out of memory**: Reduce `--gpu-memory-utilization` in run_server.sh
+3. **Unsupported speaker**: Check supported speakers via model documentation
+4. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning
+
+## Example materials
+
+??? abstract "openai_speech_client.py"
+    ``````py
+    --8<-- "examples/online_serving/qwen3_tts/openai_speech_client.py"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/qwen3_tts/run_server.sh"
+    ``````
@@ -1,4 +1,4 @@
-# LoRA Inference Examples
+# LoRA Inference(Diffusion)

 This directory contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
 The example uses the  `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.
@@ -1,4 +1,4 @@
-# Qwen3-TTS Offline Inference
+# Qwen3-TTS

 This directory contains an offline demo for running Qwen3 TTS models with vLLM Omni. It builds task-specific inputs and generates WAV files locally.

@@ -12,6 +12,9 @@ For **multi-image** input editing, use **Qwen-Image-Edit-2509** (QwenImageEditPl
 vllm serve Qwen/Qwen-Image-Edit --omni --port 8092
 ```

+!!! note
+    If you encounter Out-of-Memory (OOM) issues or have limited GPU memory, you can enable VAE slicing and tiling to reduce memory usage, --vae-use-slicing --vae-use-tiling
+
 ### Multi-Image Edit (Qwen-Image-Edit-2509)

 ```bash
@@ -1,4 +1,4 @@
-# Online LoRA Inference (Diffusion)
+# LoRA Inference(Diffusion)

 This example shows how to use **per-request LoRA** with vLLM-Omni diffusion models via the OpenAI-compatible Chat Completions API.

@@ -308,19 +308,24 @@ def get_multi_audios_query(custom_prompt: str | None = None):

 def get_use_audio_in_video_query(
    video_path: str | None = None,
-    audio_path: str | None = None,
    custom_prompt: str | None = None,
 ):
+    """Query for use_audio_in_video mode.
+
+    When use_audio_in_video=True, audio is automatically extracted from the video
+    by the server. Do NOT send a separate audio_url - this would cause a mismatch
+    between the number of audio and video items.
+    """
    question = custom_prompt or (
        "Describe the content of the video in details, then convert what the baby say into text."
    )
    video_url = get_video_url_from_path(video_path)
-    audio_url = get_audio_url_from_path(audio_path)
+    # Note: audio is extracted from video automatically when use_audio_in_video=True
+    # Do not include a separate audio_url here
    return {
        "role": "user",
        "content": [
            {"type": "video_url", "video_url": {"url": video_url}},
-            {"type": "audio_url", "audio_url": {"url": audio_url}},
            {"type": "text", "text": question},
        ],
    }
@@ -397,7 +402,6 @@ def run_multimodal_generation(args) -> None:
    elif args.query_type == "use_audio_in_video":
        prompt = query_func(
            video_path=video_path,
-            audio_path=audio_path,
            custom_prompt=custom_prompt,
        )
    else:
@@ -484,7 +488,7 @@ def parse_args():
        "--query-type",
        "-q",
        type=str,
-        default="use_mixed_modalities",
+        default="use_audio_in_video",
        choices=query_map.keys(),
        help="Query type.",
    )
@@ -1,4 +1,4 @@
-# Qwen3-TTS Online Serving
+# Qwen3-TTS

 This directory contains examples for running Qwen3-TTS models with vLLM-Omni's online serving API.

@@ -44,7 +44,7 @@ dependencies = [
    "imageio[ffmpeg]>=2.37.2",
    "onnxruntime>=1.19.0",
    "sox>=1.5.0",
-    # "vllm==0.14.0",  # TODO: fix the entrypoints overwrite problem
+    # "vllm==0.15.0",  # TODO: fix the entrypoints overwrite problem
 ]

 [project.optional-dependencies]
@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for SP subgroup construction in parallel_state.py."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed.parallel_state import RankGenerator, set_seq_parallel_pg
+
+pytestmark = [
+    pytest.mark.diffusion,
+    pytest.mark.parallel,
+]
+
+
+def _fake_new_group_factory(created_groups: list[SimpleNamespace]):
+    def _fake_new_group(ranks, *args, **kwargs):
+        group = SimpleNamespace(ranks=list(ranks))
+        created_groups.append(group)
+        return group
+
+    return _fake_new_group
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "rank, expected_ulysses, expected_ring",
+    [
+        (0, [0, 2], [0]),
+        (1, [1, 3], [1]),
+        (2, [0, 2], [2]),
+        (3, [1, 3], [3]),
+    ],
+)
+def test_set_seq_parallel_pg_uses_explicit_sp_groups(rank, expected_ulysses, expected_ring, monkeypatch):
+    created_groups: list[SimpleNamespace] = []
+    monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups))
+
+    # tp=2, sp=2 -> SP groups are non-contiguous: [0,2] and [1,3]
+    sp_group_ranks = RankGenerator(2, 2, 1, 1, 1, "tp-sp-pp-cfg-dp").get_ranks("sp")
+
+    ulysses_pg, ring_pg = set_seq_parallel_pg(
+        sp_ulysses_degree=2,
+        sp_ring_degree=1,
+        rank=rank,
+        world_size=4,
+        sp_group_ranks=sp_group_ranks,
+    )
+
+    assert ulysses_pg.ranks == expected_ulysses
+    assert ring_pg.ranks == expected_ring
+
+
+@pytest.mark.cpu
+def test_set_seq_parallel_pg_validates_sp_group_ranks(monkeypatch):
+    created_groups: list[SimpleNamespace] = []
+    monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups))
+
+    # world_size=4, sp_size=2 -> expect 2 groups, provide 1 to trigger validation
+    with pytest.raises(ValueError, match="Invalid sp_group_ranks"):
+        set_seq_parallel_pg(
+            sp_ulysses_degree=2,
+            sp_ring_degree=1,
+            rank=0,
+            world_size=4,
+            sp_group_ranks=[[0, 2]],
+        )
@@ -14,13 +14,17 @@ class MultiModalsBenchmarkMetrics(BenchmarkMetrics):
    median_audio_ttfp_ms: float = 0.0
    std_audio_ttfp_ms: float = 0.0
    percentiles_audio_ttfp_ms: list[tuple[float, float]] = None
-    total_audio_duration_ms: float = 0.0
+    total_audio_duration_s: float = 0.0
    total_audio_frames: int = 0
    audio_throughput: float = 0.0
    mean_audio_rtf: float = 0.0
    median_audio_rtf: float = 0.0
    std_audio_rtf: float = 0.0
    percentiles_audio_rtf: list[tuple[float, float]] = None
+    mean_audio_duration_s: float = 0.0
+    median_audio_duration_s: float = 0.0
+    std_audio_duration_s: float = 0.0
+    percentiles_audio_duration_s: list[tuple[float, float]] = None


 def print_metrics(
@@ -73,7 +77,7 @@ def print_text_metrics(task_type, selected_percentile_metrics, metrics: MultiMod

 def print_audio_metrics(selected_percentile_metrics, metrics: MultiModalsBenchmarkMetrics):
    print("{s:{c}^{n}}".format(s=" Audio Result ", n=50, c="="))
-    print("{:<40} {:<10.2f}".format("Total audio duration generated(s):", metrics.total_audio_duration_ms))
+    print("{:<40} {:<10.2f}".format("Total audio duration generated(s):", metrics.total_audio_duration_s))
    print("{:<40} {:<10}".format("Total audio frames generated:", metrics.total_audio_frames))
    print("{:<40} {:<10.2f}".format("Audio throughput(audio duration/s):", metrics.audio_throughput))
    for metric in selected_percentile_metrics:
@@ -92,16 +96,23 @@ def process_one_metric(
        "e2el": "End-to-end Latency",
        "audio_ttfp": "Time to First Packet",
        "audio_rtf": "Real Time Factor",
+        "audio_duration": "Audio Duration",
    }

    header = metric_header_map.get(metric_attribute_name, metric_attribute_name)
    print("{s:{c}^{n}}".format(s=header, n=50, c="-"))

    is_audio_rtf = metric_attribute_name == "audio_rtf"
+    is_audio_duration = metric_attribute_name == "audio_duration"

-    suffix = "" if is_audio_rtf else "_ms"
-    unit_suffix = "" if is_audio_rtf else " (ms)"
-
+    suffix = "_ms"
+    unit_suffix = " (ms)"
+    if is_audio_duration:
+        suffix = "_s"
+        unit_suffix = " (s)"
+    elif is_audio_rtf:
+        suffix = ""
+        unit_suffix = ""
    mean_attr_name = f"mean_{metric_attribute_name}{suffix}"
    mean_value = getattr(metrics, mean_attr_name, 0.0)
    print(f"{f'Mean {metric_attribute_name.upper()}{unit_suffix}:':<40} {mean_value:<10.2f}")
@@ -212,6 +223,9 @@ def calculate_metrics(
                good_completed += 1

    if completed == 0:
+        warnings.formatwarning = (
+            lambda msg, category, filename, lineno, line=None: f"{filename}:{lineno}: {category.__name__}: {msg}\n"
+        )
        warnings.warn(
            "All requests failed. This is likely due to a misconfiguration on the benchmark arguments.",
            stacklevel=2,
@@ -296,7 +310,11 @@ def calculate_metrics(
        std_audio_ttfp_ms=np.std(audio_ttfps or 0) * 1000,
        median_audio_ttfp_ms=np.median(audio_ttfps or 0) * 1000,
        percentiles_audio_ttfp_ms=[(p, np.percentile(audio_ttfps or 0, p) * 1000) for p in selected_percentiles],
-        total_audio_duration_ms=sum(audio_duration),
+        mean_audio_duration_s=np.mean(audio_duration or 0),
+        std_audio_duration_s=np.std(audio_duration or 0),
+        median_audio_duration_s=np.median(audio_duration or 0),
+        percentiles_audio_duration_s=[(p, np.percentile(audio_duration or 0, p)) for p in selected_percentiles],
+        total_audio_duration_s=sum(audio_duration),
        total_audio_frames=sum(audio_frames),
        audio_throughput=sum(audio_duration) / dur_s,
        mean_audio_rtf=np.mean(audio_rtfs or 0),
@@ -113,8 +113,8 @@ async def async_request_openai_chat_omni_completions(
    st = time.perf_counter()
    output.start_time = st
    most_recent_timestamp = st
+    timestamp = st
    audio_generate_time = 0.0
-    audio_first_timestamp = st
    try:
        async with session.post(url=api_url, json=payload, headers=headers) as response:
            if response.status == 200:
@@ -133,7 +133,6 @@ async def async_request_openai_chat_omni_completions(
                            continue

                        chunk = message.removeprefix("data: ")
-
                        if chunk != "[DONE]":
                            timestamp = time.perf_counter()
                            data = json.loads(chunk)
@@ -148,11 +147,11 @@ async def async_request_openai_chat_omni_completions(
                                    else:
                                        output.itl.append(timestamp - most_recent_timestamp)
                                    generated_text += content or ""
+                                    most_recent_timestamp = timestamp
                                elif modality == "audio":
                                    if output.audio_ttfp == 0.0:
-                                        audio_first_timestamp = timestamp
                                        output.audio_ttfp = timestamp - st
-                                    audio_generate_time = timestamp - audio_first_timestamp
+                                    audio_generate_time = timestamp - st
                                    if content != "":
                                        audio_bytes = base64.b64decode(content)
                                        seg = AudioSegment.from_file(io.BytesIO(audio_bytes))
@@ -164,8 +163,8 @@ async def async_request_openai_chat_omni_completions(

                            elif usage := data.get("usage"):
                                output.output_tokens = usage.get("completion_tokens")
-                            most_recent_timestamp = timestamp

+                output.latency = timestamp - st
                output.generated_text = generated_text
                if generated_audio is not None:
                    output.audio_duration = len(generated_audio) / 1000.0
@@ -181,9 +180,7 @@ async def async_request_openai_chat_omni_completions(
                    else:
                        output.audio_rtf = 0
                        logger.warning("Audio duration is zero")
-
                output.success = True
-                output.latency = most_recent_timestamp - st
            else:
                output.error = response.reason or ""
                output.success = False
@@ -473,6 +470,9 @@ async def benchmark(
            "request_goodput": metrics.request_goodput if goodput_config_dict else None,
            "output_throughput": metrics.output_throughput,
            "total_token_throughput": metrics.total_token_throughput,
+            "total_audio_duration_s": metrics.total_audio_duration_s,
+            "total_audio_frames": metrics.total_audio_frames,
+            "audio_throughput": metrics.audio_throughput,
            "input_lens": [output.prompt_len for output in outputs],
            "output_lens": actual_output_lens,
            "ttfts": [output.ttft for output in outputs],
@@ -505,8 +505,20 @@ async def benchmark(
        if metric_attribute_name not in selected_percentile_metrics:
            return
        is_audio_rtf = metric_attribute_name == "audio_rtf"
+        is_audio_duration = metric_attribute_name == "audio_duration"

-        suffix = "" if is_audio_rtf else "_ms"
+        suffix = "_ms"
+        if is_audio_duration:
+            suffix = "_s"
+        elif is_audio_rtf:
+            suffix = ""
+        mean_attr_name = f"mean_{metric_attribute_name}{suffix}"
+        mean_value = getattr(metrics, mean_attr_name, 0.0)
+        result[mean_attr_name] = mean_value
+
+        median_attr_name = f"median_{metric_attribute_name}{suffix}"
+        median_value = getattr(metrics, median_attr_name, 0.0)
+        result[median_attr_name] = median_value
        for p, value in getattr(metrics, f"percentiles_{metric_attribute_name}{suffix}"):
            p_word = str(int(p)) if int(p) == p else str(p)
            result[f"p{p_word}_{metric_attribute_name}{suffix}"] = value
@@ -67,9 +67,11 @@ class OmniGenerationScheduler(VLLMScheduler):
            if self.omni_connector is not None:
                get_chunk_for_generation(self.omni_connector, request)

-            # OMNI: Skip requests that are already finished or not in self.requests
-            # This can happen when connector marks request as finished
-            if request.status == RequestStatus.FINISHED_STOPPED or request.request_id not in self.requests:
+            # OMNI: Skip requests that are not in self.requests
+            # This can happen when connector marks request as finished and it's removed from requests
+            if request.request_id not in self.requests or (
+                self.omni_connector is None and request.status == RequestStatus.FINISHED_STOPPED
+            ):
                already_finished_reqs.add(request)
                req_index += 1
                continue
@@ -108,9 +110,10 @@ class OmniGenerationScheduler(VLLMScheduler):
            if self.omni_connector is not None:
                get_chunk_for_generation(self.omni_connector, request)

-            # OMNI: Skip requests that are already finished or not in self.requests
-            # This can happen when connector marks request as finished
-            if request.status == RequestStatus.FINISHED_STOPPED or request.request_id not in self.requests:
+            # OMNI: Skip requests that are not in self.requests
+            if request.request_id not in self.requests or (
+                self.omni_connector is None and request.status == RequestStatus.FINISHED_STOPPED
+            ):
                # Pop the finished request from waiting queue and don't schedule it
                self.waiting.pop_request()
                continue
@@ -273,19 +273,77 @@ def enable_cache_for_longcat_image(pipeline: Any, cache_config: Any) -> Callable


 def enable_cache_for_flux(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for Flux dual-transformer architecture.
-
-    Flux uses two transformers (transformer and transformer_2) that need
-    to be enabled together using BlockAdapter.
+    """Enable cache-dit for Flux.1-dev pipeline.

    Args:
        pipeline: The Flux pipeline instance.
        cache_config: DiffusionCacheConfig instance with cache configuration.
-
    Returns:
-        A refresh function that can be called to update cache context with new num_inference_steps.
+        A refresh function that can be called with a new ``num_inference_steps``
+        to update the cache context for the pipeline.
    """
-    raise NotImplementedError("cache-dit is not implemented for Flux pipeline.")
+    # Build DBCacheConfig for transformer
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    # Build ParamsModifier for transformer
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator,
+    )
+
+    logger.info(
+        f"Enabling cache-dit on Flux transformer with BlockAdapter: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+
+    # Enable cache-dit using BlockAdapter for transformer
+    cache_dit.enable_cache(
+        (
+            BlockAdapter(
+                transformer=pipeline.transformer,
+                blocks=[
+                    pipeline.transformer.transformer_blocks,
+                    pipeline.transformer.single_transformer_blocks,
+                ],
+                forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+                params_modifiers=[modifier],
+            )
+        ),
+        cache_config=db_cache_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the transformer with new num_inference_steps.
+
+        Args:
+            pipeline: The Flux pipeline instance.
+            num_inference_steps: New number of inference steps.
+        """
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context


 def enable_cache_for_sd3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
@@ -253,7 +253,7 @@ class DiffusionEngine:

        try:
            # Give worker enough time — export + compression + table can be slow
-            results = self.collective_rpc(method="stop_profile", timeout=60000)
+            results = self.collective_rpc(method="stop_profile", timeout=600)
        except Exception:
            logger.error("Failed to stop profiling on workers", exc_info=True)
            return {"traces": [], "tables": []}
@@ -503,10 +503,50 @@ def init_vae_group(


 # adapted from https://github.com/feifeibear/long-context-attention/blob/main/yunchang/globals.py
-def set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size, use_ulysses_low=True):
+def set_seq_parallel_pg(
+    sp_ulysses_degree: int,
+    sp_ring_degree: int,
+    rank: int,
+    world_size: int,
+    use_ulysses_low: bool = True,
+    sp_group_ranks: list[list[int]] | None = None,
+) -> tuple[torch.distributed.ProcessGroup, torch.distributed.ProcessGroup]:
    """
-    sp_ulysses_degree x sp_ring_degree = seq_parallel_size
-    (ulysses_degree, dp_size)
+    Initialize sequence-parallel Ulysses and Ring process groups.
+
+    This builds sequence-parallel (SP) subgroups inside each data-parallel (DP)
+    slice. The SP group size is sp_ulysses_degree * sp_ring_degree, and
+    world_size must be divisible by that size.
+
+    Args:
+        sp_ulysses_degree: Size of each Ulysses subgroup.
+        sp_ring_degree: Size of each Ring subgroup.
+        rank: Global rank of the current process.
+        world_size: Total number of processes.
+        use_ulysses_low: If True, Ulysses groups are contiguous chunks and Ring
+            groups are strided within each SP group. If False, the opposite.
+        sp_group_ranks: Optional explicit SP groups. Each entry must be a list
+            of length sp_ulysses_degree * sp_ring_degree. When provided, groups
+            are built from these ranks instead of auto-generated contiguous
+            ranges.
+
+    Returns:
+        ulyssess_pg (torch.distributed.ProcessGroup): The Ulysses process group
+            for this rank.
+        ring_pg (torch.distributed.ProcessGroup): The Ring process group for
+            this rank.
+
+    Raises:
+        ValueError: If sp_group_ranks length does not match world_size or any
+            entry has the wrong size.
+        AssertionError: If world_size is not divisible by sp_size.
+
+    Behavior:
+        - If sp_group_ranks is provided, groups are built per entry and each
+          entry is further split into Ulysses/Ring groups according to
+          use_ulysses_low.
+        - If sp_group_ranks is None, groups are auto-generated within each DP
+          slice using offsets of size sp_size.
    """
    sp_size = sp_ring_degree * sp_ulysses_degree
    dp_size = world_size // sp_size
@@ -516,40 +556,96 @@ def set_seq_parallel_pg(sp_ulysses_degree, sp_ring_degree, rank, world_size, use
    num_ulysses_pgs = sp_ring_degree  # world_size // sp_ulysses_degree
    num_ring_pgs = sp_ulysses_degree  # world_size // sp_ring_degree

-    if use_ulysses_low:
-        for dp_rank in range(dp_size):
-            offset = dp_rank * sp_size
-            for i in range(num_ulysses_pgs):
-                ulysses_ranks = list(
-                    range(
-                        i * sp_ulysses_degree + offset,
-                        (i + 1) * sp_ulysses_degree + offset,
-                    )
-                )
-                group = torch.distributed.new_group(ulysses_ranks)
-                if rank in ulysses_ranks:
-                    ulyssess_pg = group
-
-            for i in range(num_ring_pgs):
-                ring_ranks = list(range(i + offset, sp_size + offset, num_ring_pgs))
-                group = torch.distributed.new_group(ring_ranks)
-                if rank in ring_ranks:
-                    ring_pg = group
-
+    if sp_group_ranks is not None:
+        if len(sp_group_ranks) * sp_size != world_size:
+            raise ValueError(
+                f"Invalid sp_group_ranks: expected {world_size // sp_size} groups of size {sp_size}, "
+                f"but got {len(sp_group_ranks)} groups."
+            )
+        logger.info(
+            "Building SP subgroups from explicit sp_group_ranks "
+            f"(sp_size={sp_size}, ulysses={sp_ulysses_degree}, ring={sp_ring_degree}, "
+            f"use_ulysses_low={use_ulysses_low})."
+        )
+        local_sp_group = None
+        local_ulysses = None
+        local_ring = None
+        for group_ranks in sp_group_ranks:
+            if len(group_ranks) != sp_size:
+                raise ValueError(f"Invalid sp_group_ranks entry: expected size {sp_size}, got {len(group_ranks)}.")
+            if rank in group_ranks:
+                local_sp_group = list(group_ranks)
+            if use_ulysses_low:
+                # Ulysses groups are contiguous chunks; Ring groups are strided.
+                for i in range(num_ulysses_pgs):
+                    ulysses_ranks = group_ranks[i * sp_ulysses_degree : (i + 1) * sp_ulysses_degree]
+                    group = torch.distributed.new_group(ulysses_ranks)
+                    if rank in ulysses_ranks:
+                        ulyssess_pg = group
+                        local_ulysses = list(ulysses_ranks)
+                for i in range(num_ring_pgs):
+                    ring_ranks = group_ranks[i::num_ring_pgs]
+                    group = torch.distributed.new_group(ring_ranks)
+                    if rank in ring_ranks:
+                        ring_pg = group
+                        local_ring = list(ring_ranks)
+            else:
+                # Ring groups are contiguous chunks; Ulysses groups are strided.
+                for i in range(num_ring_pgs):
+                    ring_ranks = group_ranks[i * sp_ring_degree : (i + 1) * sp_ring_degree]
+                    group = torch.distributed.new_group(ring_ranks)
+                    if rank in ring_ranks:
+                        ring_pg = group
+                        local_ring = list(ring_ranks)
+                for i in range(num_ulysses_pgs):
+                    ulysses_ranks = group_ranks[i::num_ulysses_pgs]
+                    group = torch.distributed.new_group(ulysses_ranks)
+                    if rank in ulysses_ranks:
+                        ulyssess_pg = group
+                        local_ulysses = list(ulysses_ranks)
+        if local_sp_group is not None:
+            logger.info(
+                "SP group details for rank %d: sp_group=%s, ulysses_group=%s, ring_group=%s",
+                rank,
+                local_sp_group,
+                local_ulysses,
+                local_ring,
+            )
    else:
-        for dp_rank in range(dp_size):
-            offset = dp_rank * sp_size
-            for i in range(num_ring_pgs):
-                ring_ranks = list(range(i * sp_ring_degree + offset, (i + 1) * sp_ring_degree + offset))
-                group = torch.distributed.new_group(ring_ranks)
-                if rank in ring_ranks:
-                    ring_pg = group
+        if use_ulysses_low:
+            for dp_rank in range(dp_size):
+                offset = dp_rank * sp_size
+                for i in range(num_ulysses_pgs):
+                    ulysses_ranks = list(
+                        range(
+                            i * sp_ulysses_degree + offset,
+                            (i + 1) * sp_ulysses_degree + offset,
+                        )
+                    )
+                    group = torch.distributed.new_group(ulysses_ranks)
+                    if rank in ulysses_ranks:
+                        ulyssess_pg = group

-            for i in range(num_ulysses_pgs):
-                ulysses_ranks = list(range(i + offset, sp_size + offset, num_ulysses_pgs))
-                group = torch.distributed.new_group(ulysses_ranks)
-                if rank in ulysses_ranks:
-                    ulyssess_pg = group
+                for i in range(num_ring_pgs):
+                    ring_ranks = list(range(i + offset, sp_size + offset, num_ring_pgs))
+                    group = torch.distributed.new_group(ring_ranks)
+                    if rank in ring_ranks:
+                        ring_pg = group
+
+        else:
+            for dp_rank in range(dp_size):
+                offset = dp_rank * sp_size
+                for i in range(num_ring_pgs):
+                    ring_ranks = list(range(i * sp_ring_degree + offset, (i + 1) * sp_ring_degree + offset))
+                    group = torch.distributed.new_group(ring_ranks)
+                    if rank in ring_ranks:
+                        ring_pg = group
+
+                for i in range(num_ulysses_pgs):
+                    ulysses_ranks = list(range(i + offset, sp_size + offset, num_ulysses_pgs))
+                    group = torch.distributed.new_group(ulysses_ranks)
+                    if rank in ulysses_ranks:
+                        ulyssess_pg = group

    return ulyssess_pg, ring_pg

@@ -652,6 +748,7 @@ def initialize_model_parallel(
        data_parallel_size,
        "tp-sp-pp-cfg-dp",
    )
+    sp_group_ranks = rank_generator.get_ranks("sp")
    global _DP
    assert _DP is None, "data parallel group is already initialized"
    _DP = init_model_parallel_group(
@@ -685,9 +782,10 @@ def initialize_model_parallel(
        sp_ring_degree=ring_degree,
        rank=get_world_group().rank_in_group,
        world_size=dit_parallel_size,
+        sp_group_ranks=sp_group_ranks,
    )
    _SP = init_model_parallel_group(
-        group_ranks=rank_generator.get_ranks("sp"),
+        group_ranks=sp_group_ranks,
        local_rank=get_world_group().local_rank,
        backend=backend,
        parallel_mode="sequence",
@@ -72,11 +72,15 @@ def get_longcat_image_edit_pre_process_func(
            if "additional_information" not in prompt:
                prompt["additional_information"] = {}

-            if raw_image is None or isinstance(raw_image, list):
-                raise ValueError(
-                    """Received no image or a list of image. Only a single image is supported by this model."""
-                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
-                )
+            if not raw_image:  # None or empty list
+                raise ValueError("""Received no input image. This model requires one input image to run.""")
+            elif isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    raise ValueError(
+                        """Received multiple input images. Only a single image is supported by this model."""
+                    )
+                else:
+                    raw_image = raw_image[0]

            if isinstance(raw_image, str):
                image = PIL.Image.open(raw_image)
@@ -75,11 +75,15 @@ def get_qwen_image_edit_pre_process_func(
                prompt["additional_information"] = {}

            # Only handles single image
-            if raw_image is None or isinstance(raw_image, list):
-                raise ValueError(
-                    """Received no image or a list of image. Only a single image is supported by this model."""
-                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
-                )
+            if not raw_image:  # None or empty list
+                raise ValueError("""Received no input image. This model requires one input image to run.""")
+            elif isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    raise ValueError(
+                        """Received multiple input images. Only a single image is supported by this model."""
+                    )
+                else:
+                    raw_image = raw_image[0]

            if isinstance(raw_image, str):
                image = PIL.Image.open(raw_image)
@@ -628,6 +632,13 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
        first_prompt = req.prompts[0]
        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
        negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
+        if negative_prompt is None:
+            logger.warning(
+                "negative_prompt is not set. The official Qwen-Image-Edit model "
+                "may produce lower-quality results without a negative_prompt. "
+                "Qwen official repository recommends to use whitespace string as negative_prompt. "
+                "Note: some distilled variants may not be affected by this."
+            )

        # Get preprocessed image from request (pre-processing is done in DiffusionEngine)
        if not isinstance(first_prompt, str) and "preprocessed_image" in (
@@ -563,6 +563,13 @@ class QwenImageEditPlusPipeline(nn.Module, SupportImageInput, QwenImageCFGParall
        first_prompt = req.prompts[0]
        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
        negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
+        if negative_prompt is None:
+            logger.warning(
+                "negative_prompt is not set. The official Qwen-Image-Edit model "
+                "may produce lower-quality results without a negative_prompt. "
+                "Qwen official repository recommends to use whitespace string as negative_prompt. "
+                "Note: some distilled variants may not be affected by this."
+            )

        # Get preprocessed images from request (pre-processing is done in DiffusionEngine)
        if (
@@ -79,11 +79,15 @@ def get_qwen_image_layered_pre_process_func(
            if "additional_information" not in prompt:
                prompt["additional_information"] = {}

-            if raw_image is None or isinstance(raw_image, list):
-                raise ValueError(
-                    """Received no image or a list of image. Only a single image is supported by this model."""
-                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
-                )
+            if not raw_image:  # None or empty list
+                raise ValueError("""Received no input image. This model requires one input image to run.""")
+            elif isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    raise ValueError(
+                        """Received multiple input images. Only a single image is supported by this model."""
+                    )
+                else:
+                    raw_image = raw_image[0]

            if isinstance(raw_image, str):
                image = PIL.Image.open(raw_image)
@@ -1,8 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-import functools
 from collections.abc import Iterable
+from functools import lru_cache
 from math import prod
 from typing import Any

@@ -249,7 +249,7 @@ class QwenEmbedLayer3DRope(nn.Module):

        return vid_freqs, txt_freqs

-    @functools.cache
+    @lru_cache(maxsize=16)
    def _compute_video_freqs(self, frame, height, width, idx=0):
        seq_lens = frame * height * width
        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
@@ -268,7 +268,7 @@ class QwenEmbedLayer3DRope(nn.Module):
        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1)
        return freqs.clone().contiguous()

-    @functools.cache
+    @lru_cache(maxsize=16)
    def _compute_condition_freqs(self, frame, height, width):
        seq_lens = frame * height * width
        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
@@ -311,7 +311,6 @@ class QwenEmbedRope(nn.Module):
            ],
            dim=1,
        )
-        self.rope_cache = {}

        # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
        self.scale_rope = scale_rope
@@ -349,14 +348,7 @@ class QwenEmbedRope(nn.Module):
        max_vid_index = 0
        for idx, fhw in enumerate(video_fhw):
            frame, height, width = fhw
-            rope_key = f"{idx}_{height}_{width}"
-
-            if not torch.compiler.is_compiling():
-                if rope_key not in self.rope_cache:
-                    self.rope_cache[rope_key] = self._compute_video_freqs(frame, height, width, idx)
-                video_freq = self.rope_cache[rope_key]
-            else:
-                video_freq = self._compute_video_freqs(frame, height, width, idx)
+            video_freq = self._compute_video_freqs(frame, height, width, idx)
            video_freq = video_freq.to(device)
            vid_freqs.append(video_freq)

@@ -371,7 +363,7 @@ class QwenEmbedRope(nn.Module):

        return vid_freqs, txt_freqs

-    @functools.cache
+    @lru_cache(maxsize=16)
    def _compute_video_freqs(self, frame, height, width, idx=0):
        seq_lens = frame * height * width
        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
@@ -527,6 +527,16 @@ class ZImagePipeline(nn.Module):
        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
        self._num_timesteps = len(timesteps)

+        # Precompute normalized timesteps once to avoid per-step GPU->CPU sync (.item() causes cudaStreamSynchronize)
+        if isinstance(timesteps, torch.Tensor):
+            timesteps_tensor = timesteps.to(device=device, dtype=torch.float32)
+        else:
+            timesteps_tensor = torch.as_tensor(timesteps, device=device, dtype=torch.float32)
+        norm_timesteps = (1000 - timesteps_tensor) / 1000
+        t_norm_list = norm_timesteps.cpu().tolist()
+        if not isinstance(t_norm_list, list):
+            t_norm_list = [t_norm_list]
+
        # 6. Denoising loop
        for i, t in enumerate(timesteps):
            if self.interrupt:
@@ -535,8 +545,9 @@ class ZImagePipeline(nn.Module):
            # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
            timestep = t.expand(latents.shape[0])
            timestep = (1000 - timestep) / 1000
-            # Normalized time for time-aware config (0 at start, 1 at end)
-            t_norm = timestep[0].item()
+            # Normalized time for time-aware config (0 at start, 1 at end);
+            # use precomputed to avoid .item() sync per step
+            t_norm = t_norm_list[i]

            # Handle cfg truncation
            current_guidance_scale = self.guidance_scale
@@ -582,13 +593,17 @@ class ZImagePipeline(nn.Module):

                    pred = pos + current_guidance_scale * (pos - neg)

-                    # Renormalization
+                    # Renormalization (torch.where avoids GPU->CPU sync from Python if/scalar comparison)
                    if self._cfg_normalization and float(self._cfg_normalization) > 0.0:
                        ori_pos_norm = torch.linalg.vector_norm(pos)
                        new_pos_norm = torch.linalg.vector_norm(pred)
                        max_new_norm = ori_pos_norm * float(self._cfg_normalization)
-                        if new_pos_norm > max_new_norm:
-                            pred = pred * (max_new_norm / new_pos_norm)
+                        scale = torch.where(
+                            new_pos_norm > max_new_norm,
+                            (max_new_norm / new_pos_norm.clamp(min=1e-12)).to(pred.dtype),
+                            pred.new_tensor(1.0),
+                        )
+                        pred = pred * scale

                    noise_pred.append(pred)

@@ -852,7 +852,13 @@ class ZImageTransformer2DModel(CachedTransformer):

        # Match t_embedder output dtype to x for layerwise casting compatibility
        adaln_input = t.type_as(x)
-        x[torch.cat(x_inner_pad_mask)] = self.x_pad_token
+        # Use torch.where instead of x[mask]= to avoid aten::index_put_/nonzero and cudaStreamSynchronize
+        x_pad_mask = torch.cat(x_inner_pad_mask)
+        x = torch.where(
+            x_pad_mask.unsqueeze(1).expand_as(x),
+            self.x_pad_token.expand(x.shape[0], -1),
+            x,
+        )
        x = list(x.split(x_item_seqlens, dim=0))
        x_cos, x_sin = self.rope_embedder(torch.cat(x_pos_ids, dim=0))
        x_cos = list(x_cos.split(x_item_seqlens, dim=0))
@@ -875,7 +881,13 @@ class ZImageTransformer2DModel(CachedTransformer):

        cap_feats = torch.cat(cap_feats, dim=0)
        cap_feats = self.cap_embedder(cap_feats)
-        cap_feats[torch.cat(cap_inner_pad_mask)] = self.cap_pad_token
+        # Use torch.where instead of cap_feats[mask]= to avoid aten::index_put_/nonzero and cudaStreamSynchronize
+        cap_pad_mask = torch.cat(cap_inner_pad_mask)
+        cap_feats = torch.where(
+            cap_pad_mask.unsqueeze(1).expand_as(cap_feats),
+            self.cap_pad_token.expand(cap_feats.shape[0], -1),
+            cap_feats,
+        )
        cap_feats = list(cap_feats.split(cap_item_seqlens, dim=0))
        cap_cos, cap_sin = self.rope_embedder(torch.cat(cap_pos_ids, dim=0))
        cap_cos = list(cap_cos.split(cap_item_seqlens, dim=0))
@@ -13,14 +13,13 @@ logger = init_logger(__name__)

 class Scheduler:
    def initialize(self, od_config: OmniDiffusionConfig):
-        existing_context = getattr(self, "context", None)
-        if existing_context is not None and not existing_context.closed:
+        existing_mq = getattr(self, "mq", None)
+        if existing_mq is not None and not existing_mq.closed:
            logger.warning("SyncSchedulerClient is already initialized. Re-initializing.")
            self.close()

        self.num_workers = od_config.num_gpus
        self.od_config = od_config
-        self.context = zmq.Context()  # Standard synchronous context

        # Initialize single MessageQueue for all message types (generation & RPC)
        # Assuming all readers are local for now as per current launch_engine implementation
@@ -72,8 +71,5 @@ class Scheduler:

    def close(self):
        """Closes the socket and terminates the context."""
-        if hasattr(self, "context"):
-            self.context.term()
-        self.context = None
        self.mq = None
        self.result_mq = None
@@ -179,6 +179,27 @@ def try_recv_via_connector(
            return None, None


+def update_request_payload(connector: "OmniConnectorBase", req_id: str, payload_data: dict[str, Any]) -> dict[str, Any]:
+    """Update the payload data for a request in the connector.
+
+    Args:
+        connector: OmniConnectorBase instance
+        req_id: Request ID to update
+        payload_data: New payload data to store
+    """
+    origin_payload = connector.request_payload[req_id]
+    for key, value in payload_data.items():
+        if key == "finished":
+            continue
+        elif isinstance(value, torch.Tensor) and key in origin_payload:
+            payload_data[key] = torch.cat([origin_payload[key], value], dim=0)
+        elif isinstance(value, list) and key in origin_payload:
+            payload_data[key] = origin_payload[key] + value
+
+    connector.request_payload[req_id] = payload_data
+    return payload_data
+
+
 def get_chunk(
    connector: "OmniConnectorBase",
    scheduler_output: SchedulerOutput,
@@ -206,6 +227,7 @@ def get_chunk(
        payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key)
        if payload_data:
            new_req_data.additional_information = payload_data
+            connector.request_payload[req_id] = payload_data
            if payload_data.get("finished"):
                connector.finished_requests.add(req_id)

@@ -222,15 +244,13 @@ def get_chunk(
        connector_get_key = f"{req_id}_{target_stage_id}_{chunk_id}"
        payload_data = get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key)
        if payload_data:
+            payload_data = update_request_payload(connector, req_id, payload_data)
            cached_reqs.additional_information[cached_req_id] = payload_data
            if payload_data.get("finished"):
                connector.finished_requests.add(req_id)


 def get_through_connector(connector, target_stage_id, stage_id, req_id, connector_get_key):
-    # Wait for data from previous stage
-    import time
-
    # TODO: add correct check mechanism for the payload_data
    max_wait = 300
    for _ in range(max_wait):
@@ -280,12 +300,8 @@ def get_chunk_for_generation(
    if payload_data.get("finished"):
        connector.finished_requests.add(request_id)
        request.status = RequestStatus.FINISHED_STOPPED
-
-    # TODO: remove special handling for prompt token ids ?
-    if chunk_id == 0:
-        request.prompt_token_ids = payload_data.get("code_predictor_codes", [])
-    else:
-        request.prompt_token_ids += payload_data.get("code_predictor_codes", [])
+    request.prompt_token_ids = payload_data.get("code_predictor_codes", [])
+    request.num_computed_tokens = 0


 def put_chunk(
@@ -308,8 +324,6 @@ def put_chunk(
    stage_id = connector.stage_id
    next_stage_id = stage_id + 1
    request_id = request.external_req_id
-    prompt_token_ids = request.prompt_token_ids
-    connector.request_prompt_token_ids[request_id] = prompt_token_ids
    chunk_id = connector.put_requests[request_id]
    connector_put_key = f"{request_id}_{stage_id}_{chunk_id}"
    payload_data = None
@@ -318,6 +332,7 @@ def put_chunk(
    if custom_process_input_func:
        try:
            payload_data = custom_process_input_func(
+                connector=connector,
                pooling_output=pooling_output,
                request=request,
            )
@@ -325,42 +340,8 @@ def put_chunk(
            logger.error(f"Failed to use custom_process_input_func for payload extraction: {e}")

        if not payload_data:
-            logger.warning("[Stage-%d] No payload data to send for request %s", stage_id, request_id)
            return

-        if stage_id == 0 and chunk_id == 0:
-            if connector.request_payload.get(request_id) is None:
-                if not payload_data.get("finished"):
-                    connector.request_payload[request_id] = payload_data
-                    return
-            else:
-                save_payload = connector.request_payload.pop(request_id)
-                payload_data["thinker_embeddings"] = torch.cat(
-                    (save_payload.get("thinker_embeddings"), payload_data.get("thinker_embeddings")), dim=0
-                )
-                payload_data["thinker_hidden_states"] = torch.cat(
-                    (save_payload.get("thinker_hidden_states"), payload_data.get("thinker_hidden_states")), dim=0
-                )
-                logger.debug("[Stage-%d] Merged embeddings and hidden states for request %s", stage_id, request_id)
-
-        if stage_id == 1:
-            # TODO: Make parameters configurable and optimize algorithms
-            chunk_size = left_context_size = 25
-            connector.code_prompt_token_ids[request_id].append(payload_data.get("code_predictor_codes", []))
-            length = len(connector.code_prompt_token_ids[request_id])
-            chunk_length = length % chunk_size
-            if chunk_length != 0 and not payload_data.get("finished"):
-                return
-
-            context_length = chunk_length if chunk_length != 0 else chunk_size
-            end_index = min(length, left_context_size + context_length)
-            payload_data["code_predictor_codes"] = (
-                torch.tensor(connector.code_prompt_token_ids[request_id][-end_index:])
-                .transpose(0, 1)
-                .reshape(-1)
-                .tolist()
-            )
-
        success, size, metadata = connector.put(
            from_stage=str(stage_id), to_stage=str(next_stage_id), put_key=connector_put_key, data=payload_data
        )
@@ -345,14 +345,22 @@ class AsyncOmni(OmniBase):
                ):
                    yield output

-            logger.debug(f"[{self._name}] All requests completed")
-
-            # Summarize and print stats
+            logger.debug(f"[{self._name}] Request {request_id} finalized at stage-{final_stage_id_for_e2e}")
            try:
+                # Finalize E2E metrics if not already done
+                if str(request_id) not in metrics.e2e_done:
+                    metrics.on_finalize_request(
+                        final_stage_id_for_e2e,
+                        request_id,
+                        _req_start_ts.get(request_id, _wall_start_ts),
+                    )
+
+                logger.debug(f"[{self._name}] All requests completed")
+                # Summarize and print stats
                summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
                logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
            except Exception as e:
-                logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
+                logger.exception(f"[{self._name}] Request {request_id} Failed to finalized/build/log summary: {e}")
            finally:
                self.request_states.pop(request_id, None)
        except (asyncio.CancelledError, GeneratorExit):
@@ -382,9 +390,11 @@ class AsyncOmni(OmniBase):
                except asyncio.QueueEmpty:
                    await asyncio.sleep(0.001)
                    continue
-
                engine_outputs, finished, output_to_yield = self._process_single_result(
-                    result, stage, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e
+                    result,
+                    stage,
+                    stage_id,
+                    metrics,
                )
                if submit_flag and stage_id == 0:
                    submit_flag = False
@@ -421,9 +431,11 @@ class AsyncOmni(OmniBase):
            while not finished:
                result = await req_state.queue.get()
                assert stage_id == req_state.stage_id
-                req_id = result.get("request_id")
                engine_outputs, finished, output_to_yield = self._process_single_result(
-                    result, stage, stage_id, metrics, req_start_ts, wall_start_ts, final_stage_id_for_e2e
+                    result,
+                    stage,
+                    stage_id,
+                    metrics,
                )
                if output_to_yield:
                    yield output_to_yield
@@ -432,7 +444,7 @@ class AsyncOmni(OmniBase):
            stage.set_engine_outputs(engine_outputs)
            # Forward to next stage if there is one
            next_stage_id = stage_id + 1
-            if next_stage_id <= final_stage_id_for_e2e and finished:
+            if next_stage_id <= final_stage_id_for_e2e:
                next_stage: OmniStage = self.stage_list[next_stage_id]
                next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
                sp_next: SamplingParams = sampling_params_list[next_stage_id]
@@ -447,7 +459,7 @@ class AsyncOmni(OmniBase):
                        connector=connector,
                        stage_id=stage_id,
                        next_stage_id=next_stage_id,
-                        req_id=req_id,
+                        req_id=request_id,
                        next_inputs=next_inputs,
                        sampling_params=sp_next,
                        original_prompt=prompt,
@@ -461,14 +473,14 @@ class AsyncOmni(OmniBase):
                    # because continuing would cause the request to be silently dropped
                    # and the orchestrator to hang waiting for completion.
                    error_msg = (
-                        f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
+                        f"[{self._name}] Failed to send request {request_id} to stage-{next_stage_id} via connector. "
                        "Configure a connector for this edge or inspect connector logs for details."
                    )
                    logger.error(error_msg)
                    raise RuntimeError(error_msg)
-                logger.debug(f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}")
+                logger.debug(f"[{self._name}] Forwarded request {request_id} to stage-{next_stage_id}")
            else:
-                logger.debug(f"[{self._name}] Request {req_id} fully completed")
+                logger.debug(f"[{self._name}] Request {request_id} fully completed")

    def _process_single_result(
        self,
@@ -476,9 +488,6 @@ class AsyncOmni(OmniBase):
        stage: OmniStage,
        stage_id: int,
        metrics: OrchestratorMetrics,
-        req_start_ts: dict[int, float],
-        wall_start_ts: float,
-        final_stage_id_for_e2e: int,
    ) -> tuple[Any, bool, OmniRequestOutput | None]:
        """
        Process a single result dictionary from a stage.
@@ -517,23 +526,8 @@ class AsyncOmni(OmniBase):
        )

        output_to_yield = None
+
        if getattr(stage, "final_output", False):
-            logger.debug(f"[{self._name}] Request {req_id} finalized at stage-{stage_id}")
-
-            # Finalize request metrics if this is the E2E final stage and it's finished
-            try:
-                rid_key = str(req_id)
-                if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done and finished:
-                    metrics.on_finalize_request(
-                        stage_id,
-                        req_id,
-                        req_start_ts.get(req_id, wall_start_ts),
-                    )
-            except Exception as e:
-                logger.exception(
-                    f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}",
-                )
-
            # Construct output to yield
            images = []
            if stage.final_output_type == "image":
@@ -97,7 +97,7 @@ class OmniServeCommand(CLISubcommand):
        omni_config_group.add_argument(
            "--init-timeout",
            type=int,
-            default=60000,
+            default=600,
            help="The timeout for initializing the stages.",
        )
        omni_config_group.add_argument(
@@ -221,8 +221,8 @@ class OmniStage:

        # Wait for result from worker
        try:
-            # Profiling stop might take time to flush files, give it 180s
-            response = self._out_q.get(timeout=60000)
+            # Profiling stop might take time to flush files, give it 600s
+            response = self._out_q.get(timeout=600)

            if isinstance(response, dict):
                if response.get("type") == "profiler_result":
@@ -599,11 +599,16 @@ class Qwen3OmniMoeForConditionalGeneration(
            )
            update_dict["code_predictor_codes"] = code_predictor_codes
        else:
+            # decode
+            if info_dict.get("num_processed_tokens", 0) < len(info_dict.get("thinker_input_ids", [])):
+                info_dict["num_processed_tokens"] = len(info_dict.get("thinker_input_ids", [])) + 1
+
            last_talker_hidden, text_step, update_dict = self.talker_preprocess_decode(
                input_ids, input_embeds, **info_dict
            )
            update_dict["mtp_inputs"] = last_talker_hidden, text_step

+        update_dict["num_processed_tokens"] = info_dict.get("num_processed_tokens", 0) + span_len
        return input_ids, input_embeds, update_dict

    def talker_mtp(
@@ -662,7 +667,8 @@ class Qwen3OmniMoeForConditionalGeneration(
        update_dict: dict[str, dict] = {}
        # TODO(Peiqi): add voice_type support
        voice_type = self.voice_type
-
+        start_index = info_dict.get("num_processed_tokens", 0)
+        end_index = start_index + input_embeds.shape[0]
        # Read thinker outputs for prefill
        thinker_sequence_embeds = info_dict.get("thinker_embeddings").to(
            device=self._module_device(self.talker), dtype=torch.bfloat16
@@ -762,7 +768,7 @@ class Qwen3OmniMoeForConditionalGeneration(
        except Exception:
            pass

-        return req_input_ids, req_embeds, update_dict
+        return req_input_ids[start_index:end_index], req_embeds[start_index:end_index], update_dict

    def _thinker_to_talker_prefill(
        self,
@@ -858,13 +864,14 @@ class Qwen3OmniMoeForConditionalGeneration(
            (input_ids, input_embeds) for talker
        """
        thinker_embed = info_dict.get("thinker_embeddings", None)
-        if thinker_embed is None:
+        start_index = info_dict.get("num_processed_tokens", 0)
+        if start_index >= thinker_embed.shape[0]:
            if info_dict.get("finished_flag"):
                return self.tts_pad_embed.to(device)
            update_dict["finished_flag"] = True
            return self.tts_eos_embed.to(device)

-        thinker_embed = thinker_embed.to(device)
+        thinker_embed = thinker_embed[start_index : start_index + 1].to(device)
        return self.talker.text_projection(thinker_embed).to(device)

    def talker_preprocess_decode(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
@@ -336,14 +336,19 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
                mm_item_counts,
            )
        else:
-            if use_audio_in_video and "audio" in mm_prompt_updates:
+            if use_audio_in_video:
+                # When use_audio_in_video=True, audio is extracted from video and embedded
+                # in the video placeholder tokens. We should:
+                # 1. Filter out audio from prompt updates (audio has no separate placeholder)
+                # 2. Apply remaining updates (video, image, etc.)
+                # 3. Derive audio placeholders from video placeholders
                filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"}
                prompt_ids, mm_placeholders = self._apply_prompt_updates(
                    prompt_ids,
                    filtered_updates,
                )
                # Derive audio placeholders from video placeholders
-                mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates)
+                mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_item_counts)
            else:
                prompt_ids, mm_placeholders = self._apply_prompt_updates(
                    prompt_ids,
@@ -508,18 +513,27 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
    def _derive_audio_from_video_placeholders(
        self,
        placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
-        mm_prompt_updates: MultiModalPromptUpdates,
+        mm_item_counts: Mapping[str, int],
    ) -> Mapping[str, list[PlaceholderFeaturesInfo]]:
        """
        Helper to derive audio placeholders from video placeholders when
        use_audio_in_video=True.
+
+        In use_audio_in_video mode, audio is extracted from video and embedded
+        within the video placeholder tokens. This function creates audio placeholder
+        info by extracting the audio token positions from video placeholders.
+
+        Args:
+            placeholders: Current placeholders (should contain "video")
+            mm_item_counts: Counts of multimodal items from mm_items.get_all_counts()
        """
        if "video" not in placeholders:
            return placeholders

        # Validate audio and video counts match
+        # In use_audio_in_video mode, audio count comes from mm_items (extracted from video)
        num_videos = len(placeholders["video"])
-        num_audios = len(mm_prompt_updates.get("audio", []))
+        num_audios = mm_item_counts.get("audio", 0)
        if num_audios != num_videos:
            raise ValueError(
                f"use_audio_in_video requires equal number of audio and video items, got {num_audios=}, {num_videos=}"
@@ -63,7 +63,7 @@ stage_args:
    default_sampling_params:
      temperature: 0.9
      top_k: 50
-      max_tokens: 2048 # TODO: The max_tokens of the async_chunk feature cannot exceed 2048.
+      max_tokens: 4096
      seed: 42
      detokenize: False
      repetition_penalty: 1.05
@@ -84,6 +84,7 @@ def _validate_stage_inputs(stage_list, engine_input_source):


 def thinker2talker_async_chunk(
+    connector: Any,
    pooling_output: dict[str, Any],
    request: OmniEngineCoreRequest,
 ) -> list[dict[str, Any]]:
@@ -93,27 +94,50 @@ def thinker2talker_async_chunk(
    2. Split hidden states into: prompt embeddings + generated embeddings
    3. Package for talker with additional information
    """
-    all_token_ids = request.all_token_ids  # prefill + decode
-    prompt_token_ids = request.prompt_token_ids

-    # Convert ConstantList to regular list for OmniSerializer serialization
-    all_token_ids = _ensure_list(all_token_ids)
-    prompt_token_ids = _ensure_list(prompt_token_ids)
-
-    thinker_output = pooling_output
-
-    talker_additional_info = {
-        "thinker_embeddings": thinker_output.get("0").detach().cpu(),
-        "thinker_hidden_states": thinker_output.get("24").detach().cpu(),
-        "thinker_sequences": all_token_ids,
-        "thinker_input_ids": prompt_token_ids,
-        # Provide thinker-side TTS token embeddings for talker projection
-        "tts_bos_embed": thinker_output.get("tts_bos_embed").detach().cpu(),
-        "tts_eos_embed": thinker_output.get("tts_eos_embed").detach().cpu(),
-        "tts_pad_embed": thinker_output.get("tts_pad_embed").detach().cpu(),
-        "finished": torch.tensor(request.is_finished(), dtype=torch.bool),
-    }
+    request_id = request.external_req_id
+    chunk_id = connector.put_requests[request_id]
+    if chunk_id == 0:
+        all_token_ids = request.all_token_ids  # prefill + decode
+        prompt_token_ids = request.prompt_token_ids
+        # Convert ConstantList to regular list for OmniSerializer serialization
+        all_token_ids = _ensure_list(all_token_ids)
+        prompt_token_ids = _ensure_list(prompt_token_ids)
+        talker_additional_info = {
+            "thinker_embeddings": pooling_output.get("0").detach().cpu(),
+            "thinker_hidden_states": pooling_output.get("24").detach().cpu(),
+            "thinker_sequences": all_token_ids,
+            "thinker_input_ids": prompt_token_ids,
+            # Provide thinker-side TTS token embeddings for talker projection
+            "tts_bos_embed": pooling_output.get("tts_bos_embed").detach().cpu(),
+            "tts_eos_embed": pooling_output.get("tts_eos_embed").detach().cpu(),
+            "tts_pad_embed": pooling_output.get("tts_pad_embed").detach().cpu(),
+            "finished": torch.tensor(request.is_finished(), dtype=torch.bool),
+        }
+        if connector.request_payload.get(request_id) is None:
+            if not request.is_finished():
+                connector.request_payload[request_id] = talker_additional_info
+                return None
+        else:
+            save_payload = connector.request_payload.pop(request_id)
+            talker_additional_info["thinker_embeddings"] = torch.cat(
+                (save_payload.get("thinker_embeddings"), talker_additional_info.get("thinker_embeddings")), dim=0
+            )
+            talker_additional_info["thinker_hidden_states"] = torch.cat(
+                (save_payload.get("thinker_hidden_states"), talker_additional_info.get("thinker_hidden_states")),
+                dim=0,
+            )
+    else:
+        output_token_ids = request.output_token_ids
+        # Convert ConstantList to regular list for OmniSerializer serialization
+        output_token_ids = _ensure_list(output_token_ids)

+        talker_additional_info = {
+            "thinker_embeddings": pooling_output.get("0").detach().cpu(),
+            "thinker_hidden_states": pooling_output.get("24").detach().cpu(),
+            "thinker_sequences": output_token_ids,
+            "finished": torch.tensor(request.is_finished(), dtype=torch.bool),
+        }
    return talker_additional_info


@@ -182,6 +206,7 @@ def thinker2talker(


 def talker2code2wav_async_chunk(
+    connector: Any,
    pooling_output: dict[str, Any],
    request: OmniEngineCoreRequest,
 ):
@@ -189,35 +214,49 @@ def talker2code2wav_async_chunk(
    Pooling version.
    """
    if "code_predictor_codes" not in pooling_output:
-        return []
+        return None

    code_predictor_codes = pooling_output["code_predictor_codes"]

    if code_predictor_codes is None:
-        return []
+        return None
    if isinstance(code_predictor_codes, torch.Tensor):
        if code_predictor_codes.numel() == 0:
-            return []
+            return None
    elif hasattr(code_predictor_codes, "__len__"):
        if len(code_predictor_codes) == 0:
-            return []
+            return None

    if isinstance(code_predictor_codes, torch.Tensor):
        if not code_predictor_codes.any():
-            return []
+            return None
    else:
        code_tensor = torch.tensor(code_predictor_codes, dtype=torch.long)
        if not code_tensor.any():
-            return []
+            return None

    codec_codes = code_predictor_codes.to(torch.long).transpose(0, 1).cpu().to(torch.long).reshape(-1).tolist()
    if sum(codec_codes) == 0:
-        return []
+        return None

-    return {
-        "code_predictor_codes": codec_codes,
+    request_id = request.external_req_id
+    chunk_size = left_context_size = 25
+    connector.code_prompt_token_ids[request_id].append(codec_codes)
+    length = len(connector.code_prompt_token_ids[request_id])
+    chunk_length = length % chunk_size
+    if chunk_length != 0 and not request.is_finished():
+        return None
+
+    context_length = chunk_length if chunk_length != 0 else chunk_size
+    end_index = min(length, left_context_size + context_length)
+
+    info = {
+        "code_predictor_codes": (
+            torch.tensor(connector.code_prompt_token_ids[request_id][-end_index:]).transpose(0, 1).reshape(-1).tolist()
+        ),
        "finished": torch.tensor(request.is_finished(), dtype=torch.bool),
    }
+    return info


 def talker2code2wav(
@@ -733,13 +733,10 @@ class OmniNPUModelRunner(NPUModelRunner):
            # TODO(Peiqi): do we have a more elegant way to do this?
            if hasattr(self.model, "has_postprocess") and self.model.has_postprocess:
                for req_index, req_id in enumerate(self.input_batch.req_ids):
-                    if self.model_config.async_chunk:
-                        req_infos = self._get_additional_information(scheduler_output, req_id)
-                    else:
-                        req_state = self.requests.get(req_id)
-                        req_infos = (
-                            getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
-                        )
+                    req_state = self.requests.get(req_id)
+                    req_infos = (
+                        getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
+                    )
                    start_offset = int(self.query_start_loc.cpu[req_index])
                    sched_tokens = int(num_scheduled_tokens_np[req_index])
                    s, e = start_offset, start_offset + sched_tokens
@@ -780,39 +777,16 @@ class OmniNPUModelRunner(NPUModelRunner):
                start_offset = int(self.query_start_loc.cpu[req_index])
                self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)

-    def _get_additional_information(self, scheduler_output: "SchedulerOutput", req_id: str) -> dict:
-        req_infos = None
-        req_state = self.requests.get(req_id)
-        additional_information_cpu = getattr(req_state, "additional_information_cpu", None)
+    def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> None:
        for new_req in scheduler_output.scheduled_new_reqs:
-            if new_req.req_id == req_id:
-                payload_info = getattr(new_req, "additional_information", None)
-                if payload_info is not None:
-                    return payload_info
+            payload_info = getattr(new_req, "additional_information", None)
+            self._merge_additional_information_update(new_req.req_id, payload_info)

        if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"):
            cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {})
-            if isinstance(cached_infos, dict) and req_id in cached_infos:
-                req_infos = cached_infos[req_id]
-                if not isinstance(req_infos, dict):
-                    req_infos = None
-
-        if req_infos is None or req_infos.get("last_talker_hidden", None) is None:
-            if req_infos is None:
-                additional_information_cpu.pop("thinker_embeddings", None)
-                req_infos = additional_information_cpu
-            else:
-                req_infos["last_talker_hidden"] = additional_information_cpu.get("last_talker_hidden", None)
-                req_infos["num_processed_thinker_tokens"] = additional_information_cpu.get(
-                    "num_processed_thinker_tokens", 0
-                )
-            if not isinstance(req_infos, dict):
-                req_infos = None
-
-        if req_infos is None:
-            logger.warning(f"No additional_information found for req_id: {req_id}")
-
-        return req_infos
+            if isinstance(cached_infos, dict):
+                for req_id, req_infos in cached_infos.items():
+                    self._merge_additional_information_update(req_id, req_infos)

    def _preprocess(
        self,
@@ -928,15 +902,12 @@ class OmniNPUModelRunner(NPUModelRunner):
            # Overlay custom prompt_embeds per request for the prompt portion;
            # collect additional_information (tensor/list) for prefill portion only
            decode_req_ids = []
+            if self.vllm_config.model_config.async_chunk:
+                self._update_additional_information(scheduler_output)
            for req_index, req_id in enumerate(self.input_batch.req_ids):
                # Try to get additional_information from multiple sources
-                if self.vllm_config.model_config.async_chunk:
-                    req_infos = self._get_additional_information(scheduler_output, req_id)
-                else:
-                    req_state = self.requests.get(req_id)
-                    req_infos = (
-                        getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
-                    )
+                req_state = self.requests.get(req_id)
+                req_infos = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
                start_offset = int(self.query_start_loc.cpu[req_index])
                sched_tokens = int(num_scheduled_tokens_np[req_index])
                s, e = start_offset, start_offset + sched_tokens
@@ -946,6 +917,7 @@ class OmniNPUModelRunner(NPUModelRunner):
                req_input_ids, req_embeds, update_dict = self.model.preprocess(
                    input_ids=input_ids[s:e], input_embeds=inputs_embeds[s:e], **req_infos
                )
+
                if hasattr(self.model, "talker_mtp") and span_len == 1:
                    last_talker_hidden, text_step = update_dict.pop("mtp_inputs")
                    decode_slice = slice(len(decode_req_ids), len(decode_req_ids) + 1)
@@ -45,7 +45,8 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):

    @classmethod
    def supports_torch_inductor(cls) -> bool:
-        return True
+        # TODO: Enable this when torch compile bugs are resolved
+        return False

    @classmethod
    def get_default_stage_config_path(cls) -> str:
@@ -58,7 +58,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
            self.input_batch.remove_request(req_id)
            # update the request state in self.input_batch
            self.input_batch.add_request(req_state)
-            self._init_mrope_positions(req_state)
+            # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
+            if self.uses_mrope:
+                self._init_mrope_positions(req_state)

    @torch.inference_mode()
    def execute_model(
@@ -832,13 +832,10 @@ class OmniGPUModelRunner(GPUModelRunner):
            # TODO(Peiqi): do we have a more elegant way to do this?
            if hasattr(self.model, "has_postprocess") and self.model.has_postprocess:
                for req_index, req_id in enumerate(self.input_batch.req_ids):
-                    if self.model_config.async_chunk:
-                        req_infos = self._get_additional_information(scheduler_output, req_id)
-                    else:
-                        req_state = self.requests.get(req_id)
-                        req_infos = (
-                            getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
-                        )
+                    req_state = self.requests.get(req_id)
+                    req_infos = (
+                        getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
+                    )
                    start_offset = int(self.query_start_loc.cpu[req_index])
                    sched_tokens = int(num_scheduled_tokens_np[req_index])
                    s, e = start_offset, start_offset + sched_tokens
@@ -879,39 +876,16 @@ class OmniGPUModelRunner(GPUModelRunner):
                start_offset = int(self.query_start_loc.cpu[req_index])
                self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)

-    def _get_additional_information(self, scheduler_output: "SchedulerOutput", req_id: str) -> dict:
-        req_infos = None
-        req_state = self.requests.get(req_id)
-        additional_information_cpu = getattr(req_state, "additional_information_cpu", None)
+    def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> None:
        for new_req in scheduler_output.scheduled_new_reqs:
-            if new_req.req_id == req_id:
-                payload_info = getattr(new_req, "additional_information", None)
-                if payload_info is not None:
-                    return payload_info
+            payload_info = getattr(new_req, "additional_information", None)
+            self._merge_additional_information_update(new_req.req_id, payload_info)

        if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"):
            cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {})
-            if isinstance(cached_infos, dict) and req_id in cached_infos:
-                req_infos = cached_infos[req_id]
-                if not isinstance(req_infos, dict):
-                    req_infos = None
-
-        if req_infos is None or req_infos.get("last_talker_hidden", None) is None:
-            if req_infos is None:
-                additional_information_cpu.pop("thinker_embeddings", None)
-                req_infos = additional_information_cpu
-            else:
-                req_infos["last_talker_hidden"] = additional_information_cpu.get("last_talker_hidden", None)
-                req_infos["num_processed_thinker_tokens"] = additional_information_cpu.get(
-                    "num_processed_thinker_tokens", 0
-                )
-            if not isinstance(req_infos, dict):
-                req_infos = None
-
-        if req_infos is None:
-            logger.warning(f"No additional_information found for req_id: {req_id}")
-
-        return req_infos
+            if isinstance(cached_infos, dict):
+                for req_id, req_infos in cached_infos.items():
+                    self._merge_additional_information_update(req_id, req_infos)

    def _preprocess(
        self,
@@ -1027,15 +1001,11 @@ class OmniGPUModelRunner(GPUModelRunner):
            # Overlay custom prompt_embeds per request for the prompt portion;
            # collect additional_information (tensor/list) for prefill portion only
            decode_req_ids = []
+            if self.vllm_config.model_config.async_chunk:
+                self._update_additional_information(scheduler_output)
            for req_index, req_id in enumerate(self.input_batch.req_ids):
-                # Try to get additional_information from multiple sources
-                if self.vllm_config.model_config.async_chunk:
-                    req_infos = self._get_additional_information(scheduler_output, req_id)
-                else:
-                    req_state = self.requests.get(req_id)
-                    req_infos = (
-                        getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
-                    )
+                req_state = self.requests.get(req_id)
+                req_infos = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
                start_offset = int(self.query_start_loc.cpu[req_index])
                sched_tokens = int(num_scheduled_tokens_np[req_index])
                s, e = start_offset, start_offset + sched_tokens
@@ -1045,6 +1015,7 @@ class OmniGPUModelRunner(GPUModelRunner):
                req_input_ids, req_embeds, update_dict = self.model.preprocess(
                    input_ids=input_ids[s:e], input_embeds=inputs_embeds[s:e], **req_infos
                )
+
                if hasattr(self.model, "talker_mtp") and span_len == 1:
                    last_talker_hidden, text_step = update_dict.pop("mtp_inputs")
                    decode_slice = slice(len(decode_req_ids), len(decode_req_ids) + 1)