[Rebase] Rebase to vllm 0.23.0 (#4286 )

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
[XPU] update Dreamzero to support any HW (#4399 )
2026-06-13 16:43:02 +00:00 · 2026-06-13 15:32:39 +00:00 · 2026-06-13 18:12:37 +08:00 · 2026-06-13 15:42:16 +08:00
47 changed files with 1021 additions and 278 deletions
@@ -10,7 +10,7 @@ steps:
          DOCKER_BUILDKIT: "1"
          # Buildkite will automatically replace this with the actual commit hash
          VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.22.0"
+          VLLM_VERSION: "v0.23.0"
        priority: 100
        timeout_in_minutes: 60
        soft_fail: false
@@ -2,7 +2,7 @@
 # The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
 # be updated on every rebase to match the new upstream vllm commit/tag.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.22.0
+ARG VLLM_BASE_TAG=v0.23.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -16,6 +16,7 @@ RUN apt-get update && \

 RUN uv pip install --system ".[dev]"

+
 RUN ln -sf /usr/bin/python3 /usr/bin/python

 ENTRYPOINT []
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.22.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.23.0
 FROM ${BASE_IMAGE}

 ARG COMMON_WORKDIR=/app
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.23.0
 FROM ${BASE_IMAGE} AS base

 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"

-ARG VLLM_VERSION=v0.22.0
+ARG VLLM_VERSION=v0.23.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm

@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr

 ### Pre-built wheels

-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, and 0.22.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, and 0.23.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).

 === "NVIDIA CUDA"

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**

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

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

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

 ```bash
-uv pip install vllm==0.22.0 --torch-backend=auto
+uv pip install vllm==0.23.0 --torch-backend=auto
 ```

-The 0.22.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.23.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.22.0
+git checkout v0.23.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.23.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.23.0/vllm-0.23.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
    --env "HF_TOKEN=$HF_TOKEN" \
    -p 8091:8091 \
    --ipc=host \
-    vllm/vllm-omni:v0.22.0 \
+    vllm/vllm-omni:v0.23.0 \
    vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.

 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
+uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722
 ```

 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]

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

 ```bash
-uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
+uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722
 ```

-The pre-built 0.22.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.23.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.22.0
+git checkout v0.23.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --env "HF_TOKEN=$HF_TOKEN" \
  -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.22.0 \
+  vllm/vllm-omni-rocm:v0.23.0 \
  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```

@@ -149,7 +149,7 @@ docker run --rm -it \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  --env "HF_TOKEN=$HF_TOKEN" \
  --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.22.0
+  vllm/vllm-omni-rocm:v0.23.0
 ```

 # --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.22.0 --torch-backend=auto
+uv pip install vllm==0.23.0 --torch-backend=auto

 # On ROCm
-uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
+uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722

 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
    It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.

-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.22.0` and vLLM Omni `0.22.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.23.0` and vLLM Omni `0.23.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.

 ## Offline Inference

@@ -42,7 +42,7 @@ th {
 | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
 | `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
-| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | | | |
+| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
@@ -54,7 +54,6 @@ recipes/
 | [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
 | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
 | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
-|

 Within a single recipe file, include different hardware support sections such
 as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like
@@ -1,82 +0,0 @@
-# Ming-omni-tts 0.5B
-
-> Offline and online dense Ming text-to-speech on ROCm
-
-## Summary
-
- Vendor: inclusionAI
- Model: `inclusionAI/Ming-omni-tts-0.5B`
- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference
- Maintainer: Community
-
-## References
-
- [Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
- [Upstream repository](https://github.com/inclusionAI/Ming-omni-tts)
- [Offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
- [Online example](../../examples/online_serving/text_to_speech/ming_tts/)
-
-## Hardware Support
-
-## ROCm
-
-### 1x AMD `gfx942`
-
-#### Environment
-
- OS: Ubuntu 22.04.5 LTS, x86_64
- Python: 3.12.13
- ROCm / HIP: 7.2.53211
- vLLM version: 0.22.0
- vLLM-Omni version or commit: 0.1.dev1873 / `99c35c410`
- Docker image: `vllm/vllm-omni-rocm:v0.22.0`
-
-#### Command
-
-From the vLLM-Omni repository root:
-
-```bash
-docker run --rm \
-    --group-add=video \
-    --ipc=host \
-    --cap-add=SYS_PTRACE \
-    --security-opt seccomp=unconfined \
-    --device /dev/kfd \
-    --device /dev/dri \
-    -v ~/.cache/huggingface:/root/.cache/huggingface \
-    -v "$PWD":/app/vllm-omni \
-    -w /app/vllm-omni \
-    -e VLLM_ROCM_USE_AITER=0 \
-    -p 8091:8091 \
-    vllm/vllm-omni-rocm:v0.22.0 \
-    --model inclusionAI/Ming-omni-tts-0.5B \
-    --deploy-config vllm_omni/deploy/ming_tts.yaml \
-    --omni \
-    --port 8091 \
-    --enforce-eager
-```
-
-#### Verification
-
-```bash
-python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
-    --text "我觉得社会企业同个人都有责任" \
-    --instruction-json '{"方言":"广粤话"}' \
-    --ref-audio /path/to/yue_prompt.wav \
-    --max-new-tokens 200 \
-    --output dialect.wav
-```
-
-`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
-for the dialect case.
-
-## Notes
-
- The official ROCm image includes the platform dependencies.
- See the [ROCm installation guide](../../docs/getting_started/installation/gpu.md)
-  for interactive and source-build workflows.
- The [offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
-  can be run from an interactive container.
- The tested environment uses `--enforce-eager`.
- Output is mono 44.1 kHz audio.
@@ -0,0 +1,181 @@
+# Ming-omni-tts
+
+> Online serving for voice synthesis, cloning, and design
+
+## Summary
+
+- Vendor: inclusionAI
+- Model: `inclusionAI/Ming-omni-tts-0.5B`
+- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
+- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API
+- Maintainer: Community
+
+## References
+
+- [Huggingface Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
+- Upstream repository [inclusionAI/Ming-omni-tts](https://github.com/inclusionAI/Ming-omni-tts)
+
+
+## Hardware Support
+
+This recipe documents a validated ROCm configuration and a CUDA configuration for the dense 0.5B two-stage TTS pipeline deployment.
+Other hardware is welcome as community validation lands.
+
+## CUDA
+
+### 1x H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- CUDA Driver Version: 590.48.01
+- CUDA: 13.0
+- vLLM version: 0.22.0
+- vLLM-Omni version or commit: 0342827d
+
+#### Command
+
+Launch the two-stage talker:
+
+```bash
+vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --omni \
+    --port 8091 \
+    --enforce-eager
+```
+
+#### Verification
+
+Basic synthesis (save the WAV bytes):
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+      "model": "inclusionAI/Ming-omni-tts-0.5B",
+      "input": "你好，这是 Ming 在线语音合成测试。",
+      "response_format": "wav"
+    }' --output ming_tts_basic.wav
+```
+
+The following requests condition on reference clips from upstream audio assets. The server resolves `ref_audio` as an `http(s)://` URL, a `file://` path, or an inline base64 `data:` URL, so the examples point it straight at the public cookbook fixtures in the vendor repo's [`inclusionAI/Ming-omni-tts/tree/main/data/wavs`](https://github.com/inclusionAI/Ming-omni-tts/tree/main/data/wavs) directory:
+
+```bash
+BASE=https://raw.githubusercontent.com/inclusionAI/Ming-omni-tts/main/data/wavs
+```
+
+And note that the server latency will include automatically resolving and downloading input files.
+
+Reference-audio zero-shot cloning:
+
+```bash
+REF_AUDIO="$BASE/10002287-00000094.wav"
+jq -n --arg ref_audio "$REF_AUDIO" \
+    '{
+       model: "inclusionAI/Ming-omni-tts-0.5B",
+       input: "我们的愿景是构建未来服务业的数字化基础设施，为世界带来更多微小而美好的改变。",
+       ref_audio: $ref_audio,
+       ref_text: "在此奉劝大家别乱打美白针。",
+       max_new_tokens: 200,
+       response_format: "wav"
+     }' \
+| curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d @- --output ming_tts_zero_shot.wav
+```
+
+Podcast-style multi-speaker generation uses one reference clip plus one transcript per speaker, with the `speaker_N` labels in `input`/`ref_text` lined up with the reference order:
+
+```bash
+REF_A="$BASE/CTS-CN-F2F-2019-11-11-423-012-A.wav"
+REF_B="$BASE/CTS-CN-F2F-2019-11-11-423-012-B.wav"
+
+INPUT=' speaker_1:你可以说一下，就大概说一下，可能虽然我也不知道，我看过那部电影没有。
+ speaker_2:就是那个叫什么，变相一节课的嘛。
+ speaker_1:嗯。
+ speaker_2:一部搞笑的电影。
+ speaker_1:一部搞笑的。'
+
+REF_TEXT=' speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试，做个，当服务员还要去笔试了
+ speaker_2:对啊，这真的很奇怪，就是 单纯的因，单纯自己工资不高，只是因为可能人家那个店比较出名一点，就对你苛刻要求'
+
+jq -n \
+    --arg input "$INPUT" \
+    --arg ref_text "$REF_TEXT" \
+    --arg a "$REF_A" \
+    --arg b "$REF_B" \
+    '{
+       model: "inclusionAI/Ming-omni-tts-0.5B",
+       input: $input,
+       ref_audio: [$a, $b],
+       ref_text: $ref_text,
+       response_format: "wav"
+     }' \
+| curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d @- --output ming_tts_podcast_style.wav
+```
+
+## ROCm
+
+### 1x AMD MI300X
+
+#### Environment
+
+- OS: Ubuntu 22.04.5 LTS, x86_64
+- Python: 3.12.13
+- ROCm / HIP: 7.2.53211
+- vLLM version: 0.22.0
+- vLLM-Omni version or commit: 0.1.dev1873 / `99c35c410`
+- Docker image: `vllm/vllm-omni-rocm:v0.22.0`
+
+#### Command
+
+From the vLLM-Omni repository root:
+
+```bash
+docker run --rm \
+    --group-add=video \
+    --ipc=host \
+    --cap-add=SYS_PTRACE \
+    --security-opt seccomp=unconfined \
+    --device /dev/kfd \
+    --device /dev/dri \
+    -v ~/.cache/huggingface:/root/.cache/huggingface \
+    -v "$PWD":/app/vllm-omni \
+    -w /app/vllm-omni \
+    -e VLLM_ROCM_USE_AITER=0 \
+    -p 8091:8091 \
+    vllm/vllm-omni-rocm:v0.22.0 \
+    --model inclusionAI/Ming-omni-tts-0.5B \
+    --deploy-config vllm_omni/deploy/ming_tts.yaml \
+    --omni \
+    --port 8091 \
+    --enforce-eager
+```
+
+#### Verification
+
+```bash
+python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
+    --text "我觉得社会企业同个人都有责任" \
+    --instruction-json '{"方言":"广粤话"}' \
+    --ref-audio /path/to/yue_prompt.wav \
+    --max-new-tokens 200 \
+    --output dialect.wav
+```
+
+`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
+for the dialect case.
+
+## Notes
+
+- The official ROCm image includes the platform dependencies.
+- See the [ROCm installation guide](../../docs/getting_started/installation/gpu.md) for interactive and source-build workflows.
+- The reference clips above are a subset of the upstream [`inclusionAI/Ming-omni-tts/tree/main/data/wavs`](https://github.com/inclusionAI/Ming-omni-tts/tree/main/data/wavs) cookbook fixtures.
+- To clone from local audio instead of a URL, pass a `file://` URI (launch the server with `--allowed-local-media-path <dir>`), or base64-encode the clip into a `data:` URL written to a payload file and send it with `curl -d @payload.json`. Putting the base64 directly on the `curl` command line overflows the OS argument limit (`bash: /usr/bin/curl: Argument list too long`).
+- The tested environment uses `--enforce-eager`.
+- Non-streaming responses return WAV bytes; streaming responses return PCM.
+- Output is mono 44.1 kHz audio.
@@ -1,10 +1,15 @@
 # Common dependencies for all platforms
-# transformers: keep aligned with upstream vLLM constraint.
+# transformers: kept exactly aligned with the upstream vLLM constraint.
 # vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
-# 5.9.0 breaks models that define _keys_to_ignore_on_load_unexpected as a list
-# (e.g. OpenMOSS-Team/MOSS-TTS-Nano) because it changed
-# _adjust_missing_and_unexpected_keys from list+list to set|set.
-transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0, < 5.9.0
+# The previous extra "< 5.9.0" exclusion was removed when re-aligning with
+# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from
+# list+list to set|set, breaking models that declare
+# _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano)
+# with "TypeError: unsupported operand type(s) for |: 'list' and 'set'".
+# That breakage is now handled in code instead of by pinning:
+# MossTTSNanoForGeneration wraps weight loading in
+# transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py).
+transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0
@@ -42,10 +42,12 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
                        yield chunk

            class MockResponse:
-                def __init__(self):
-                    self.status = 200
-                    self.reason = "OK"
-                    self.content = _Content(SSE_CHUNKS)
+                def __init__(self, status=200, reason="OK", chunks=None):
+                    self.status = status
+                    self.reason = reason
+                    self.content = _Content(chunks or SSE_CHUNKS)
+                    self._json_data = None
+                    self._text_data = None

                async def __aenter__(self):
                    return self
@@ -53,20 +55,49 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
                async def __aexit__(self, exc_type, exc, tb):
                    return False

+                async def text(self):
+                    if self._text_data is not None:
+                        return self._text_data
+                    return json.dumps({"tokens": [1, 2, 3, 4, 5]})
+
+                async def json(self):
+                    if self._json_data is not None:
+                        return self._json_data
+                    return {"tokens": [1, 2, 3, 4, 5]}
+
+                def raise_for_status(self):
+                    pass
+
            class MockClientSession:
                def __init__(self, *args, **kwargs):
                    pass

+                async def __aenter__(self):
+                    return self
+
+                async def __aexit__(self, exc_type, exc, tb):
+                    return False
+
                def post(self, url=None, *args, **kwargs):
                    if url is not None:
                        POST_CALLS.append(url)
-                    return MockResponse()
+                    return MockResponse(status=200)
+
+                async def get(self, url=None, *args, **kwargs):
+                    if url is not None:
+                        POST_CALLS.append(url)
+                    return MockResponse(status=200)

                async def close(self):
                    return None

-            import vllm_omni.benchmarks.patch.patch as patch_mod
+            # Patch globally so modules loaded after this also see the mock
+            import aiohttp
+            aiohttp.ClientSession = MockClientSession
+            aiohttp.TCPConnector = lambda *args, **kwargs: object()

+            # Also patch in the patch module
+            import vllm_omni.benchmarks.patch.patch as patch_mod
            patch_mod.aiohttp.ClientSession = MockClientSession
            patch_mod.aiohttp.TCPConnector = lambda *args, **kwargs: object()

@@ -136,8 +167,13 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):

    expected_url = f"http://127.0.0.1:{port}/v1/chat/completions"
    requested_urls = calls["requested_urls"]
-    sent_requests = len(requested_urls)
+    # Count only benchmark requests (to the chat completions endpoint),
+    # excluding tokenize/detokenize alignment calls from the upstream.
+    bench_requests = [url for url in requested_urls if url == expected_url]
+    sent_bench_requests = len(bench_requests)

-    assert result["completed"] == sent_requests == num_prompts
-    assert requested_urls
-    assert all(url == expected_url for url in requested_urls), f"Unexpected target URLs: {requested_urls}"
+    assert result["completed"] == sent_bench_requests == num_prompts, (
+        f"completed={result['completed']}, bench_requests={sent_bench_requests}, all_requested_urls={requested_urls}"
+    )
+    assert bench_requests
+    assert all(url == expected_url for url in bench_requests), f"Unexpected target URLs: {bench_requests}"
@@ -578,12 +578,27 @@ def test_rmsnorm_fused_native_parity_2d_input():


 def test_rmsnorm_fused_native_parity_custom_weight():
-    """Verify fused and native RMSNorm parity with non-default weights."""
+    """Verify fused and native RMSNorm parity with non-default weights.
+
+    Note: bfloat16 requires wider tolerance because the fused CUDA kernel
+    (vllm._custom_ops.rms_norm) computes in bf16 precision internally,
+    while forward_native uses float32 intermediates. With random custom
+    weights (not the default ones-weight), this precision gap exceeds 1e-3.
+
+    The gap surfaced after upstream vLLM routed RMSNorm through the vLLM IR
+    op (ir.ops.rms_norm) — see upstream commit 40bb17502
+    "[vLLM IR] 1/N Implement IR skeleton and rms_norm op (#33825)" — which
+    changed which fused kernel/precision path the layer dispatches to.
+    """
    from vllm_omni.diffusion.layers.norm import RMSNorm

    hidden_size = 64
    eps = 1e-6

+    # Per-dtype tolerances: bf16 fused kernel has inherent precision limit
+    atol_map = {torch.float32: 1e-3, torch.float16: 1e-3, torch.bfloat16: 2e-2}
+    rtol_map = {torch.float32: 1e-3, torch.float16: 1e-3, torch.bfloat16: 1e-2}
+
    for dtype in [torch.float32, torch.float16, torch.bfloat16]:
        torch.manual_seed(789)
        norm = RMSNorm(hidden_size, eps=eps).cuda()
@@ -594,4 +609,4 @@ def test_rmsnorm_fused_native_parity_custom_weight():
        out_fused = norm._forward_fused(x)
        out_native = norm.forward_native(x)

-        torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3)
+        torch.testing.assert_close(out_fused, out_native, atol=atol_map[dtype], rtol=rtol_map[dtype])
@@ -9,6 +9,7 @@ import torch
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.ipc import (
    _SHM_TENSOR_THRESHOLD,
+    DIFFUSION_RPC_RESULT_ENVELOPE,
    _pack_value_if_large,
    _unpack_if_shm_handle,
    pack_diffusion_output_shm,
@@ -81,6 +82,31 @@ def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
    torch.testing.assert_close(output.output[1], frames[1])


+def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
+    tensor = torch.arange(300_000, dtype=torch.float32)
+    envelope = {
+        "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+        "result": DiffusionOutput(output=tensor),
+        "rank_statuses": [{"rank": 0, "ok": True}],
+    }
+
+    packed = pack_diffusion_output_shm(envelope)
+
+    assert packed is envelope
+    result = packed["result"]
+    assert isinstance(result, DiffusionOutput)
+    assert result.output["__tensor_shm__"] is True
+    assert packed["rank_statuses"] == [{"rank": 0, "ok": True}]
+
+    unpacked = unpack_diffusion_output_shm(packed)
+
+    assert unpacked is envelope
+    result = unpacked["result"]
+    assert isinstance(result, DiffusionOutput)
+    torch.testing.assert_close(result.output, tensor)
+    assert unpacked["rank_statuses"] == [{"rank": 0, "ok": True}]
+
+
 def test_pack_value_keeps_tensor_at_threshold_inline() -> None:
    tensor = torch.arange(
        _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(),
@@ -17,8 +17,10 @@ from vllm.v1.engine.exceptions import EngineDeadError
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
+from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE
 from vllm_omni.diffusion.sched import RequestScheduler
 from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

@@ -305,6 +307,80 @@ class TestSerialEngineOperations:
        assert len(results) == 1
        assert results[0].error == "rank0"

+    def test_collective_rpc_all_rank_status_error_propagation(self):
+        engine, _, _, res_q = _make_engine(num_gpus=2)
+
+        res_q.put(
+            {
+                "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+                "method": "add_lora",
+                "result": True,
+                "rank_statuses": [
+                    {"rank": 0, "ok": True, "bool_result": True},
+                    {
+                        "rank": 1,
+                        "ok": False,
+                        "error": "rank1 boom",
+                        "error_type": "RuntimeError",
+                        "traceback": "rank1 traceback",
+                    },
+                ],
+            }
+        )
+
+        with pytest.raises(RuntimeError) as excinfo:
+            engine.collective_rpc("add_lora")
+        error = str(excinfo.value)
+        assert "rank 1" in error
+        assert "rank1 boom" in error
+        assert "rank1 traceback" in error
+
+    def test_collective_rpc_all_rank_bool_false_is_aggregated(self):
+        engine, _, _, res_q = _make_engine(num_gpus=2)
+
+        res_q.put(
+            {
+                "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+                "method": "remove_lora",
+                "result": True,
+                "rank_statuses": [
+                    {"rank": 0, "ok": True, "bool_result": True},
+                    {"rank": 1, "ok": True, "bool_result": False},
+                ],
+            }
+        )
+
+        assert engine.collective_rpc("remove_lora") == [False]
+
+    def test_collective_rpc_collects_rank_status_only_for_control_plane_all_rank_rpc(self):
+        executor, req_q, res_q = _make_executor(num_gpus=2)
+
+        res_q.put(_tagged_output("forward"))
+        result = executor.collective_rpc(
+            "execute_stepwise",
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        forward_rpc = req_q.get_nowait()
+
+        assert result.error == "forward"
+        assert forward_rpc["exec_all_ranks"] is True
+        assert forward_rpc["collect_rank_status"] is False
+
+        res_q.put(
+            {
+                "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+                "method": "remove_lora",
+                "result": True,
+                "rank_statuses": [{"rank": 0, "ok": True, "bool_result": True}],
+            }
+        )
+        assert executor.collective_rpc("remove_lora") == [True]
+        control_rpc = req_q.get_nowait()
+
+        assert control_rpc["exec_all_ranks"] is True
+        assert control_rpc["collect_rank_status"] is True
+
    def test_serial_add_req_then_collective_rpc(self):
        engine, _, req_q, res_q = _make_engine()
        wt = _start_worker(req_q, res_q, count=2)
@@ -348,6 +424,115 @@ class TestSerialEngineOperations:
            engine.collective_rpc("anything")


+class TestWorkerProcRpcRankStatus:
+    def _make_worker_proc(self, has_result_mq: bool = True):
+        proc = object.__new__(WorkerProc)
+        proc.gpu_id = 0
+        proc.result_mq = object() if has_result_mq else None
+        proc.worker = SimpleNamespace(execute_method=Mock(return_value=True))
+        return proc
+
+    def test_execute_rpc_returns_rank_status_envelope(self, monkeypatch):
+        proc = self._make_worker_proc()
+
+        monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2)
+
+        def _all_gather_object(out, local):
+            out[0] = local
+            out[1] = {
+                "rank": 1,
+                "ok": False,
+                "error": "rank1 boom",
+                "error_type": "RuntimeError",
+                "traceback": "trace",
+                "bool_result": None,
+            }
+
+        monkeypatch.setattr(torch.distributed, "all_gather_object", _all_gather_object)
+
+        result, should_reply = proc.execute_rpc(
+            {
+                "method": "remove_lora",
+                "args": (),
+                "kwargs": {},
+                "output_rank": 0,
+                "exec_all_ranks": True,
+                "collect_rank_status": True,
+            }
+        )
+
+        assert should_reply is True
+        assert result["type"] == DIFFUSION_RPC_RESULT_ENVELOPE
+        assert result["result"] is True
+        assert result["rank_statuses"][0]["rank"] == 0
+        assert result["rank_statuses"][1]["rank"] == 1
+        assert result["rank_statuses"][1]["ok"] is False
+
+    def test_execute_rpc_local_exception_is_reported_in_envelope(self, monkeypatch):
+        proc = self._make_worker_proc()
+        proc.worker.execute_method = Mock(side_effect=RuntimeError("local boom"))
+        monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
+
+        result, should_reply = proc.execute_rpc(
+            {
+                "method": "add_lora",
+                "args": (),
+                "kwargs": {},
+                "output_rank": 0,
+                "exec_all_ranks": True,
+                "collect_rank_status": True,
+            }
+        )
+
+        assert should_reply is True
+        assert result["type"] == DIFFUSION_RPC_RESULT_ENVELOPE
+        assert len(result["rank_statuses"]) == 1
+        status = result["rank_statuses"][0]
+        assert status["rank"] == 0
+        assert status["ok"] is False
+        assert status["error"] == "local boom"
+        assert status["error_type"] == "RuntimeError"
+        assert status["bool_result"] is None
+        assert "local boom" in status["traceback"]
+
+    def test_execute_rpc_rejects_collect_rank_status_without_all_ranks(self):
+        proc = self._make_worker_proc()
+
+        with pytest.raises(ValueError, match="collect_rank_status requires exec_all_ranks=True"):
+            proc.execute_rpc(
+                {
+                    "method": "ping",
+                    "args": (),
+                    "kwargs": {},
+                    "output_rank": 0,
+                    "exec_all_ranks": False,
+                    "collect_rank_status": True,
+                }
+            )
+
+        proc.worker.execute_method.assert_not_called()
+
+    def test_execute_rpc_non_collect_exception_preserves_original_type(self):
+        proc = self._make_worker_proc()
+        original = ValueError("local boom")
+        proc.worker.execute_method = Mock(side_effect=original)
+
+        with pytest.raises(ValueError) as excinfo:
+            proc.execute_rpc(
+                {
+                    "method": "bad",
+                    "args": (),
+                    "kwargs": {},
+                    "output_rank": 0,
+                    "exec_all_ranks": False,
+                    "collect_rank_status": False,
+                }
+            )
+
+        assert excinfo.value is original
+
+
 # ───────── error handling: EngineDeadError propagation through layers ─────


@@ -140,6 +140,7 @@ class TestHiggsAudioV3OnlineHappyPath:
            }
        )

+    @pytest.mark.skip(reason="issue#4411")
    @pytest.mark.core_model
    @pytest.mark.tts
    @hardware_test(res={"cuda": "H100"}, num_cards=1)
@@ -1193,6 +1193,7 @@ async def benchmark(
    ramp_up_end_rps: int | None = None,
    ready_check_timeout_sec: int = 600,
    ssl_context: ssl.SSLContext | bool | None = None,
+    self_timed: bool = False,
 ):
    try:
        request_func = ASYNC_REQUEST_FUNCS[endpoint_type]
@@ -1355,6 +1356,7 @@ async def benchmark(
        ramp_up_strategy,
        ramp_up_start_rps,
        ramp_up_end_rps,
+        self_timed,
    ):
        if ramp_up_strategy is not None:
            current_int_rps = int(current_request_rate)
@@ -1443,6 +1445,7 @@ async def benchmark(
            defs.AVERAGE_PIXELS_PER_IMAGE: getattr(metrics, defs.AVERAGE_PIXELS_PER_IMAGE),
            defs.MEAN_DENOISE_STEP_LATENCY_MS: getattr(metrics, defs.MEAN_DENOISE_STEP_LATENCY_MS),
            "input_lens": [output.prompt_len for output in outputs],
+            "start_times": [output.start_time for output in outputs],
            "output_lens": actual_output_lens,
            "ttfts": [output.ttft for output in outputs],
            "itls": [output.itl for output in outputs],
@@ -712,6 +712,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):

        self._omits_kv_transfer_cache.pop(request.request_id, None)

+        # [Upstream compat] Discard request from in-flight prefills set added
+        # upstream for routed-experts in-flight reservation tracking.
+        # Use getattr for safety with test __new__ code paths.
+        getattr(self, "_inflight_prefills", set()).discard(request)
+
        # 1. Standard cleanup parts from base _free_request
        connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request)

@@ -6,8 +6,7 @@ from typing import TYPE_CHECKING, Any

 from vllm.utils.import_utils import resolve_obj_by_qualname

-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.data import OmniDiffusionConfig

 if TYPE_CHECKING:
    from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
@@ -69,11 +68,6 @@ class DiffusionExecutor(ABC):
        """Initialize the executor (e.g., launch workers, setup IPC)."""
        pass

-    @abstractmethod
-    def add_req(self, requests: OmniDiffusionRequest) -> DiffusionOutput:
-        """Add requests to the execution queue."""
-        pass
-
    @abstractmethod
    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
        """Execute request-mode work from a scheduler output."""
@@ -16,8 +16,7 @@ from vllm.v1.engine.exceptions import EngineDeadError

 from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
-from vllm_omni.diffusion.ipc import unpack_diffusion_output_shm
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, unpack_diffusion_output_shm
 from vllm_omni.diffusion.worker import WorkerProc

 if TYPE_CHECKING:
@@ -133,6 +132,57 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                    raise EngineDeadError()
                continue

+    @staticmethod
+    def _raise_for_rpc_error_dict(response: Any) -> None:
+        if isinstance(response, dict) and response.get("status") == "error":
+            raise RuntimeError(
+                f"Worker failed with error '{response.get('error')}', "
+                "please check the stack trace above for the root cause"
+            )
+
+    @staticmethod
+    def _unwrap_rpc_result_envelope(response: Any) -> Any:
+        if not (isinstance(response, dict) and response.get("type") == DIFFUSION_RPC_RESULT_ENVELOPE):
+            return response
+
+        rank_statuses = response.get("rank_statuses") or []
+        failed = [status for status in rank_statuses if not status.get("ok", False)]
+        if failed:
+            details = "; ".join(
+                f"rank {status.get('rank')}: {status.get('error_type') or 'Error'}: {status.get('error')}"
+                for status in failed
+            )
+            tracebacks = "\n\n".join(
+                f"rank {status.get('rank')} traceback:\n{status['traceback']}"
+                for status in failed
+                if status.get("traceback")
+            )
+            if tracebacks:
+                details = f"{details}\n\n{tracebacks}"
+            method = response.get("method", "<unknown>")
+            raise RuntimeError(f"RPC '{method}' failed on worker rank(s): {details}")
+
+        result = response.get("result")
+        if isinstance(result, bool):
+            # Only bool-returning RPCs participate in the all-rank AND.
+            # Non-bool results leave bool_result unset and are ignored here.
+            bool_results = [
+                status.get("bool_result") for status in rank_statuses if status.get("bool_result") is not None
+            ]
+            if bool_results and not all(bool_results):
+                return False
+        return result
+
+    @staticmethod
+    def _handle_rpc_response(response: Any) -> Any:
+        MultiprocDiffusionExecutor._raise_for_rpc_error_dict(response)
+        response = MultiprocDiffusionExecutor._unwrap_rpc_result_envelope(response)
+        # After unwrapping, a worker method result may itself be the same
+        # {"status": "error"} shape produced by worker_busy_loop transport
+        # failures. Preserve the pre-envelope error handling for that case.
+        MultiprocDiffusionExecutor._raise_for_rpc_error_dict(response)
+        return response
+
    def _launch_workers(self, broadcast_handle, wake_events):
        od_config = self.od_config
        logger.info("Starting server...")
@@ -257,38 +307,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
        """Register a callback invoked when a worker process dies."""
        self._failure_callbacks.append(callback)

-    def add_req(self, request: OmniDiffusionRequest) -> DiffusionOutput:
-        self._ensure_open()
-        rpc_request = {
-            "type": "rpc",
-            "method": "generate",
-            "args": (request,),
-            "kwargs": {},
-            "output_rank": 0,
-            "exec_all_ranks": True,
-        }
-
-        try:
-            self._broadcast_mq.enqueue(rpc_request)
-            response = self._result_mq.dequeue()
-
-            try:
-                unpack_diffusion_output_shm(response)
-            except Exception as e:
-                logger.warning("SHM unpack failed (data may already be inline): %s", e)
-
-            if isinstance(response, dict) and response.get("status") == "error":
-                raise RuntimeError(
-                    f"Worker failed with error '{response.get('error')}', "
-                    "please check the stack trace above for the root cause"
-                )
-            if not isinstance(response, DiffusionOutput):
-                raise RuntimeError(f"Unexpected response type for generate: {type(response)!r}")
-            return response
-        except Exception as e:
-            logger.error(f"Generate call failed: {e}")
-            raise
-
    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
        """Adapt request-mode scheduler output to worker execute_model RPC."""
        from vllm_omni.diffusion.worker.utils import RunnerOutput
@@ -358,16 +376,21 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
        deadline = None if timeout is None else time.monotonic() + timeout
        kwargs = kwargs or {}

-        # Prepare RPC request message
-        # When unique_reply_rank is None, all workers must execute the RPC
-        # but only rank 0 can reply (it's the only one with a result_mq).
+        # Prepare RPC request message. When unique_reply_rank is None, all
+        # workers must execute the RPC but only rank 0 can reply (it's the
+        # only one with a result_mq). Collect detailed rank statuses only for
+        # this control-plane all-rank path; forward-path exec_all_ranks RPCs
+        # avoid the per-step host object gather.
+        execute_all_ranks = unique_reply_rank is None or exec_all_ranks
+        collect_rank_status = unique_reply_rank is None
        rpc_request = {
            "type": "rpc",
            "method": method,
            "args": args,
            "kwargs": kwargs,
            "output_rank": unique_reply_rank if unique_reply_rank is not None else 0,
-            "exec_all_ranks": unique_reply_rank is None or exec_all_ranks,
+            "exec_all_ranks": execute_all_ranks,
+            "collect_rank_status": collect_rank_status,
        }

        try:
@@ -386,12 +409,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                except Exception as e:
                    logger.warning("SHM unpack failed (data may already be inline): %s", e)

-                # Check if response indicates an error
-                if isinstance(response, dict) and response.get("status") == "error":
-                    raise RuntimeError(
-                        f"Worker failed with error '{response.get('error')}', "
-                        "please check the stack trace above for the root cause"
-                    )
+                response = MultiprocDiffusionExecutor._handle_rpc_response(response)

                responses.append(response)

@@ -18,6 +18,7 @@ import torch
 from vllm_omni.diffusion.data import DiffusionOutput

 _SHM_TENSOR_THRESHOLD = 1_000_000  # 1 MB
+DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result"


 def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]:
@@ -130,15 +131,26 @@ def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
    return output


+def _is_rpc_result_envelope(output: object) -> bool:
+    return isinstance(output, dict) and output.get("type") == DIFFUSION_RPC_RESULT_ENVELOPE
+
+
 def pack_diffusion_output_shm(output: object) -> object:
    """Replace large tensors in diffusion worker outputs with SHM handles.

    Supports either a bare ``DiffusionOutput`` or a wrapper object carrying one
-    in ``.result`` (for example ``RunnerOutput``).
+    in ``.result`` (for example ``RunnerOutput``), or an RPC result envelope
+    carrying the diffusion output in ``["result"]``.
    """
    if isinstance(output, DiffusionOutput):
        return _pack_diffusion_fields(output)

+    if _is_rpc_result_envelope(output):
+        result = output.get("result")
+        if isinstance(result, DiffusionOutput):
+            output["result"] = _pack_diffusion_fields(result)
+        return output
+
    result = getattr(output, "result", None)
    if isinstance(result, DiffusionOutput):
        output.result = _pack_diffusion_fields(result)
@@ -158,6 +170,12 @@ def unpack_diffusion_output_shm(output: object) -> object:
    if isinstance(output, DiffusionOutput):
        return _unpack_diffusion_fields(output)

+    if _is_rpc_result_envelope(output):
+        result = output.get("result")
+        if isinstance(result, DiffusionOutput):
+            output["result"] = _unpack_diffusion_fields(result)
+        return output
+
    result = getattr(output, "result", None)
    if isinstance(result, DiffusionOutput):
        output.result = _unpack_diffusion_fields(result)
@@ -617,6 +617,8 @@ class CausalWanModel(nn.Module):
    Architecture (14B): 40 layers, dim=5120, heads=40, ffn=13824
    """

+    _layerwise_offload_blocks_attrs = ["blocks"]
+
    def __init__(
        self,
        model_type: str = "t2v",
@@ -49,6 +49,7 @@ from vllm_omni.diffusion.models.dreamzero.utils import (
 )
 from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.platforms import current_omni_platform

 logger = logging.getLogger(__name__)
 MAX_DREAMZERO_SESSIONS = 64
@@ -363,12 +364,10 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
        latents: tuple[torch.Tensor, torch.Tensor],
        do_true_cfg: bool,
    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Post-step sync: .contiguous() + cuda.synchronize()"""
+        """Post-step sync: .contiguous() + platform synchronize."""
        latents = tuple(t.contiguous() for t in latents)
        if do_true_cfg and get_classifier_free_guidance_world_size() > 1:
-            device = next((t.device for t in latents if t.is_cuda), None)
-            if device is not None:
-                torch.cuda.current_stream(device).synchronize()
+            current_omni_platform.synchronize()
        return latents

    # -----------------------------------------------------------------------
@@ -4,7 +4,6 @@
 from typing import Any

 import vllm.forward_context as _vllm_fc
-from vllm.model_executor.layers.fused_moe import FusedMoE
 from vllm.utils.import_utils import resolve_obj_by_qualname

 from vllm_omni.platforms import current_omni_platform
@@ -25,22 +24,93 @@ def _set_forward_context_num_tokens(num_tokens: int) -> None:
        forward_context.in_profile_run = False


-class HunyuanFusedMoEDefault(FusedMoE):
-    def __init__(self, *, prefix: str = "", **kwargs: Any) -> None:
+class HunyuanFusedMoEDefault:
+    """Adapter that configures the upstream FusedMoE ``MoERunner`` for HunyuanImage3.
+
+    Upstream commit dc68bd8c41 refactored FusedMoE from a class (``nn.Module``)
+    into a factory function that returns a ``MoERunner`` instance, whose expert
+    weights live in a ``routed_experts`` submodule
+    (``...experts.routed_experts.w13_weight`` / ``...w2_weight``).
+
+    This adapter builds that runner, installs the omni-specific forward-context
+    setup and one-shot kernel-initialisation hook, and returns the runner
+    *directly* from ``__new__`` so the parent MoE block registers it as a real
+    ``nn.Module`` submodule.
+
+    Returning the runner (rather than wrapping it in a plain object that holds it
+    in an attribute) is required for correctness: a non-Module wrapper hides the
+    runner's parameters from ``named_parameters()``, so ``load_weights`` cannot
+    find ``...experts.routed_experts.w13_weight`` and raises ``KeyError`` during
+    weight loading.
+    """
+
+    def __new__(cls, *, prefix: str = "", **kwargs: Any) -> Any:
        # Current vLLM FusedMoE handles output reduction internally.
        kwargs.pop("reduce_results", None)
-        super().__init__(prefix=prefix, **kwargs)
-        self._prefix = prefix
-        self._init_hook_handle = self.register_forward_pre_hook(self._initialize_kernel_hook, with_kwargs=True)
+        # FusedMoE is now a factory function — call it to get a MoERunner.
+        from vllm.model_executor.layers.fused_moe import FusedMoE as _FusedMoE

-    def _initialize_kernel_hook(self, module: Any, args: Any, kwargs: Any) -> None:
-        if self.quant_method and getattr(self.quant_method, "moe_kernel", None) is None:
-            self.quant_method.process_weights_after_loading(self)
-        self._init_hook_handle.remove()
+        moe_runner = _FusedMoE(prefix=prefix, **kwargs)

-    def forward(self, hidden_states: Any, router_logits: Any) -> Any:
-        _set_forward_context_num_tokens(hidden_states.shape[0])
-        return super().forward(hidden_states, router_logits)
+        # Set ForwardContext.num_tokens before each forward. After the rebase to
+        # vLLM 0.18.0 FusedMoE requires this; without it MoE routing is silently
+        # incorrect. Previously done in the wrapper's forward(); now a pre-hook
+        # on the runner so we can return the runner directly.
+        def _num_tokens_pre_hook(module: Any, args: Any, kwargs: Any) -> None:
+            hidden_states = kwargs.get("hidden_states")
+            if hidden_states is None and args:
+                hidden_states = args[0]
+            if hidden_states is not None:
+                _set_forward_context_num_tokens(hidden_states.shape[0])
+
+        moe_runner.register_forward_pre_hook(_num_tokens_pre_hook, with_kwargs=True)
+
+        # One-shot lazy kernel initialisation on the first forward (no-op unless
+        # the runner exposes an uninitialised quant_method). Mirrors the prior
+        # wrapper behaviour exactly, just bound to the runner module.
+        init_handle: Any = None
+
+        def _kernel_init_pre_hook(module: Any, args: Any, kwargs: Any) -> None:
+            nonlocal init_handle
+            quant_method = getattr(module, "quant_method", None)
+            if quant_method is not None and getattr(quant_method, "moe_kernel", None) is None:
+                quant_method.process_weights_after_loading(module)
+            if init_handle is not None:
+                init_handle.remove()
+
+        init_handle = moe_runner.register_forward_pre_hook(_kernel_init_pre_hook, with_kwargs=True)
+
+        return moe_runner
+
+    @staticmethod
+    def make_expert_params_mapping(
+        model: Any,
+        ckpt_gate_proj_name: str,
+        ckpt_down_proj_name: str,
+        ckpt_up_proj_name: str,
+        num_experts: int,
+        num_redundant_experts: int = 0,
+    ) -> list[tuple[str, str, int, str]]:
+        """Delegate to the upstream standalone function.
+
+        Upstream vLLM refactored ``FusedMoE`` from a class (which had
+        ``make_expert_params_mapping`` as a classmethod) to a factory
+        function.  The method was moved to a standalone function
+        ``fused_moe_make_expert_params_mapping`` in
+        ``vllm.model_executor.layers.fused_moe``.
+        """
+        from vllm.model_executor.layers.fused_moe import (
+            fused_moe_make_expert_params_mapping,
+        )
+
+        return fused_moe_make_expert_params_mapping(
+            model,
+            ckpt_gate_proj_name=ckpt_gate_proj_name,
+            ckpt_down_proj_name=ckpt_down_proj_name,
+            ckpt_up_proj_name=ckpt_up_proj_name,
+            num_experts=num_experts,
+            num_redundant_experts=num_redundant_experts,
+        )


 class HunyuanFusedMoE:
@@ -12,6 +12,7 @@ import gc
 import multiprocessing as mp
 import os
 import signal
+import traceback
 from collections.abc import Iterable, Iterator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
@@ -41,7 +42,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
    initialize_model_parallel,
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
-from vllm_omni.diffusion.ipc import pack_diffusion_output_shm
+from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
@@ -740,6 +741,18 @@ class WorkerProc:
        """Receive messages from broadcast queue."""
        return self.mq.dequeue(indefinite=True)

+    def _gather_rpc_rank_statuses(self, status: dict[str, Any]) -> list[dict[str, Any]]:
+        if not torch.distributed.is_initialized():
+            return [status]
+
+        world_size = torch.distributed.get_world_size()
+        statuses: list[dict[str, Any] | None] = [None] * world_size
+        torch.distributed.all_gather_object(statuses, status)
+        missing_ranks = [rank for rank, rank_status in enumerate(statuses) if rank_status is None]
+        if missing_ranks:
+            logger.warning("RPC rank status gather returned missing entries for ranks: %s", missing_ranks)
+        return [s for s in statuses if s is not None]
+
    def execute_rpc(self, rpc_request: dict) -> tuple[object | None, bool]:
        """Execute an RPC request and indicate whether to reply."""
        method = rpc_request["method"]
@@ -747,6 +760,10 @@ class WorkerProc:
        kwargs = rpc_request.get("kwargs", {})
        output_rank = rpc_request.get("output_rank")
        exec_all_ranks = rpc_request.get("exec_all_ranks", False)
+        collect_rank_status = rpc_request.get("collect_rank_status", False)
+
+        if collect_rank_status and not exec_all_ranks:
+            raise ValueError("collect_rank_status requires exec_all_ranks=True so all ranks enter the status gather")

        should_execute = exec_all_ranks or output_rank is None or output_rank == self.gpu_id
        should_reply = (output_rank is None or output_rank == self.gpu_id) and self.result_mq is not None
@@ -754,13 +771,52 @@ class WorkerProc:
        if not should_execute:
            return None, False

+        result = None
+        rpc_exception: Exception | None = None
+        status: dict[str, Any] = {
+            "rank": self.gpu_id,
+            "ok": True,
+            "error": None,
+            "error_type": None,
+            "traceback": None,
+            "bool_result": None,
+        }
+
        try:
            # Use execute_method from WorkerWrapperBase for consistent method resolution
            result = self.worker.execute_method(method, *args, **kwargs)
-            return result, should_reply
        except Exception as e:
            logger.error(f"Error executing RPC: {e}", exc_info=True)
-            raise e
+            rpc_exception = e
+            status.update(
+                {
+                    "ok": False,
+                    "error": str(e),
+                    "error_type": type(e).__name__,
+                    "traceback": traceback.format_exc(),
+                }
+            )
+
+        if isinstance(result, bool):
+            status["bool_result"] = result
+
+        if collect_rank_status:
+            rank_statuses = self._gather_rpc_rank_statuses(status)
+            if should_reply:
+                return (
+                    {
+                        "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+                        "method": method,
+                        "result": result,
+                        "rank_statuses": rank_statuses,
+                    },
+                    True,
+                )
+            return None, False
+
+        if rpc_exception is not None:
+            raise rpc_exception
+        return result, should_reply

    def worker_busy_loop(self) -> None:
        """Main busy loop for Multiprocessing Workers."""
@@ -801,7 +857,7 @@ class WorkerProc:
                except Exception as e:
                    logger.error(f"Error processing RPC: {e}", exc_info=True)
                    if self.result_mq is not None:
-                        self.return_result(DiffusionOutput.from_exception(e))
+                        self.return_result({"status": "error", "error": str(e)})

            elif isinstance(msg, dict) and msg.get("type") == "shutdown":
                logger.info("Worker %s: Received shutdown message", self.gpu_id)
@@ -114,6 +114,7 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):
        self._listener_ready = threading.Event()
        self._local_buffers: dict[str, Any] = {}
        self._local_buffers_lock = threading.Lock()
+        self._put_lock = threading.Lock()
        self._req_local = threading.local()
        self._worker_local = threading.local()
        self._last_ttl_check: float = _time_mod.monotonic()
@@ -382,6 +383,14 @@ class MooncakeTransferEngineConnector(OmniConnectorBase):

        put_key = self._make_key(put_key, from_stage, to_stage)

+        # Serialize concurrent put() calls on the same connector to prevent
+        # races in the alloc -> copy -> store-metadata flow at the Mooncake
+        # C++ engine level.
+        with self._put_lock:
+            return self._put_impl(put_key, data)
+
+    def _put_impl(self, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]:
+        """Internal put implementation, called under _put_lock."""
        try:
            src_addr = 0
            size = 0
@@ -2480,6 +2480,28 @@ class AsyncOmniEngine:
                logger.exception("[AsyncOmniEngine] Failed to close OmniCoordinatorRuntime during shutdown")
            self._coordinator_runtime = None

+        # ── Release CuMem allocator memory pool ──────────────────────────────
+        # When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory
+        # Management) allocator holds model weights in a singleton memory pool
+        # that lives in the parent process.  Killing the engine-core subprocess
+        # does NOT release this pool — the weights stay resident on the GPU
+        # and can cause CUDA OOM for subsequent engine instances (especially
+        # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB).
+        #
+        # Discard mode (level=2) is correct at shutdown: there is no benefit to
+        # keeping a CPU backup when the engine is being torn down.
+        try:
+            from vllm.device_allocator.cumem import CuMemAllocator, cumem_available
+
+            if cumem_available:
+                allocator = CuMemAllocator.get_instance()
+                # Sleep at level 2 discards all pool memory from the GPU
+                # without creating CPU backups — cheapest and fastest.
+                allocator.sleep()
+                logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown")
+        except Exception:
+            pass
+
    def _try_shutdown(self, *args, **kwargs) -> None:
        try:
            self.shutdown()
@@ -122,9 +122,11 @@ class OmniRequestState(RequestState):
        *args,
        **kwargs,
    ):
+        # arrival_time is always passed as a keyword argument by the sole caller
+        # (from_new_request).  The *args fallback via positional index was removed
+        # because upstream RequestState.__init__ reorders parameters between
+        # releases, making args[12] fragile.
        arrival_time = kwargs.get("arrival_time")
-        if arrival_time is None and len(args) > 12:
-            arrival_time = args[12]
        super().__init__(*args, **kwargs)
        self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0))
        # Omni-specific: multimodal output accumulation
@@ -669,6 +671,14 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
        native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None
        previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0

+        # NOTE: We pass ``None`` for  *iteration_stats* to the parent so that
+        # the upstream's ``_update_stats_from_output`` logs stats via its own
+        # ``RequestStateStats`` (req_state.stats) without touching the omni
+        # ``IterationStats`` object.  Below we re-apply the same update using
+        # *native_stats* (the omni-specific ``RequestStateStats`` attached to
+        # ``OmniRequestState``), which correctly feeds the omni metrics path.
+        # If the upstream adds new logic in ``_update_stats_from_output`` (e.g.
+        # additional stat fields), this override must be kept in sync.
        super()._update_stats_from_output(
            req_state,
            engine_core_output,
@@ -4,7 +4,7 @@ import argparse
 import typing

 from vllm.entrypoints.cli.types import CLISubcommand
-from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
+from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG

 from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase

@@ -23,7 +23,7 @@ def main():
        if "VLLM_LOGGING_COLOR" not in os.environ:
            os.environ["VLLM_LOGGING_COLOR"] = "1"

-        from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup
+        from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup

        import vllm_omni.entrypoints.cli.benchmark.main
        import vllm_omni.entrypoints.cli.serve
@@ -17,7 +17,7 @@ from typing import Any
 import uvloop
 from vllm.entrypoints.cli.types import CLISubcommand
 from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
-from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
+from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm.logger import init_logger

 from vllm_omni.entrypoints.cli.logo import log_logo
@@ -34,7 +34,6 @@ from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
 from vllm.entrypoints.chat_utils import load_chat_template
 from vllm.entrypoints.launcher import serve_http, terminate_if_errored
-from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer
 from vllm.entrypoints.openai.api_server import build_app as build_openai_app
 from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server
@@ -56,10 +55,7 @@ from vllm.entrypoints.openai.engine.protocol import (
 )
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
-from vllm.entrypoints.openai.orca_metrics import metrics_header
 from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
-from vllm.entrypoints.openai.server_utils import get_uvicorn_log_config
-from vllm.entrypoints.openai.utils import validate_json_request
 from vllm.entrypoints.pooling.classify.serving import ServingClassification
 from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding
 from vllm.entrypoints.pooling.pooling.serving import ServingPooling
@@ -71,6 +67,16 @@ from vllm.entrypoints.serve.disagg.serving import ServingTokens
 from vllm.entrypoints.serve.instrumentator.basic import base
 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
+from vllm.entrypoints.serve.utils.api_utils import (
+    load_aware_call,
+    process_lora_modules,
+    validate_json_request,
+    with_cancellation,
+)
+from vllm.entrypoints.serve.utils.error_response import create_error_response
+from vllm.entrypoints.serve.utils.orca_metrics import metrics_header
+from vllm.entrypoints.serve.utils.request_logger import RequestLogger
+from vllm.entrypoints.serve.utils.server_utils import get_uvicorn_log_config
 from vllm.entrypoints.speech_to_text.realtime.serving import OpenAIServingRealtime
 from vllm.entrypoints.speech_to_text.transcription.serving import (
    OpenAIServingTranscription,
@@ -78,12 +84,6 @@ from vllm.entrypoints.speech_to_text.transcription.serving import (
 from vllm.entrypoints.speech_to_text.translation.serving import (
    OpenAIServingTranslation,
 )
-from vllm.entrypoints.utils import (
-    create_error_response,
-    load_aware_call,
-    process_lora_modules,
-    with_cancellation,
-)
 from vllm.logger import init_logger
 from vllm.tasks import POOLING_TASKS
 from vllm.tool_parsers import ToolParserManager
@@ -69,8 +69,8 @@ from vllm.entrypoints.openai.parser.harmony_utils import (
    parse_chat_output,
 )
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
-from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls
-from vllm.entrypoints.utils import should_include_usage
+from vllm.entrypoints.serve.utils.api_utils import should_include_usage
+from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls
 from vllm.inputs import PromptType
 from vllm.logger import init_logger
 from vllm.outputs import RequestOutput
@@ -843,6 +843,27 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
            if value is not None and hasattr(obj, key):
                setattr(obj, key, value)

+    def _should_check_for_unstreamed_tool_arg_tokens(
+        self,
+        delta_message: Any,
+        output: Any,
+    ) -> bool:
+        """Check whether the streaming generator should flush unstreamed
+        tool-arg tokens at finish-time.
+
+        This method was moved from OpenAIServingChat to the tool-parser layer
+        by upstream commit 9affc17a05.  Omni's independently-maintained
+        ``chat_completion_stream_generator`` still calls it, so we keep a
+        local copy with the pre-9affc17a05 semantics.
+        """
+        return (
+            output.finish_reason is not None
+            and self.enable_auto_tools
+            and self.tool_parser is not None
+            and delta_message is not None
+            and delta_message.tool_calls
+        )
+
    def _build_sampling_params_list_from_request(
        self,
        request: ChatCompletionRequest,
@@ -2794,6 +2815,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                    if isinstance(ar_text, str) and ar_text.strip():
                        cot_output = ar_text

+        req_out = getattr(result, "request_output", None)
+        if req_out:
+            prompt_obj = getattr(req_out, "prompt", None)
+            if isinstance(prompt_obj, dict):
+                extra = prompt_obj.get("extra", {})
+                if isinstance(extra, dict):
+                    ar_text = extra.get("ar_generated_text")
+                    if isinstance(ar_text, str) and ar_text.strip():
+                        cot_output = ar_text
+
        return self._flatten_diffusion_images(images), stage_durations, peak_memory_mb, cot_output

    async def _stream_diffusion_image_chunks(
@@ -87,29 +87,17 @@ from vllm.config.multimodal import (
    ImageDummyOptions,
    VideoDummyOptions,
 )
+from vllm.inputs import ModalityData, MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMRoPE, SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.utils import (
+    _merge_multimodal_embeddings as merge_multimodal_embeddings,
+)
 from vllm.model_executor.models.utils import (
    flatten_bn,
    init_vllm_registered_model,
    maybe_prefix,
 )
-
-try:
-    from vllm.model_executor.models.utils import merge_multimodal_embeddings
-except ImportError:
-    from vllm.model_executor.models.utils import (
-        _merge_multimodal_embeddings as merge_multimodal_embeddings,
-    )
-try:
-    from vllm.utils import flatten_2d_lists
-except ImportError:
-
-    def flatten_2d_lists(nested):
-        """vLLM >=0.18 removed vllm.utils.flatten_2d_lists; keep MiniCPM-o vision path working."""
-        return [x for row in nested for x in row]
-
-
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import (
    ImageItem,
@@ -117,17 +105,6 @@ from vllm.multimodal.inputs import (
    MultiModalKwargsItems,
    NestedTensors,
 )
-
-# vllm >=0.21 moved ModalityData / MultiModalDataDict from
-# vllm.multimodal.inputs to vllm.inputs (still re-exported from
-# vllm.multimodal.parse). Fall back to the old location for older vllm.
-try:
-    from vllm.inputs import ModalityData, MultiModalDataDict
-except ImportError:
-    from vllm.multimodal.inputs import (  # type: ignore[no-redef]
-        ModalityData,
-        MultiModalDataDict,
-    )
 from vllm.multimodal.parse import (
    AudioItem,
    AudioProcessorItems,
@@ -141,37 +118,19 @@ from vllm.multimodal.parse import (
    VideoProcessorItems,
 )
 from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    PromptReplacement,
    PromptUpdate,
    PromptUpdateDetails,
 )
-
-try:
-    from vllm.multimodal.processing import BaseDummyInputsBuilder
-except ImportError:
-    from vllm.multimodal.profiling import BaseDummyInputsBuilder
 from vllm.sequence import IntermediateTensors


 def encode_tokens(tokenizer, prompt: str) -> list[int]:
-    """Tokenize ``prompt`` without adding special tokens.
-
-    Prefer vllm's ``vllm.transformers_utils.tokenizer.encode_tokens`` when
-    available; falls back to the equivalent ``tokenizer.encode`` call so we
-    keep working across vllm versions that do not ship that helper.
-
-    The import is performed lazily inside the function so the docs builder
-    (griffe / mkdocs strict mode) does not try to statically resolve the
-    upstream alias, which can fail across vllm releases that rename or move
-    the helper.
-    """
-    try:
-        from vllm.transformers_utils.tokenizer import encode_tokens as _impl
-    except ImportError:
-        return tokenizer.encode(prompt, add_special_tokens=False)
-    return _impl(tokenizer, prompt)
+    """Tokenize ``prompt`` without adding special tokens."""
+    return tokenizer.encode(prompt, add_special_tokens=False)


 from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -2855,7 +2814,7 @@ class MiniCPMO45OmniLLMProcessingInfo(BaseProcessingInfo):
        if self.ctx.tokenizer is not None:
            return self.ctx.tokenizer
        if not hasattr(self, "_lazy_tokenizer") or self._lazy_tokenizer is None:
-            from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
+            from vllm.tokenizers import cached_tokenizer_from_config

            self._lazy_tokenizer = cached_tokenizer_from_config(self.ctx.model_config)
        return self._lazy_tokenizer
@@ -34,7 +34,10 @@ from vllm.config import VllmConfig
 from vllm.logger import init_logger

 from vllm_omni.model_executor.models.output_templates import OmniOutput
-from vllm_omni.model_executor.models.utils import reinit_rotary_inv_freq
+from vllm_omni.model_executor.models.utils import (
+    reinit_rotary_inv_freq,
+    transformers_keys_to_ignore_compat,
+)

 logger = init_logger(__name__)

@@ -193,7 +196,10 @@ class MossTTSNanoForGeneration(nn.Module):
        logger.info("Loading MOSS-TTS-Nano LM from %s (dtype=%s)", self.model_path, tts_dtype)
        from transformers import AutoModelForCausalLM

-        with _hf_load_without_tp_warmup():
+        # Two orthogonal trust_remote_code transformers fixes:
+        #  - _hf_load_without_tp_warmup: tp_plan=None warmup crash (transformers 5.8.x)
+        #  - transformers_keys_to_ignore_compat: keys list-vs-set (transformers 5.9)
+        with _hf_load_without_tp_warmup(), transformers_keys_to_ignore_compat():
            lm = AutoModelForCausalLM.from_pretrained(
                self.model_path,
                trust_remote_code=True,
@@ -233,7 +239,10 @@ class MossTTSNanoForGeneration(nn.Module):
        logger.info("Loading MOSS-Audio-Tokenizer-Nano from %s", codec_path)
        from transformers import AutoModel

-        with _hf_load_without_tp_warmup():
+        # Two orthogonal trust_remote_code transformers fixes:
+        #  - _hf_load_without_tp_warmup: tp_plan=None warmup crash (transformers 5.8.x)
+        #  - transformers_keys_to_ignore_compat: keys list-vs-set (transformers 5.9)
+        with _hf_load_without_tp_warmup(), transformers_keys_to_ignore_compat():
            audio_tokenizer = AutoModel.from_pretrained(
                codec_path,
                trust_remote_code=True,
@@ -5,7 +5,7 @@ import torch.nn.functional as F
 from torch import nn
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.layers.fused_moe.layer import MoERunner
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.model_executor.models.qwen3_moe import (
@@ -149,14 +149,14 @@ class Qwen3MoeForCausalLM(_BaseQwen3MoeForCausalLM):
        # Set MoE hyperparameters for individual experts
        self.expert_weights = []

-        self.moe_layers: list[FusedMoE] = []
+        self.moe_layers: list[MoERunner] = []
        example_layer = None
        for layer in self.model.layers:
            if isinstance(layer, PPMissingLayer):
                continue

            assert isinstance(layer, Qwen3MoeDecoderLayer)
-            if isinstance(layer.mlp, FusedMoE):
+            if isinstance(layer.mlp, MoERunner):
                example_layer = layer.mlp
                self.moe_layers.append(layer.mlp)

@@ -903,12 +903,12 @@ class Qwen3TTSPromptEmbedsBuilder:

            * ``talker_prompt``: ``[prompt_len, hidden]`` prefill embedding.
            * ``trailing_text_hidden``: ``[T, hidden]`` queue of text-side
-                embeddings consumed one-per-decode-step (streaming mode) or a
-                single pad row (non-streaming).
+                    embeddings consumed one-per-decode-step (streaming mode)
+                    or a single pad row (non-streaming).
            * ``ref_code_len``: Number of ref codec frames consumed when the
-                Base task ran in in-context mode (``None`` otherwise).
+                    Base task ran in in-context mode (``None`` otherwise).
            * ``ref_code``: The ``[T, Q]`` ref codec tensor used for ICL
-                (``None`` outside Base / ICL).
+                    (``None`` outside Base / ICL).
        """
        config = self._config
        talker_config = self._talker_config
@@ -1,9 +1,65 @@
 from collections.abc import Callable
+from contextlib import contextmanager

 import torch
 from vllm.model_executor.models.utils import maybe_prefix


+@contextmanager
+def transformers_keys_to_ignore_compat():
+    """Make ``trust_remote_code`` weight loading robust to the transformers 5.9
+    ``_keys_to_ignore_on_load_unexpected`` list-vs-set change.
+
+    transformers 5.9 rewrote ``PreTrainedModel._adjust_missing_and_unexpected_keys``
+    from ``(attr or []) + patterns`` (list concatenation) to
+    ``(attr or set()) | patterns`` (set union). Remote-code models such as
+    ``OpenMOSS-Team/MOSS-TTS-Nano`` still declare
+    ``_keys_to_ignore_on_load_unexpected`` as a *list*, so ``list | set`` raises
+    ``TypeError: unsupported operand type(s) for |: 'list' and 'set'`` and the
+    engine core dies during model load.
+
+    Wrap any ``from_pretrained(..., trust_remote_code=True)`` call whose remote
+    code may declare the attribute as a list. The guard keeps such models
+    loadable regardless of which transformers version is installed, while
+    preserving the model's ignore patterns.
+    """
+    try:
+        from transformers.modeling_utils import PreTrainedModel
+    except Exception:  # pragma: no cover - transformers is always present here
+        yield
+        return
+
+    orig = getattr(PreTrainedModel, "_adjust_missing_and_unexpected_keys", None)
+    if orig is None:
+        yield
+        return
+
+    def _wrapper(self, *args, **kwargs):
+        try:
+            return orig(self, *args, **kwargs)
+        except TypeError as exc:
+            if "unsupported operand type" not in str(exc):
+                raise
+            attr = getattr(self, "_keys_to_ignore_on_load_unexpected", None)
+            if not isinstance(attr, (list, tuple)):
+                raise
+            # transformers >=5.9 combines patterns with ``set | ...``. The
+            # ignore-pattern combine runs before any mutation of ``loading_info``
+            # (see PreTrainedModel._adjust_missing_and_unexpected_keys), so it is
+            # safe to coerce the list to a set and re-run the original once.
+            self._keys_to_ignore_on_load_unexpected = set(attr)
+            try:
+                return orig(self, *args, **kwargs)
+            finally:
+                self._keys_to_ignore_on_load_unexpected = attr
+
+    PreTrainedModel._adjust_missing_and_unexpected_keys = _wrapper
+    try:
+        yield
+    finally:
+        PreTrainedModel._adjust_missing_and_unexpected_keys = orig
+
+
 def add_prefix_to_loaded_weights(weights: set[str], prefix: str) -> set[str]:
    """
    Add a prefix to the names of the loaded weights.
@@ -6,6 +6,7 @@ and also outputs sampled tokens.

 from __future__ import annotations

+import gc
 from collections.abc import Mapping
 from contextlib import nullcontext
 from copy import copy
@@ -268,6 +269,63 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
        self._capture_talker_mtp_graphs()
        return result

+    def shutdown(self) -> None:
+        """Release omni-specific GPU resources before upstream shutdown.
+
+        Order of operations (must match upstream's expectation):
+          1. Unfreeze Python GC so model weights are collected immediately
+             when self.model is set to None (upstream Worker.init_device
+             calls gc.freeze() / freeze_gc_heap()).
+          2. Destroy omni-specific CUDA graphs (talker MTP) so references to
+             model parameters are released before self.model = None.
+          3. Clear GPU-side buffers (input_ids, inputs_embeds) and per-request
+             caches that may hold GPU tensor references.
+          4. Call CUDAGraphWrapper.clear_all_graphs() unconditionally (not just
+             on ROCm) to ensure all CUDA graphs including talker MTP are
+             released before model weight teardown.
+          5. Call BreakableCUDAGraphWrapper.clear_all_graphs() as well, to
+             match the upstream ROCm-only pattern but also protect CUDA.
+          6. Delegate to upstream GPUModelRunner.shutdown() which sets
+             self.model = None, clears KV caches, resets workspace, etc.
+
+        This prevents abrupt GPU memory release during EngineCore subprocess
+        exit that can trigger GPU OOM signals when the parent process
+        concurrently cleans up its own GPU state.
+        """
+        # 1. Unfreeze GC so model weights and GPU tensors are collected
+        #    immediately when references are dropped (upstream Worker.shutdown
+        #    also does this before any teardown).
+        gc.unfreeze()
+
+        # 2. Destroy talker MTP CUDA graph wrapper to release captured graphs.
+        if hasattr(self, "talker_mtp") and self.talker_mtp is not None:
+            self.talker_mtp = None
+        self.has_talker_mtp = False
+
+        # 3. Clear GPU-side buffers (small tensors, but every MiB helps).
+        if hasattr(self, "input_ids") and self.input_ids is not None:
+            self.input_ids = None
+        if hasattr(self, "inputs_embeds") and self.inputs_embeds is not None:
+            self.inputs_embeds = None
+
+        # 4. Clear per-request caches that may hold GPU tensor references.
+        if hasattr(self, "_downstream_payload_cache"):
+            self._downstream_payload_cache.clear()
+        if hasattr(self, "model_intermediate_buffer"):
+            self.model_intermediate_buffer.clear()
+
+        # 5. Release all CUDA graphs unconditionally (upstream only does this
+        #    on ROCm; on CUDA the graphs are only freed by Python GC during
+        #    interpreter shutdown, which is too late to prevent memory spikes).
+        from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper
+        from vllm.compilation.cuda_graph import CUDAGraphWrapper
+
+        CUDAGraphWrapper.clear_all_graphs()
+        BreakableCUDAGraphWrapper.clear_all_graphs()
+
+        # 6. Delegate to upstream shutdown (model = None, KV caches, workspace).
+        super().shutdown()
+
    def _capture_talker_mtp_graphs(self) -> None:
        from vllm.compilation.cuda_graph import CUDAGraphWrapper

@@ -9,12 +9,12 @@ from vllm.utils.mem_utils import MemorySnapshot, format_gib
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.utils import report_usage_stats
 from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
-from vllm.v1.worker.utils import request_memory
 from vllm.v1.worker.workspace import init_workspace_manager

 from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask
 from vllm_omni.worker.base import OmniGPUWorkerBase
 from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+from vllm_omni.worker.memory_utils import request_memory_tolerant
 from vllm_omni.worker.mixins import OmniWorkerMixin

 logger = init_logger(__name__)
@@ -84,7 +84,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):

            # take current memory snapshot
            self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
-            self.requested_memory = request_memory(init_snapshot, self.cache_config)
+            self.requested_memory = request_memory_tolerant(init_snapshot, self.cache_config)
            logger.debug("worker init memory snapshot: %r", self.init_snapshot)
            logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory))
        else:
@@ -718,6 +718,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin

                cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np)
                self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
+                self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1])
                self.query_start_loc.copy_to_gpu()

                self.input_batch.block_table.commit_block_table(num_reqs_padded)
@@ -9,11 +9,11 @@ from vllm.utils.mem_utils import MemorySnapshot, format_gib
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.utils import report_usage_stats
 from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
-from vllm.v1.worker.utils import request_memory
 from vllm.v1.worker.workspace import init_workspace_manager

 from vllm_omni.worker.base import OmniGPUWorkerBase
 from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner
+from vllm_omni.worker.memory_utils import request_memory_tolerant
 from vllm_omni.worker.mixins import OmniWorkerMixin

 logger = init_logger(__name__)
@@ -83,7 +83,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):

            # take current memory snapshot
            self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
-            self.requested_memory = request_memory(init_snapshot, self.cache_config)
+            self.requested_memory = request_memory_tolerant(init_snapshot, self.cache_config)
            logger.debug("worker init memory snapshot: %r", self.init_snapshot)
            logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory))
        else:
@@ -1006,6 +1006,7 @@ class OmniGPUModelRunner(GPUModelRunner):

                cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np)
                self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
+                self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1])
                self.query_start_loc.copy_to_gpu()

                self.input_batch.block_table.commit_block_table(num_reqs_padded)
@@ -0,0 +1,55 @@
+"""GPU memory utilities for vLLM Omni workers.
+
+Includes a tolerant version of the upstream request_memory() that handles
+multi-stage GPU sharing by capping the memory budget to available free
+memory instead of raising ValueError.
+"""
+
+from __future__ import annotations
+
+import math
+
+from vllm.config import CacheConfig
+from vllm.logger import init_logger
+from vllm.utils.mem_utils import MemorySnapshot, format_gib
+
+logger = init_logger(__name__)
+
+
+def request_memory_tolerant(
+    init_snapshot: MemorySnapshot,
+    cache_config: CacheConfig,
+) -> int:
+    """Calculate the amount of memory required for this stage.
+
+    Like upstream ``request_memory()`` but tolerates multi-stage GPU sharing:
+    if ``free_memory < requested_memory`` (because another stage on the same
+    GPU has already consumed memory), caps the requested budget to the actual
+    free memory instead of raising ``ValueError``.  The downstream
+    ``OmniGPUWorkerBase.determine_available_memory()`` already does per-process
+    NVML accounting and correctly computes the KV cache budget regardless.
+
+    Logs a warning when the budget is capped so operators can detect
+    under-provisioned GPU memory.
+    """
+    requested_memory = math.ceil(init_snapshot.total_memory * cache_config.gpu_memory_utilization)
+
+    if init_snapshot.free_memory < requested_memory:
+        capped = init_snapshot.free_memory
+        logger.warning(
+            "Free memory on device %s (%s/%s GiB) on startup is less than "
+            "desired GPU memory utilization (%.2f, %s GiB). "
+            "Capping requested memory to available free memory (%s GiB). "
+            "This is expected when multiple Omni stages share a GPU; "
+            "the per-process NVML accounting in determine_available_memory() "
+            "will compute the correct KV cache budget.",
+            init_snapshot.device_,
+            format_gib(init_snapshot.free_memory),
+            format_gib(init_snapshot.total_memory),
+            cache_config.gpu_memory_utilization,
+            format_gib(requested_memory),
+            format_gib(capped),
+        )
+        return capped
+
+    return requested_memory
Author	SHA1	Message	Date
Zhou Taichang	7b837944a1	[Rebase] Rebase to vllm 0.23.0 (#4286 ) Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-06-13 16:43:02 +00:00
Chendi.Xue	a9ee92ddec	[XPU] update Dreamzero to support any HW (#4399 ) Signed-off-by: Chendi Xue <chendi.xue@intel.com>	2026-06-13 15:32:39 +00:00
Yuanheng Zhao	31886d6ece	[Doc][Recipe] Update CUDA verifications for inclusionAI/Ming-omni-tts-0.5B (#4324 ) Signed-off-by: Yuanheng <jonathan.zhaoyh@gmail.com>	2026-06-13 18:12:37 +08:00
Hongsheng Liu	fa94f41807	[Bugfix] Surface all-rank diffusion RPC failures (#4403 ) Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>	2026-06-13 15:42:16 +08:00