Revert "[BugFix] ignore mm data from stages to async omni (#954)"

This reverts commit 6fb4058566.

.buildkite/pipeline.yml

@@ -18,25 +18,11 @@ steps:
   #     queue: "ascend"
 
   - label: "Simple Unit Test"
-    depends_on: image-build
+    depends_on: ~
     commands:
-      - pytest -v -s tests/entrypoints/
-      - pytest -v -s tests/diffusion/cache/
-      - pytest -v -s tests/diffusion/lora/
-      - pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
-      - pytest -v -s tests/worker/
-      - pytest -v -s tests/distributed/omni_connectors/test_kv_flow.py
+      - ".buildkite/scripts/simple_test.sh"
     agents:
-      queue: "gpu_1_queue"
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
+      queue: "cpu_queue_premerge"
 
   - label: "Diffusion Model Test"
     timeout_in_minutes: 20
@@ -77,7 +63,6 @@ steps:
     depends_on: image-build
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -164,7 +149,7 @@ steps:
     timeout_in_minutes: 20
     depends_on: image-build
     commands:
-      - pytest -s -v tests/diffusion/test_diffusion_worker.py
+      - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -178,13 +163,34 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-
-  - label: "Benchmark Test"
+  - label: "Omni Model Test"
     timeout_in_minutes: 15
     depends_on: image-build
     commands:
+      - export VLLM_LOGGING_LEVEL=DEBUG
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/benchmarks/test_serve_cli.py
+      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Omni Model Test with H100"
+    timeout_in_minutes: 30
+    depends_on: image-build
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
+      - pytest -s -v tests/e2e/online_serving/test_async_omni.py
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -214,69 +220,12 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: "Omni Model Test"
-    timeout_in_minutes: 15
-    depends_on: image-build
-    commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-    - docker#v5.2.0:
-        image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-        always-pull: true
-        propagate-environment: true
-        environment:
-        - "HF_HOME=/fsx/hf_cache"
-        volumes:
-        - "/fsx/hf_cache:/fsx/hf_cache"
-
-  # - label: "Omni Model Test with H100"
-  #   timeout_in_minutes: 30
-  #   depends_on: image-build
-  #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-  #     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-  #     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
-  #     - pytest -s -v tests/e2e/online_serving/test_async_omni.py
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 2
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
-
   - label: "Diffusion Image Edit Test with H100 (1 GPU)"
     timeout_in_minutes: 20
     depends_on: image-build
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+      - pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -305,38 +254,3 @@ steps:
                 hostPath:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
-
-  # - label: "Bagel Text2Img Model Test with H100"
-  #   timeout_in_minutes: 30
-  #   depends_on: image-build
-  #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 1
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate

.buildkite/scripts/hardware_ci/run-amd-test.sh

@@ -116,9 +116,6 @@ if [[ $commands == *"--shard-id="* ]]; then
         --shm-size=16gb \
         --group-add "$render_gid" \
         --rm \
-        -e MIOPEN_DEBUG_CONV_DIRECT=0 \
-        -e MIOPEN_DEBUG_CONV_GEMM=0 \
-        -e VLLM_ROCM_USE_AITER=1 \
         -e HIP_VISIBLE_DEVICES="${GPU}" \
         -e HF_TOKEN \
         -e AWS_ACCESS_KEY_ID \
@@ -151,9 +148,6 @@ else
           --shm-size=16gb \
           --group-add "$render_gid" \
           --rm \
-          -e MIOPEN_DEBUG_CONV_DIRECT=0 \
-          -e MIOPEN_DEBUG_CONV_GEMM=0 \
-          -e VLLM_ROCM_USE_AITER=1 \
           -e HF_TOKEN \
           -e AWS_ACCESS_KEY_ID \
           -e AWS_SECRET_ACCESS_KEY \

.buildkite/scripts/simple_test.sh

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)"
+UV_BIN="uv"
+PYTHON_BOOTSTRAP="${PYTHON:-python3}"
+VENV_DIR="${VENV_DIR:-${REPO_ROOT}/.venv-simple-test}"
+TARGET_PY_VERSION="3.12"
+
+if ! command -v "${PYTHON_BOOTSTRAP}" >/dev/null 2>&1; then
+  PYTHON_BOOTSTRAP="python"
+fi
+
+cd "${REPO_ROOT}"
+
+# Ensure pip-installed scripts (including uv) are on PATH
+PYTHON_SCRIPTS="$(${PYTHON_BOOTSTRAP} -c 'import sysconfig; print(sysconfig.get_path("scripts"))')"
+USER_BASE="$(${PYTHON_BOOTSTRAP} -m site --user-base 2>/dev/null || true)"
+PATH="${PYTHON_SCRIPTS}:${PATH}"
+if [[ -n "${USER_BASE}" ]]; then
+  PATH="${USER_BASE}/bin:${PATH}"
+fi
+
+if ! command -v "${UV_BIN}" >/dev/null 2>&1; then
+  "${PYTHON_BOOTSTRAP}" -m pip install --upgrade pip
+  "${PYTHON_BOOTSTRAP}" -m pip install uv
+  hash -r
+fi
+
+TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
+if [[ -z "${TARGET_PY_PATH}" ]]; then
+  "${UV_BIN}" python install "${TARGET_PY_VERSION}"
+  TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
+fi
+
+if [[ -z "${TARGET_PY_PATH}" ]]; then
+  echo "Failed to locate python ${TARGET_PY_VERSION}" >&2
+  exit 1
+fi
+
+if [[ ! -d "${VENV_DIR}" ]]; then
+  "${UV_BIN}" venv --python "${TARGET_PY_PATH}" "${VENV_DIR}"
+fi
+
+VENV_PYTHON="${VENV_DIR}/bin/python"
+[[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; }
+
+"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.14.0
+"${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
+
+"${VENV_PYTHON}" -m pytest -v -s tests/entrypoints/
+"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/
+"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/lora/
+"${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
+"${VENV_PYTHON}" -m pytest -v -s tests/worker/
+"${VENV_PYTHON}" -m pytest -v -s tests/distributed/omni_connectors/test_kv_flow.py

.buildkite/test-amd.yaml

@@ -8,32 +8,14 @@ steps:
   grade: Blocking
   commands:
     - export GPU_ARCHS=gfx942
+    - export MIOPEN_DEBUG_CONV_DIRECT=0
+    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - export VLLM_ROCM_USE_AITER=1
+    - export VLLM_ROCM_USE_AITER_MHA=1
+    - export VLLM_ROCM_USE_AITER_LINEAR=0
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
 
-- label: "Diffusion Images API LoRA E2E"
-  timeout_in_minutes: 20
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
-
-- label: "Diffusion Model CPU offloading Test"
-  timeout_in_minutes: 20
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
-
 - label: "Diffusion Cache Backend Test"
   timeout_in_minutes: 15
   agent_pool: mi325_1
@@ -44,30 +26,25 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export MIOPEN_DEBUG_CONV_DIRECT=0
+    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - export VLLM_ROCM_USE_AITER=1
+    - export VLLM_ROCM_USE_AITER_MHA=1
+    - export VLLM_ROCM_USE_AITER_LINEAR=0
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
 
-- label: "Diffusion Sequence Parallelism Test"
-  timeout_in_minutes: 20
+- label: "Diffusion Parallelism Test"
+  timeout_in_minutes: 15
   agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export MIOPEN_DEBUG_CONV_DIRECT=0
+    - export MIOPEN_DEBUG_CONV_GEMM=0
     - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
 
-- label: "Diffusion Tensor Parallelism Test"
-  timeout_in_minutes: 20
-  agent_pool: mi325_2
-  depends_on: amd-build
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
-
 - label: "Diffusion GPU Worker Test"
   timeout_in_minutes: 20
   agent_pool: mi325_2
@@ -75,7 +52,9 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - pytest -s -v tests/diffusion/test_diffusion_worker.py
+    - export MIOPEN_DEBUG_CONV_DIRECT=0
+    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
 
 - label: "Omni Model Test Qwen2-5-Omni"
   timeout_in_minutes: 15
@@ -87,6 +66,12 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export MIOPEN_DEBUG_CONV_DIRECT=0
+    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - export VLLM_ROCM_USE_AITER=1
+    - export VLLM_ROCM_USE_AITER_MHA=1
+    - export VLLM_ROCM_USE_AITER_LINEAR=0
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
 
 - label: "Omni Model Test Qwen3-Omni"
@@ -98,10 +83,9 @@ steps:
   commands:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-    - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
-    - pytest -s -v tests/e2e/online_serving/test_async_omni.py
+    - export MIOPEN_DEBUG_CONV_DIRECT=0
+    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py
 
 - label: "Diffusion Image Edit Test"
   timeout_in_minutes: 15
@@ -113,4 +97,10 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+    - export MIOPEN_DEBUG_CONV_DIRECT=0
+    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - export VLLM_ROCM_USE_AITER=1
+    - export VLLM_ROCM_USE_AITER_MHA=1
+    - export VLLM_ROCM_USE_AITER_LINEAR=0
+    - export VLLM_ROCM_USE_AITER_RMSNORM=0
+    - pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py

README.md

@@ -15,7 +15,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *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).
 - [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.
 

docker/Dockerfile.ci

@@ -1,21 +1,20 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.15.0
+ARG VLLM_BASE_TAG=v0.14.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 []

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.15.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.14.0
 FROM ${BASE_IMAGE} AS final
 
 ARG COMMON_WORKDIR=/app

docs/.nav.yml

@@ -7,7 +7,6 @@ nav:
   - Serving:
     - OpenAI-Compatible API:
       - Image Generation: serving/image_generation_api.md
-      - Image Edit: serving/image_edit_api.md
   - Examples:
     - examples/README.md
     - Offline Inference:

docs/cli/README.md

@@ -22,21 +22,3 @@ If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --stage-configs-path /path/to/stage_configs_file
 ```
-
-
-## bench
-
-Run benchmark tests for online serving throughput.
-Available Commands:
-
-```bash
-vllm bench serve --omni \
-    --model Qwen/Qwen2.5-Omni-7B \
-    --host server-host \
-    --port server-port \
-    --random-input-len 32 \
-    --random-output-len 4  \
-    --num-prompts  5
-```
-
-See [vllm bench serve](./bench/serve.md) for the full reference of all available arguments.

docs/cli/bench/serve.md

@@ -1,359 +0,0 @@
-# vLLM-Omni Benchmark CLI Guide
-The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
-
-## Notes
-We currently only support using the "openai-chat-omni" backend.
-
-## Basic Parameter Description
-You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
-- `--omni`  
-  Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
-
-- `--backend`  
-  Specify the backend adapter as openai-chat-omni, using OpenAI Chat compatible API behavior as the protocol. Currently only openai-chat-omni is supported.
-
-- `--model`  
-  The model identifier to load, filled according to the models supported by vLLM-Omni.
-
-- `--endpoint`  
-  The API endpoint exposed externally, to which clients send their requests.
-
-- `--dataset-name`  
-  The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio).
-
-- `--num-prompts`  
-  The total number of requests to send, an integer.
-
-- `--max-concurrency`  
-  "Maximum number of concurrent requests. This can be used "
-        "to help simulate an environment where a higher level component "
-        "is enforcing a maximum number of concurrent requests. While the "
-        "--request-rate argument controls the rate at which requests are "
-        "initiated, this argument will control how many are actually allowed "
-        "to execute at a time. This means that when used in combination, the "
-        "actual request rate may be lower than specified with --request-rate, "
-        "if the server is not processing requests fast enough to keep up."
-
-- `--request-rate`  
-  "Number of requests per second. If this is inf, "
-        "then all the requests are sent at time 0. "
-        "Otherwise, we use Poisson process or gamma distribution "
-        "to synthesize the request arrival times."
-
-- `--ignore-eos`  
-  "Set ignore_eos flag when sending the benchmark request."
-
-- `--metric-percentiles`  
-  Comma-separated list of percentiles for selected metrics. "
-        "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "
-        "Default value is \"99\"."
-        "Use \"--percentile-metrics\" to select metrics.
-
-- `--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". '
-
-- `--save-result`  
-Specify to save benchmark results to a json file
-
-- `--save-detailed`  
-"When saving the results, whether to include per request "
-        "information such as response, error, ttfs, tpots, etc."
-
-- `--result-dir`  
- "Specify directory to save benchmark json results."
-        "If not specified, results are saved in the current directory."
-
-- `--result-filename`  
-"Specify the filename to save benchmark json results."
-        "If not specified, results will be saved in "
-        "{label}-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
-
-- `--random-prefix-len`  
-  Number of fixed prefix tokens before the random context in a request.
-  The total input length is the sum of random-prefix-len and a random
-  context length sampled from [input_len * (1 - range_ratio),
-  input_len * (1 + range_ratio)].Only the random and random-mm modes
-  support this parameter.
-
-- `--random-input-len`  
-  Number of input tokens per request.Only the random and random-mm modes support this parameter.
-
-- `--random-output-len`  
-  Number of output tokens per request.Only the random and random-mm modes support this parameter.
-
-- `--random-range-ratio`  
-  Range ratio for sampling input/output length,
-  used only for random sampling. Must be in the range [0, 1) to define
-  a symmetric sampling range
-  [length * (1 - range_ratio), length * (1 + range_ratio)].
-  Only the random and random-mm modes support this parameter.
-
-- `--random-mm-base-items-per-request`  
-  Base number of multimodal items per request for random-mm.
-  Actual per-request count is sampled around this base using
-  --random-mm-num-mm-items-range-ratio.
-  Only the random-mm mode supports this parameter.
-
-- `--random-mm-limit-mm-per-prompt`  
-  Per-modality hard caps for items attached per request, e.g.
-  '{"image": 3, "video": 1, "audio": 1}'. The sampled per-request item
-  count is clamped to the sum of these limits. When a modality
-  reaches its cap, its buckets are excluded and probabilities are
-  renormalized.
-  Only the random-mm mode supports this parameter.
-
-- `--random-mm-num-mm-items-range-ratio`  
-  Range ratio r in [0, 1] for sampling items per request.
-  We sample uniformly from the closed integer range
-  [floor(n*(1-r)), ceil(n*(1+r))]
-  where n is the base items per request.
-  r=0 keeps it fixed; r=1 allows 0 items. The maximum is clamped
-  to the sum of per-modality limits from
-  --random-mm-limit-mm-per-prompt.
-  An error is raised if the computed min exceeds the max.
-  Only the random-mm mode supports this parameter.
-
-- `--random-mm-bucket-config`  
-  The bucket config is a dictionary mapping a multimodal item
-  sampling configuration to a probability.
-  Currently allows for 3 modalities: audio, images and videos.
-  A bucket key is a tuple of (height, width, num_frames)
-  The value is the probability of sampling that specific item.
-  Example:
-  --random-mm-bucket-config
-  "{(256, 256, 1): 0.5, (720, 1280, 16): 0.4, (0, 1, 5): 0.10}"
-  First item: images with resolution 256x256 w.p. 0.5
-  Second item: videos with resolution 720x1280 and 16 frames
-  Third item: audios with 1s duration and 5 channels w.p. 0.1
-  OBS.: If the probabilities do not sum to 1, they are normalized.
-  Only the random-mm mode supports this parameter
-
-## Usage Examples
-
-### Online Benchmark
-<details class="admonition abstract" markdown="1">
-<summary>Show more</summary>
-
-First start serving your model:
-
-```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni
-```
-
-Then run the benchmarking for sharegpt:
-
-```bash
-# download dataset
-# wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
-vllm bench serve \
-  --omni \
-  --port 43845 \
-  --model /home/models/Qwen/Qwen3-Omni-30B-A3B-Instruct \
-  --endpoint /v1/chat/completions \
-  --backend openai-chat-omni \
-  --num-prompts 2 \
-  --dataset-name sharegpt \
-  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
-  --percentile-metrics ttft,tpot,itl,e2el
-```
-If successful, you will see the following output:
-```text
-============ Serving Benchmark Result ============
-Successful requests:                     2
-Failed requests:                         0
-Benchmark duration (s):                  81.63
-Request throughput (req/s):              0.02
-Peak concurrent requests:                2.00
-----------------End-to-end Latency----------------
-Mean E2EL (ms):                          56966.13
-Median E2EL (ms):                        56966.13
-P99 E2EL (ms):                           81016.80
-================== Text Result ===================
-Total input tokens:                      36
-Total generated tokens:                  5926
-Output token throughput (tok/s):         72.60
-Peak output token throughput (tok/s):    103.00
-Peak concurrent requests:                2.00
-Total Token throughput (tok/s):          73.04
----------------Time to First Token----------------
-Mean TTFT (ms):                          124.76
-Median TTFT (ms):                        124.76
-P99 TTFT (ms):                           156.10
------Time per Output Token (excl. 1st token)------
-Mean TPOT (ms):                          481.30
-Median TPOT (ms):                        481.30
-P99 TPOT (ms):                           947.55
----------------Inter-token Latency----------------
-Mean ITL (ms):                           25.11
-Median ITL (ms):                         0.33
-P99 ITL (ms):                            25.17
-================== Audio Result ==================
-Total audio duration generated(s):       3.95
-Total audio frames generated:            94890
-Audio throughput(audio duration/s):      0.05
-==================================================
-```
-
-Or run the benchmarking for random:
-
-```bash
-vllm bench serve \
-  --omni \
-  --port 43845 \
-  --endpoint /v1/chat/completions \
-  --backend openai-chat-omni \
-  --model /home/models/Qwen/Qwen3-Omni-30B-A3B-Instruct \
-  --dataset-name random \
-  --num-prompts 2 \
-  --random-prefix-len 5 \
-  --random-input-len 10 \
-  --random-output-len 100 \
-  --percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf \
-  --ignore-eos
-```
-
-If successful, you will see the following output:
-
-```text
-============ Serving Benchmark Result ============
-Successful requests:                     2
-Failed requests:                         0
-Benchmark duration (s):                  24.35
-Request throughput (req/s):              0.08
-Peak concurrent requests:                2.00
-----------------End-to-end Latency----------------
-Mean E2EL (ms):                          22576.23
-Median E2EL (ms):                        22576.23
-P99 E2EL (ms):                           24205.72
-================== 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
-Peak concurrent requests:                2.00
-Total Token throughput (tok/s):          369.76
----------------Time to First Token----------------
-Mean TTFT (ms):                          125.16
-Median TTFT (ms):                        125.16
-P99 TTFT (ms):                           155.88
------Time per Output Token (excl. 1st token)------
-Mean TPOT (ms):                          5.01
-Median TPOT (ms):                        5.01
-P99 TPOT (ms):                           5.42
----------------Inter-token Latency----------------
-Mean ITL (ms):                           34.15
-Median ITL (ms):                         0.01
-P99 ITL (ms):                            376.19
-================== Audio Result ==================
-Total audio duration generated(s):       3.95
-Total audio frames generated:            94890
-Audio throughput(audio duration/s):      0.16
----------------Time to First Packet---------------
-Mean AUDIO_TTFP (ms):                    11756.89
-Median AUDIO_TTFP (ms):                  11756.89
-P99 AUDIO_TTFP (ms):                     20854.25
------------------Real Time Factor-----------------
-Mean AUDIO_RTF:                          3.75
-Median AUDIO_RTF:                        3.75
-P99 AUDIO_RTF:                           7.39
-==================================================
-```
-Notes:
-We use (audio generation time - first packet latency) / audio duration to calculate RTF.
-
-</details>
-
-### Multi-Modal Benchmark
-
-<details class="admonition abstract" markdown="1">
-<summary>Show more</summary>
-
-Benchmark the performance of multi-modal requests in vLLM-Omni.
-
-Generate synthetic image、video、audio inputs alongside random text prompts to stress-test vision models without external datasets.
-
-Notes:
-
-- Works only with online benchmark via the OpenAI backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
-
-Start the server (example):
-
-```bash
-vllm serve Qwen/Qwen2.5-Omni-7B --omni
-```
-
-It is recommended to use the flag `--ignore-eos` to simulate real responses. You can set the size of the output via the arg `random-output-len`.
-
-Then run the benchmarking script:
-```bash
-vllm bench serve \
-    --omni \
-  --dataset-name random-mm \
-  --port 40849 \
-  --model /home/models/Qwen/Qwen3-Omni-30B-A3B-Instruct \
-  --endpoint /v1/chat/completions \
-  --backend openai-chat-omni \
-  --request-rate 1 \
-  --num-prompts 1 \
-  --random-input-len 10 \
-  --random-range-ratio 0.0 \
-  --random-mm-base-items-per-request 2 \
-  --random-mm-num-mm-items-range-ratio 0 \
-  --random-mm-limit-mm-per-prompt '{"image":1,"video":1, "audio": 1}' \
-  --random-mm-bucket-config '{"(32, 32, 1)": 0.5, "(0, 1, 1)": 0.1, "(32, 32, 2)":0.4}' \
-  --ignore-eos \
-  --percentile-metrics ttft,tpot,itl \
-  --random-output-len 2 \
-  --extra_body '{"modalities": ["text"]}'
-```
-
-If successful, you will see the following output:
-
-```text
-============ Serving Benchmark Result ============
-Successful requests:                     1
-Failed requests:                         0
-Request rate configured (RPS):           1.00
-Benchmark duration (s):                  1.21
-Request throughput (req/s):              0.83
-Peak concurrent requests:                1.00
-================== Text Result ===================
-Total input tokens:                      10
-Total generated tokens:                  3
-Output token throughput (tok/s):         2.49
-Peak output token throughput (tok/s):    3.00
-Peak concurrent requests:                1.00
-Total Token throughput (tok/s):          10.77
----------------Time to First Token----------------
-Mean TTFT (ms):                          179.74
-Median TTFT (ms):                        179.74
-P99 TTFT (ms):                           179.74
------Time per Output Token (excl. 1st token)------
-Mean TPOT (ms):                          12.76
-Median TPOT (ms):                        12.76
-P99 TPOT (ms):                           12.76
----------------Inter-token Latency----------------
-Mean ITL (ms):                           12.76
-Median ITL (ms):                         12.76
-P99 ITL (ms):                            25.24
-================== Audio Result ==================
-Total audio duration generated(s):       0.00
-Total audio frames generated:            0
-Audio throughput(audio duration/s):      0.00
-==================================================
-```
-
-Behavioral notes:
-
-- If the requested base item count cannot be satisfied under the provided per-prompt limits, the tool raises an error rather than silently clamping.
-
-How sampling works:
-
-- Determine per-request item count k by sampling uniformly from the integer range defined by `--random-mm-base-items-per-request` and `--random-mm-num-mm-items-range-ratio`, then clamp k to at most the sum of per-modality limits.
-- For each of the k items, sample a bucket (H, W, T) according to the normalized probabilities in `--random-mm-bucket-config`, while tracking how many items of each modality have been added.
-- If a modality (e.g., image) reaches its limit from `--random-mm-limit-mm-per-prompt`, all buckets of that modality are excluded and the remaining bucket probabilities are renormalized before continuing.
-This should be seen as an edge case, and if this behavior can be avoided by setting `--random-mm-limit-mm-per-prompt` to a large number. Note that this might result in errors due to engine config `--limit-mm-per-prompt`.
-- The resulting request contains synthetic image data in `multi_modal_data` (OpenAI Chat format). When `random-mm` is used with the OpenAI Chat backend, prompts remain text and MM content is attached via `multi_modal_data`.
-</details>

docs/configuration/stage_configs.md

@@ -40,7 +40,7 @@ stage_args:
     engine_args: # Engine arguments for a certain engine
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
-      worker_type: ar # The specific worker used
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker # The specific worker used
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
       gpu_memory_utilization: 0.8 # The gpu memory allocation for the stage within a single chip
       enforce_eager: true  # Now we only support eager mode
@@ -66,7 +66,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -92,7 +92,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: generation
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -8,7 +8,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
@@ -34,7 +34,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -60,7 +60,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: generation
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true

docs/contributing/model/adding_diffusion_model.md

@@ -204,151 +204,7 @@ Key point for writing the example:
 
 + Save or display the generated results so users can validate the integration.
 
-## Step 5: TeaCache Coefficient Estimation (Optional)
-
-If your model supports TeaCache acceleration, you need to estimate the polynomial coefficients for optimal caching performance.
-
-### 5.1 Add Extractor Function
-
-First, implement an extractor function in `vllm_omni/diffusion/cache/teacache/extractors.py`. The extractor extracts the modulated input and defines how to run transformer blocks:
-
-```python
-def extract_your_model_context(
-    module: nn.Module,
-    hidden_states: torch.Tensor,
-    timestep: torch.Tensor,
-    **kwargs: Any,
-) -> CacheContext:
-    # 1. Preprocessing
-    temb = module.time_embed(timestep)
-
-    # 2. Extract modulated input (for cache decision)
-    modulated_input = module.transformer_blocks[0].norm1(hidden_states, temb)
-
-    # 3. Define transformer execution
-    def run_transformer_blocks():
-        h = hidden_states
-        for block in module.transformer_blocks:
-            h = block(h, temb=temb)
-        return (h,)
-
-    # 4. Define postprocessing
-    def postprocess(h):
-        return module.proj_out(module.norm_out(h, temb))
-
-    return CacheContext(
-        modulated_input=modulated_input,
-        hidden_states=hidden_states,
-        encoder_hidden_states=None,
-        temb=temb,
-        run_transformer_blocks=run_transformer_blocks,
-        postprocess=postprocess,
-    )
-```
-
-Register it in `EXTRACTOR_REGISTRY`:
-```python
-EXTRACTOR_REGISTRY = {
-    ...
-    "YourTransformer2DModel": extract_your_model_context,
-}
-```
-
-### 5.2 Add Adapter for Coefficient Estimation
-
-Add an adapter in `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py`:
-
-```python
-class YourModelAdapter:
-    @staticmethod
-    def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any:
-        # Load your pipeline
-        ...
-
-    @staticmethod
-    def get_transformer(pipeline: Any) -> tuple[Any, str]:
-        return pipeline.transformer, "YourTransformer2DModel"
-
-    @staticmethod
-    def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
-        registry = HookRegistry.get_or_create(transformer)
-        registry.register_hook(hook._HOOK_NAME, hook)
-
-_MODEL_ADAPTERS["YourModel"] = YourModelAdapter
-```
-
-### 5.3 Run Coefficient Estimation
-
-Use the provided script to estimate coefficients:
-
-```python
-from vllm_omni.diffusion.cache.teacache.coefficient_estimator import (
-    TeaCacheCoefficientEstimator,
-)
-from datasets import load_dataset
-from tqdm import tqdm
-
-# Load model
-estimator = TeaCacheCoefficientEstimator(
-    model_path="/path/to/model",
-    model_type="Bagel",  # Your model type
-    device="cuda",
-)
-
-# Load prompts (paper suggests ~70 prompts)
-dataset = load_dataset("nateraw/parti-prompts", split="train")
-prompts = dataset["Prompt"][:70]
-
-# Collect data
-for prompt in tqdm(prompts):
-    estimator.collect_from_prompt(prompt, num_inference_steps=50)
-
-# Estimate coefficients
-coeffs = estimator.estimate(poly_order=4)
-print(f"Coefficients: {coeffs}")
-```
-
-### 5.4 Interpreting Coefficient Estimation Results
-
-The estimator outputs statistics and polynomial coefficients. Here's how to interpret them:
-
-**Example Output:**
-```
-Data statistics:
-Count: 48
-Input Diffs (x): min=1.1089e-02, max=5.2555e-02, mean=2.8435e-02
-Output Diffs (y): min=2.8242e-02, max=2.9792e-01, mean=7.0312e-02
-Coefficients: [1333131.29, -168644.23, 7950.51, -163.75, 1.26]
-```
-
-**What to Check:**
-- **Count**: Number of timestep pairs analyzed. Should be at least 30-50 for reliable estimation. Low count suggests insufficient prompts or inference steps.
-- **Input/Output Ranges**: Verify output differences correlate with input differences. If ranges seem unusual, check your prompt diversity.
-- **Coefficient Magnitude**: Extremely large values (>1e8) may indicate numerical instability - try collecting more diverse data.
-
-**Troubleshooting:**
-- If results seem unreliable, try:
-  - Increasing number of prompts (100+ recommended)
-  - Using more diverse prompts from multiple datasets
-  - Adjusting `num_inference_steps` (try 20, 50, 100)
-
-### 5.5 Add Coefficients to Config
-
-Add the estimated coefficients to `vllm_omni/diffusion/cache/teacache/config.py`:
-
-```python
-_MODEL_COEFFICIENTS = {
-    ...
-    "YourTransformer2DModel": [
-        1.04730573e+06,  # a4
-        -1.34150749e+05, # a3
-        6.51517806e+03,  # a2
-        -1.41209108e+02, # a1
-        1.17241808e+00,  # a0
-    ],
-}
-```
-## Step 6: Open a Pull Request
+## Step 5: Open a Pull Request
 
 When submitting a pull request to add support for a new model, please include the following information in the PR description:
 

docs/contributing/profiling.md

@@ -2,7 +2,7 @@
 
 > **Warning:** Profiling incurs significant overhead. Use only for development and debugging, never in production.
 
-vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** and **diffusion models**.
+vLLM-Omni uses the PyTorch Profiler to analyze performance across both **Multi-Stage LLMs** and **Diffusion Models**.
 
 ### 1. Set the Output Directory
 Before running any script, set this environment variable. The system detects this and automatically saves traces here.
@@ -11,7 +11,7 @@ Before running any script, set this environment variable. The system detects thi
 export VLLM_TORCH_PROFILER_DIR=./profiles
 ```
 
-### 2. Profiling Omni-Modality Models
+### 2. Start Profiling
 
 It is best to limit profiling to one iteration to keep trace files manageable.
 
@@ -77,12 +77,12 @@ omni_llm.close()
 
 **Examples**:
 
-1. **Qwen2.5-Omni**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
+1. **Qwen-omni 2.5**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
 
-2. **Qwen3-Omni**:   [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
+2. **Qwen-omni 3.0**:   [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
 
 
-### 3. Profiling diffusion models
+**For Diffusion Models as a single stage**
 
 Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
 
@@ -131,10 +131,27 @@ python image_to_video.py \
 
 2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**:   [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
 
-> **Note:**
-As of now, asynchronous (online) profiling is not fully supported in vLLM-Omni. While start_profile() and stop_profile() methods exist, they are only reliable in offline inference scripts (e.g., the provided end2end.py examples). Do not use them in server-mode or streaming scenarios—traces may be incomplete or fail to flush.
 
-### 4. Analyzing Omni Traces
+**Online Inference(Async)**
+
+For online serving using AsyncOmni, the methods are asynchronous. This allows you to toggle profiling dynamically without restarting the server.
+
+```python
+from vllm_omni import AsyncOmni
+
+# Inside an async function:
+async_omni = AsyncOmni.from_engine_args(engine_args)
+
+await async_omni.start_profile()
+
+async for output in async_omni.generate(prompt, sampling_params, request_id):
+    # Process outputs...
+    pass
+
+await async_omni.stop_profile()
+```
+
+### 3. Analyzing Omni Traces
 
 Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
 
@@ -146,4 +163,4 @@ Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
 - [Perfetto](https://ui.perfetto.dev/)(recommended)
 - ```chrome://tracing```(Chrome only)
 
-**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/stable/contributing/profiling/)
+**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/latest/dev/profiling.html)

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

@@ -17,7 +17,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 # --8<-- [start:pre-built-wheels]
 
 #### Installation of vLLM
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
@@ -93,7 +93,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.14.0 \
+    vllm/vllm-omni:v0.14.0rc1 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

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

@@ -80,7 +80,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.14.0 \
+  vllm/vllm-omni-rocm:v0.14.0rc1 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -99,7 +99,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.14.0
+  vllm/vllm-omni-rocm:v0.14.0rc1
 ```
 
 # --8<-- [end:pre-built-images]

docs/serving/image_edit_api.md

@@ -1,205 +0,0 @@
-# Image Edit API
-
-vLLM-Omni provides an OpenAI DALL-E compatible API for image edit using diffusion models.
-
-Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
-
-## Quick Start
-
-### Start the Server
-
-For example...
-
-```bash
-# Qwen-Image
-vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
-
-
-### Generate Images
-
-**Using curl:**
-
-```bash
-curl -s -D >(grep -i x-request-id >&2) \
-  -o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
-  -X POST "http://localhost:8000/v1/images/edits" \
-  -F "model=xxx" \
-  -F "image=@./xx.png" \
-  -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
-  -F "size=1024x1024" \
-  -F "output_format=png"
-```
-
-
-**Using OpenAI SDK:**
-
-```python
-import base64
-from openai import OpenAI
-from pathlib import Path
-client = OpenAI(
-    api_key="None",
-    base_url="http://localhost:8000/v1"
-)
-
-input_image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png"
-
-result = client.images.edit(
-    image=[],
-    model="Qwen-Image-Edit-2511",
-    prompt="Change the bears in the two input images into walking together.",
-    size='512x512',
-    stream=False,
-    output_format='jpeg',
-    # url格式
-    extra_body={
-        "url": [input_image_url1,input_image_url],
-        "num_inference_steps": 50,
-        "guidance_scale": 1,
-        "seed": 777,
-    }
-)
-
-image_base64 = result.data[0].b64_json
-image_bytes = base64.b64decode(image_base64)
-
-# Save the image to a file
-with open("edit_out_http.jpeg", "wb") as f:
-    f.write(image_bytes)
-```
-
-## API Reference
-
-### Endpoint
-
-```
-POST /v1/images/edits
-Content-Type: multipart/form-data
-```
-
-### Request Parameters
-
-#### OpenAI Standard Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `prompt` | string | **required** | A text description of the desired image |
-| `model` | string | server's model | Model to use (optional, should match server if specified) |
-| `image` | string or array | **required** | The image(s) to edit. |
-| `n` | integer | 1 | Number of images to generate (1-10) |
-| `size` | string | "auto" | Image dimensions in WxH format (e.g., "1024x1024", "512x512"), when set to auto, it decide size from first input image. |
-| `response_format` | string | "b64_json" | Response format (only "b64_json" supported) |
-| `user` | string | null | User identifier for tracking |
-| `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
-| `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
-| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
-
-#### vllm-omni Extension Parameters
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `url` | string or array | None | The image(s) to edit. |
-| `negative_prompt` | string | null | Text describing what to avoid in the image |
-| `num_inference_steps` | integer | model defaults | Number of diffusion steps |
-| `guidance_scale` | float | model defaults | Classifier-free guidance scale (typically 0.0-20.0) |
-| `true_cfg_scale` | float | model defaults | True CFG scale (model-specific parameter, may be ignored if not supported) |
-| `seed` | integer | null | Random seed for reproducibility |
-
-### Response Format
-
-```json
-{
-  "created": 1701234567,
-  "data": [
-    {
-      "b64_json": "<base64-encoded PNG>",
-      "url": null,
-      "revised_prompt": null
-    }
-  ],
-  "output_format": null,
-  "size": null,
-}
-```
-
-## Examples
-
-### Multiple Images input
-
-```bash
-curl -s -D >(grep -i x-request-id >&2) \
-  -o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
-  -X POST "http://localhost:8000/v1/images/edits" \
-  -F "model=xxx" \
-  -F "image=@xx.png" \
-  -F "image=@xx.png"
-  -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
-  -F "size=1024x1024" \
-  -F "output_format=png"
-```
-
-
-## Parameter Handling
-
-The API passes parameters directly to the diffusion pipeline without model-specific transformation:
-
-- **Default values**: When parameters are not specified, the underlying model uses its own defaults
-- **Pass-through design**: User-provided values are forwarded directly to the diffusion engine
-- **Minimal validation**: Only basic type checking and range validation at the API level
-
-### Parameter Compatibility
-
-The API passes parameters directly to the diffusion pipeline without model-specific validation.
-
-- Unsupported parameters may be silently ignored by the model
-- Incompatible values will result in errors from the underlying pipeline
-- Recommended values vary by model - consult model documentation
-
-**Best Practice:** Start with the model's recommended parameters, then adjust based on your needs.
-
-## Error Responses
-
-### 400 Bad Request
-
-Invalid parameters (e.g., model mismatch):
-
-```json
-{
-  "detail": "Invalid size format: '1024x'. Expected format: 'WIDTHxHEIGHT' (e.g., '1024x1024')."
-}
-```
-
-### 422 Unprocessable Entity
-
-Validation errors (missing required fields):
-
-```json
-{
-  "detail": "Field 'image' or 'url' is required"
-}
-```
-
-## Troubleshooting
-
-### Server Not Running
-
-```bash
-# Check if server is responding
-curl -X http://localhost:8000/v1/images/edit \
-  -F "prompt='test'"
-```
-
-### Out of Memory
-
-If you encounter OOM errors:
-1. Reduce image size: `"size": "512x512"`
-2. Reduce inference steps: `"num_inference_steps": 25`
-
-## Development
-
-Enable debug logging to see prompts and generation details:
-
-```bash
-vllm serve Qwen/Qwen-Image-Edit-2511 --omni \
-  --uvicorn-log-level debug
-```

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -1,32 +1,19 @@
 # CPU Offloading for Diffusion Model
 
 ## Overview
-
-vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models, allowing you to run larger models on GPUs with limited VRAM:
-
-1. **Model-level (Component) Offloading**: Swaps entire model components (DiT transformer, VAE, encoders) between GPU and CPU.
-2. **Layerwise (Blockwise) Offloading**: Keeps only a single or a few transformer blocks on GPU at a time, with compute - memory copy overlap.
-
-Both approaches use pinned memory for faster CPU-GPU transfers. For now, the two offloading strategies could not be used at the same time.
-
-
-## Model-level CPU Offloading
-
-### Implementation
-
 CPU offload lets the diffusion worker move large model components between GPU and CPU memory on demand. It keeps the DiT transformer resident on GPU only while it is actively running, and swaps it out when encoders modules need the device. This reduces peak VRAM usage so bigger checkpoints run on smaller GPUs, or multiple requests can share the same GPU.
 
-**Execution Flow**:
+## Execution Model
 1. Text encoders run on GPU while the DiT transformer is offloaded to CPU.
 2. Before denoising, weights are prefetched back to GPU, honoring pinned-memory copies for speed.
 3. After the diffusion step, the transformer returns to CPU and the process repeats as needed.
 
 Transfers use pinned host buffers, and the worker coordinates swaps via mutex-style hooks so components never compete for memory.
 
-### Configuration
+## Configuration
 You can enable CPU offload in two ways:
 
-1. **Python API**: set `enable_cpu_offload=True`.
+- **Python API**: set `enable_cpu_offload=True`.
 
 ```python
 from vllm_omni import Omni
@@ -36,66 +23,7 @@ if __name__ == "__main__":
     m = Omni(model="Qwen/Qwen-Image",enable_cpu_offload=True)
 ```
 
-2. **CLI**: pass `--enable-cpu-offload` to the diffusion service entrypoint.
+- **CLI**: pass `--enable-cpu-offload` to the diffusion service entrypoint.
 
-### Limitations
+## Known Limitations
 - Cold start latency increases for over one minute for some models(e.g., Qwen-Image)
-
-
-## Layerwise (Blockwise) Offloading
-
-### Implementation
-Layerwise offload operates at transformer block granularity, keeping a single transformer block, or a specified number of blocks, on GPU while others stay in CPU memory.
-
-Unlike full model-wise CPU offload which swaps entire components like DiT and encoders, layerwise offloading applies a sliding window way of loading and offloading weights between gpu and cpu: while block `i` computes, block `i+1` get prefetched asynchronously via pinned memory. In this way, only partial blocks(s) reside on GPU at any moment during inference, so that greatly decrease the memory occupancy.
-
-**Execution Flow**:
-
-1. During model initialization, all components are loaded to CPU first. Then components other than DiT model(s) in the pipeline, such as VAE and encoders, are moved to GPU. The weights of target transformer blocks are collected as contiguous tensors per layer on CPU with pinned memory; and non-block modules (embeddings, norms, etc) in the DiT model are moved to and stay on GPU.
-2. The first block(s) are transferred to GPU during initialization of `LayerwiseOffloader`, before the first denoising step of the very first request.
-3. As each block executes, the next block prefetches on a separate CUDA stream for compute - memory copy overlap. After execution, the current block is immediately freed from GPU memory.
-4. When the last block completes, the first block prefetches for the next denoising step.
-
-
-Example of hook executions of a DiT model with n layers, by default keep a single layer on GPU:
-| Layer (block) idx | forward pre-hook               | forward          | forward post-hook         |
-|-------------------|--------------------------------|------------------|---------------------------|
-| layer-0           | prefetch layer 1 (copy stream) | compute layer 0  | free layer-0 gpu weights  |
-| layer-1           | prefetch layer 2 (copy stream) | compute layer 1  | free layer-1 gpu weights  |
-| layer-2           | prefetch layer 3 (copy stream) | compute layer 2  | free layer-2 gpu weights  |
-| ...               | ...                            | ...              | ...                       |
-| layer-(n-1)       | **prefetch layer 0 (copy stream)** | compute layer (n-1) | free layer (n-1) gpu weights  |
-
-
-### Configuration
-
-1. **Python API**: set `enable_layerwise_offload=True` and optionally `layerwise_num_gpu_layers`.
-
-```python
-from vllm_omni import Omni
-
-if __name__ == "__main__":
-    m = Omni(
-        model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-        enable_layerwise_offload=True,
-        ...
-    )
-```
-
-2. **CLI**: pass `--enable-layerwise-offload` and `--layerwise-num-gpu-layers` to the diffusion service entrypoint.
-
-### Supported Models
-
-| Architecture | Models | Example HF Models | DiT Model Cls | Blocks Attr Name |
-|--------------|--------|-------------------|----------|----------|
-| `QwenImagePipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | "transformer_blocks" |
-| `Wan22Pipeline` | Wan2.2 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | "blocks" |
-
-NOTE: Models must define `_layerwise_offload_blocks_attr` class attribute so that the layerwise offloader finds the target transformer blocks.
-
-### Limitations
-- Cold start latency increases because of
-    1) components are loaded to CPU first at the very first during initialization,  
-    2) weight consolidation and pinning
-- Performance depends on CPU <-> GPU interconnection (e.g., PCIe bandwidth).
-- Support single GPU only for now

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -18,19 +18,20 @@ The following table shows which models are currently supported by parallelism me
 
 ### ImageGen
 
-| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
-|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|
-| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |
-| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |
-| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |
-| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |
-| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ❌        |
-| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |
-| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |
+| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
+|-------|------------------|------------|---------|--------------|--------------------------|
+| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ❌ |
+| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ❌ |
+| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ |
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ |
+| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) |
+| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
+| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ |
+| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ❌ | ✅ |
+
 
 !!! note "TP Limitations for Diffusion Models"
     We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
@@ -372,9 +373,9 @@ def forward(self, hidden_states, ...):
 
 ### CFG-Parallel
 
-#### Offline Inference
+##### Offline Inference
 
-CFG-Parallel is enabled through `DiffusionParallelConfig(cfg_parallel_size=2)`, which runs one rank for the positive branch and one rank for the negative branch.
+CFG-Parallel is enabled through `DiffusionParallelConfig(cfg_parallel_size=...)`. The recommended configuration is `cfg_parallel_size=2` (one rank for the positive branch and one rank for the negative branch).
 
 An example of offline inference using CFG-Parallel (image-to-image) is shown below:
 
@@ -382,21 +383,19 @@ An example of offline inference using CFG-Parallel (image-to-image) is shown bel
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 
-image_path = "path_to_image.png"
 omni = Omni(
     model="Qwen/Qwen-Image-Edit",
     parallel_config=DiffusionParallelConfig(cfg_parallel_size=2),
 )
-input_image = Image.open(image_path).convert("RGB")
 
 outputs = omni.generate(
     {
         "prompt": "turn this cat to a dog",
         "negative_prompt": "low quality, blurry",
-        "multi_modal_data": {"image": input_image},
     },
     OmniDiffusionSamplingParams(
         true_cfg_scale=4.0,
+        pil_image=input_image,
         num_inference_steps=50,
     ),
 )
@@ -404,28 +403,7 @@ outputs = omni.generate(
 
 Notes:
 
-- CFG-Parallel is only effective when a `negative_prompt` is provided AND a guidance scale (or `cfg_scale`) is greater than 1.
-
-See `examples/offline_inference/image_to_image/image_edit.py` for a complete working example.
-```bash
-cd examples/offline_inference/image_to_image/
-python image_edit.py \
-  --model "Qwen/Qwen-Image-Edit" \
-  --image "qwen_image_output.png" \
-  --prompt "turn this cat to a dog" \
-  --negative_prompt "low quality, blurry" \
-  --cfg_scale 4.0 \
-  --output "edited_image.png" \
-  --cfg_parallel_size 2
-```
-
-#### Online Serving
-
-You can enable CFG-Parallel in online serving for diffusion models via `--cfg-parallel-size`:
-
-```bash
-vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
-```
+- CFG-Parallel is only effective when **true CFG** is enabled (i.e., `true_cfg_scale > 1` and a `negative_prompt` is provided).
 
 #### How to parallelize a pipeline
 
@@ -438,130 +416,58 @@ In `QwenImagePipeline`, each diffusion step runs two denoiser forward passes seq
 
 CFG-Parallel assigns these two branches to different ranks in the **CFG group** and synchronizes the results.
 
-vLLM-omni provides `CFGParallelMixin` base class that encapsulates the CFG parallel logic. By inheriting from this mixin and calling its methods, pipelines can easily implement CFG parallel without writing repetitive code.
-
-**Key Methods in CFGParallelMixin:**
-- `predict_noise_maybe_with_cfg()`: Automatically handles CFG parallel noise prediction
-- `scheduler_step_maybe_with_cfg()`: Scheduler step with automatic CFG rank synchronization
-
-**Example Implementation:**
+Below is an example of CFG-Parallel implementation:
 
 ```python
-class QwenImageCFGParallelMixin(CFGParallelMixin):
-    """
-    Base Mixin class for Qwen Image pipelines providing shared CFG methods.
-    """
-
-    def diffuse(
+def diffuse(
         self,
-        prompt_embeds: torch.Tensor,
-        prompt_embeds_mask: torch.Tensor,
-        negative_prompt_embeds: torch.Tensor,
-        negative_prompt_embeds_mask: torch.Tensor,
-        latents: torch.Tensor,
-        img_shapes: torch.Tensor,
-        txt_seq_lens: torch.Tensor,
-        negative_txt_seq_lens: torch.Tensor,
-        timesteps: torch.Tensor,
-        do_true_cfg: bool,
-        guidance: torch.Tensor,
-        true_cfg_scale: float,
-        image_latents: torch.Tensor | None = None,
-        cfg_normalize: bool = True,
-        additional_transformer_kwargs: dict[str, Any] | None = None,
-    ) -> torch.Tensor:
-        self.transformer.do_true_cfg = do_true_cfg
+        ...
+        ):
+    # Enable CFG-parallel: rank0 computes positive, rank1 computes negative.
+    cfg_parallel_ready = do_true_cfg and get_classifier_free_guidance_world_size() > 1
 
-        for i, t in enumerate(timesteps):
-            timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
+    self.transformer.do_true_cfg = do_true_cfg
 
-            # Prepare kwargs for positive (conditional) prediction
-            positive_kwargs = {
-                "hidden_states": latents,
-                "timestep": timestep / 1000,
-                "guidance": guidance,
-                "encoder_hidden_states_mask": prompt_embeds_mask,
-                "encoder_hidden_states": prompt_embeds,
-                "img_shapes": img_shapes,
-                "txt_seq_lens": txt_seq_lens,
-            }
+    if cfg_parallel_ready:
+        cfg_group = get_cfg_group()
+        cfg_rank = get_classifier_free_guidance_rank()
 
-            # Prepare kwargs for negative (unconditional) prediction
-            if do_true_cfg:
-                negative_kwargs = {
-                    "hidden_states": latents,
-                    "timestep": timestep / 1000,
-                    "guidance": guidance,
-                    "encoder_hidden_states_mask": negative_prompt_embeds_mask,
-                    "encoder_hidden_states": negative_prompt_embeds,
-                    "img_shapes": img_shapes,
-                    "txt_seq_lens": negative_txt_seq_lens,
-                }
-            else:
-                negative_kwargs = None
+        if cfg_rank == 0:
+            local_pred = self.transformer(
+                hidden_states=latents,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                encoder_hidden_states_mask=prompt_embeds_mask,
+                encoder_hidden_states=prompt_embeds,
+                img_shapes=img_shapes,
+                txt_seq_lens=txt_seq_lens,
+                attention_kwargs=self.attention_kwargs,
+                return_dict=False,
+            )[0]
+        else:
+            local_pred = self.transformer(
+                hidden_states=latents,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                encoder_hidden_states_mask=negative_prompt_embeds_mask,
+                encoder_hidden_states=negative_prompt_embeds,
+                img_shapes=img_shapes,
+                txt_seq_lens=negative_txt_seq_lens,
+                attention_kwargs=self.attention_kwargs,
+                return_dict=False,
+            )[0]
 
-            # Predict noise with automatic CFG parallel handling
-            # - In CFG parallel mode: rank0 computes positive, rank1 computes negative
-            # - Automatically gathers results and combines them on rank0
-            noise_pred = self.predict_noise_maybe_with_cfg(
-                do_true_cfg=do_true_cfg,
-                true_cfg_scale=true_cfg_scale,
-                positive_kwargs=positive_kwargs,
-                negative_kwargs=negative_kwargs,
-                cfg_normalize=cfg_normalize,
-            )
-
-            # Step scheduler with automatic CFG synchronization
-            # - Only rank0 computes the scheduler step
-            # - Automatically broadcasts updated latents to all ranks
-            latents = self.scheduler_step_maybe_with_cfg(
-                noise_pred, t, latents, do_true_cfg
-            )
-
-        return latents
-```
-
-**How it works:**
-1. Prepare separate `positive_kwargs` and `negative_kwargs` for conditional and unconditional predictions
-2. Call `predict_noise_maybe_with_cfg()` which:
-   - Detects if CFG parallel is enabled (`get_classifier_free_guidance_world_size() > 1`)
-   - Distributes computation: rank0 processes positive, rank1 processes negative
-   - Gathers predictions and combines them using `combine_cfg_noise()` on rank0
-   - Returns combined noise prediction (only valid on rank0)
-3. Call `scheduler_step_maybe_with_cfg()` which:
-   - Only rank0 computes the scheduler step
-   - Broadcasts the updated latents to all ranks for synchronization
-
-**How to customize**
-
-Some pipelines may need to customize the following functions in `CFGParallelMixin`:
-1. You may need to edit `predict_noise` function for custom behaviors.
-```python
-def predict_noise(self, *args, **kwargs):
-    """
-    Forward pass through transformer to predict noise.
-
-    Subclasses should override this if they need custom behavior,
-    but the default implementation calls self.transformer.
-    """
-    return self.transformer(*args, **kwargs)[0]
-
-```
-2. The default normalization function after combining the noise predictions from both branches is as follows. You may need to customize it.
-```python
-def cfg_normalize_function(self, noise_pred, comb_pred):
-    """
-    Normalize the combined noise prediction.
-
-    Args:
-        noise_pred: positive noise prediction
-        comb_pred: combined noise prediction after CFG
-
-    Returns:
-        Normalized noise prediction tensor
-    """
-    cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
-    noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
-    noise_pred = comb_pred * (cond_norm / noise_norm)
-    return noise_pred
+        gathered = cfg_group.all_gather(local_pred, separate_tensors=True)
+        if cfg_rank == 0:
+            noise_pred = gathered[0]
+            neg_noise_pred = gathered[1]
+            comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+            cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+            noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+            noise_pred = comb_pred * (cond_norm / noise_norm)
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+        cfg_group.broadcast(latents, src=0)
+    else:
+        # fallback: run positive then negative sequentially on one rank
+        ...
 ```

docs/user_guide/diffusion_acceleration.md

@@ -39,23 +39,23 @@ The following table shows which models are currently supported by each accelerat
 
 | Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
 |-------|------------------|:----------:|:-----------:|:-----------:|:----------------:|:----------------:|
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ |
+| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ |
 | **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ |
+| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ |
 
 ### VideoGen
 
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention |CFG-Parallel |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:----------------:|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention |
+|-------|------------------|:--------:|:---------:|:----------:|:--------------:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ❌ | ❌ |
 
 
 ## Performance Benchmarks

docs/user_guide/examples/offline_inference/bagel.md

@@ -1,179 +0,0 @@
-# BAGEL-7B-MoT
-
-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.
-
-## Run examples
-
-**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
-
-Get into the bagel folder
-
-```bash
-cd examples/offline_inference/bagel
-```
-
-### Modality Control
-
-BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
-
-#### Text to Image (text2img)
-
-- **Pipeline**: Text → Thinker  → DiT → VAE Decode → Image
-- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
-- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
-
-Generate images from text prompts:
-
-```bash
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality text2img \
-                  --prompts "A cute cat"
-```
-
-#### Image to Image (img2img)
-
-- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
-- **Stages Used**: Stage 1 (DiT) only
-- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
-
-Transform images based on text prompts:
-
-```bash
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality img2img \
-                  --image-path /path/to/image.jpg \
-                  --prompts "Let the woman wear a blue dress"
-```
-
-#### Image to Text (img2text)
-
-- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
-- **Stages Used**: Stage 0 (Thinker) only
-- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
-
-Generate text descriptions from images:
-
-```bash
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality img2text \
-                  --image-path /path/to/image.jpg \
-                  --prompts "Describe this image in detail"
-```
-
-#### Text to Text (text2text)
-
-- **Pipeline**: Text → Thinker → Text Output
-- **Stages Used**: Stage 0 (Thinker) only
-- **Special**: No visual components involved, operates as pure language model
-
-Pure text generation:
-
-```bash
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality text2text \
-                  --prompts "What is the capital of France?"
-
-# You can load prompts from a text file (one prompt per line):  
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality text2text \
-                  --txt-prompts /path/to/prompts.txt
-```
-
-### Inference Steps
-
-Control the number of inference steps for image generation:
-
-```bash
-# You can adjust steps to 100 to improve image quality
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality text2img \
-                  --steps 50 \
-                  --prompts "A cute cat"
-```
-
-### Key arguments
-
-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)
-
-#### 📌 Command Line Arguments (end2end.py)
-
-| Argument               | Type   | Default                       | Description                                                  |
-| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
-| `--model`              | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name                                           |
-| `--modality`           | choice | `text2img`                    | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
-| `--prompts`            | list   | `None`                        | Input text prompts directly                                  |
-| `--txt-prompts`        | string | `None`                        | Path to txt file with one prompt per line                    |
-| `--image-path`         | string | `None`                        | Input image path (for `img2img`/`img2text`)                  |
-| `--steps`              | int    | `50`                          | Number of inference steps                                    |
-| `--stage-configs-path` | string | `None`                        | Custom stage config file path                                |
-| `--worker-backend`     | choice | `process`                     | Worker backend: `process` or `ray`                           |
-| `--ray-address`        | string | `None`                        | Ray cluster address                                          |
-| `--enable-stats`       | flag   | `False`                       | Enable statistics logging                                    |
-| `--init-sleep-seconds` | int    | `20`                          | Initialization sleep time                                    |
-| `--batch-timeout`      | int    | `5`                           | Batch timeout                                                |
-| `--init-timeout`       | int    | `300`                         | Initialization timeout                                       |
-
-------
-
-#### ⚙️ Stage Configuration Parameters (bagel.yaml)
-
- **Stage 0 - Thinker (LLM Stage)**
-
-| Parameter                        | Value                           | Description              |
-| :------------------------------- | :------------------------------ | :----------------------- |
-| `stage_type`                     | `llm`                           | Stage type               |
-| `devices`                        | `"0"`                           | GPU device ID            |
-| `max_batch_size`                 | `1`                             | Maximum batch size       |
-| `model_stage`                    | `thinker`                       | Model stage identifier   |
-| `model_arch`                     | `BagelForConditionalGeneration` | Model architecture       |
-| `gpu_memory_utilization`         | `0.4`                           | GPU memory utilization   |
-| `tensor_parallel_size`           | `1`                             | Tensor parallel size     |
-| `max_num_batched_tokens`         | `32768`                         | Maximum batched tokens   |
-| `omni_kv_config.need_send_cache` | `true`                          | Whether to send KV cache |
-
-------
-
-**Stage 1 - DiT (Diffusion Stage)**
-
-| Parameter                        | Value       | Description                 |
-| :------------------------------- | :---------- | :-------------------------- |
-| `stage_type`                     | `diffusion` | Stage type                  |
-| `devices`                        | `"0"`       | GPU device ID               |
-| `max_batch_size`                 | `1`         | Maximum batch size          |
-| `model_stage`                    | `dit`       | Model stage identifier      |
-| `gpu_memory_utilization`         | `0.4`       | GPU memory utilization      |
-| `omni_kv_config.need_recv_cache` | `true`      | Whether to receive KV cache |
-| `engine_input_source`            | `[0]`       | Input source from Stage 0   |
-
-------
-
-#### 🔗 Runtime Configuration
-
-| Parameter             | Value   | Description                      |
-| :-------------------- | :------ | :------------------------------- |
-| `window_size`         | `-1`    | Window size (-1 means unlimited) |
-| `max_inflight`        | `1`     | Maximum inflight requests        |
-| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB)   |
-
-## FAQ
-
-- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
-
-```bash
-sudo apt update
-sudo apt install ffmpeg
-```
-
-- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
-
-| Stage               | VRAM                         |
-| :------------------ | :--------------------------- |
-| Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
-| Stage-1 (DiT)       | **26.50 GiB**                |
-| Total               | **~42 GiB + KV Cache**       |

docs/user_guide/examples/offline_inference/image_to_image.md

@@ -47,7 +47,6 @@ 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.

docs/user_guide/examples/offline_inference/image_to_video.md

@@ -52,7 +52,6 @@ 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).

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -95,7 +95,6 @@ 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).

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -28,7 +28,6 @@ Key arguments:
 - `--num_frames`: Number of frames (Wan default is 81).
 - `--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.
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.

docs/user_guide/examples/online_serving/bagel.md

@@ -1,232 +0,0 @@
-# BAGEL-7B-MoT
-
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/bagel>.
-
-## 🛠️ Installation
-
-Please refer to [README.md](../../../README.md)
-
-## Run examples (BAGEL-7B-MoT)
-
-**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
-
-### Launch the Server
-
-```bash
-# Use default configuration
-vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
-```
-
-Or use the convenience script:
-
-```bash
-cd /workspace/vllm-omni/examples/online_serving/bagel
-bash run_server.sh
-```
-
-If you have a custom stage configs file, launch the server with the command below:
-
-```bash
-vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
-```
-
-### Send Multi-modal Request
-
-Get into the bagel folder:
-
-```bash
-cd examples/online_serving/bagel
-```
-
-Send request via Python
-
-```bash
-python openai_chat_client.py --prompt "A cute cat" --modality text2img
-```
-
-The Python client supports the following command-line arguments:
-
-- `--prompt` (or `-p`): Text prompt for generation (default: `A cute cat`)
-- `--output` (or `-o`): Output file path for image results (default: `bagel_output.png`)
-- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
-- `--image-url` (or `-i`): Input image URL or local file path (for img2img/img2text modes)
-- `--modality` (or `-m`): Task modality (default: `text2img`). Options: `text2img`, `img2img`, `img2text`, `text2text`
-- `--height`: Image height in pixels (default: 512)
-- `--width`: Image width in pixels (default: 512)
-- `--steps`: Number of inference steps (default: 25)
-- `--seed`: Random seed (default: 42)
-- `--negative`: Negative prompt for image generation
-
-Example with custom parameters:
-
-```bash
-python openai_chat_client.py \
-    --prompt "A futuristic city" \
-    --modality text2img \
-    --height 768 \
-    --width 768 \
-    --steps 50 \
-    --seed 42 \
-    --negative "blurry, low quality"
-```
-
-## Modality Control
-
-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)
-
-| Modality    | Input        | Output | Description                            |
-| ----------- | ------------ | ------ | -------------------------------------- |
-| `text2img`  | Text         | Image  | Generate images from text prompts      |
-| `img2img`   | Image + Text | Image  | Transform images using text guidance   |
-| `img2text`  | Image + Text | Text   | Generate text descriptions from images |
-| `text2text` | Text         | Text   | Pure text generation                   |
-
-### Text to Image (text2img)
-
-Generate images from text prompts:
-
-**Using Python client**
-
-```bash
-python openai_chat_client.py \
-    --prompt "A beautiful sunset over mountains" \
-    --modality text2img \
-    --output sunset.png \
-    --steps 50
-```
-
-**Using curl**
-
-```bash
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}]}],
-    "modalities": ["image"],
-    "height": 512,
-    "width": 512,
-    "num_inference_steps": 50,
-    "seed": 42
-  }'
-```
-
-
-### Image to Image (img2img)
-
-Transform images based on text prompts:
-
-**Using Python client**
-
-```bash
-python openai_chat_client.py \
-    --prompt "Make the cat stand up" \
-    --modality img2img \
-    --image-url /path/to/input.jpg \
-    --output transformed.png
-```
-
-**Using curl**
-
-```bash
-IMAGE_BASE64=$(base64 -w 0 cat.jpg)
-
-cat <<EOF > payload.json
-{
-    "messages": [{
-      "role": "user",
-      "content": [
-        {"type": "text", "text": "<|im_start|>Make the cat stand up<|im_end|>"},
-        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
-      ]
-    }],
-    "modalities": ["image"],
-    "height": 512,
-    "width": 512,
-    "num_inference_steps": 50,
-    "seed": 42
-}
-EOF
-
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d @payload.json
-
-```
-
-### Image to Text (img2text)
-
-Generate text descriptions from images:
-
-**Using Python client**
-
-```bash
-python openai_chat_client.py \
-    --prompt "Describe this image in detail" \
-    --modality img2text \
-    --image-url /path/to/image.jpg
-```
-
-**Using curl**
-
-```bash
-IMAGE_BASE64=$(base64 -w 0 cat.jpg)
-
-cat <<EOF > payload.json
-{
-  "messages": [{
-    "role": "user",
-    "content": [
-      {"type": "text", "text": "<|im_start|>user\n<|image_pad|>\nDescribe this image in detail<|im_end|>\n<|im_start|>assistant\n"},
-      {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
-    ]
-  }],
-  "modalities": ["text"]
-}
-EOF
-
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d @payload.json
-```
-
-### Text to Text (text2text)
-
-Pure text generation:
-
-**Using Python client**
-
-```bash
-python openai_chat_client.py \
-    --prompt "What is the capital of France?" \
-    --modality text2text
-```
-
-**Using curl**
-
-```bash
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}]
-    "modalities": ["text"]
-  }'
-```
-
-## FAQ
-
-- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
-
-```bash
-sudo apt update
-sudo apt install ffmpeg
-```
-
-- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
-
-| Stage               | VRAM                         |
-| :------------------ | :--------------------------- |
-| Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
-| Stage-1 (DiT)       | **26.50 GiB**                |
-| Total               | **~42 GiB + KV Cache**       |

examples/offline_inference/bagel/README.md

@@ -1,177 +0,0 @@
-# BAGEL-7B-MoT
-
-## 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.
-
-## Run examples
-
-**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
-
-Get into the bagel folder
-
-```bash
-cd examples/offline_inference/bagel
-```
-
-### Modality Control
-
-BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
-
-#### Text to Image (text2img)
-
-- **Pipeline**: Text → Thinker  → DiT → VAE Decode → Image
-- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
-- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
-
-Generate images from text prompts:
-
-```bash
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality text2img \
-                  --prompts "A cute cat"
-```
-
-#### Image to Image (img2img)
-
-- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
-- **Stages Used**: Stage 1 (DiT) only
-- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
-
-Transform images based on text prompts:
-
-```bash
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality img2img \
-                  --image-path /path/to/image.jpg \
-                  --prompts "Let the woman wear a blue dress"
-```
-
-#### Image to Text (img2text)
-
-- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
-- **Stages Used**: Stage 0 (Thinker) only
-- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
-
-Generate text descriptions from images:
-
-```bash
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality img2text \
-                  --image-path /path/to/image.jpg \
-                  --prompts "Describe this image in detail"
-```
-
-#### Text to Text (text2text)
-
-- **Pipeline**: Text → Thinker → Text Output
-- **Stages Used**: Stage 0 (Thinker) only
-- **Special**: No visual components involved, operates as pure language model
-
-Pure text generation:
-
-```bash
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality text2text \
-                  --prompts "What is the capital of France?"
-
-# You can load prompts from a text file (one prompt per line):  
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality text2text \
-                  --txt-prompts /path/to/prompts.txt
-```
-
-### Inference Steps
-
-Control the number of inference steps for image generation:
-
-```bash
-# You can adjust steps to 100 to improve image quality
-python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
-                  --modality text2img \
-                  --steps 50 \
-                  --prompts "A cute cat"
-```
-
-### Key arguments
-
-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)
-
-#### 📌 Command Line Arguments (end2end.py)
-
-| Argument               | Type   | Default                       | Description                                                  |
-| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
-| `--model`              | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name                                           |
-| `--modality`           | choice | `text2img`                    | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
-| `--prompts`            | list   | `None`                        | Input text prompts directly                                  |
-| `--txt-prompts`        | string | `None`                        | Path to txt file with one prompt per line                    |
-| `--image-path`         | string | `None`                        | Input image path (for `img2img`/`img2text`)                  |
-| `--steps`              | int    | `50`                          | Number of inference steps                                    |
-| `--stage-configs-path` | string | `None`                        | Custom stage config file path                                |
-| `--worker-backend`     | choice | `process`                     | Worker backend: `process` or `ray`                           |
-| `--ray-address`        | string | `None`                        | Ray cluster address                                          |
-| `--enable-stats`       | flag   | `False`                       | Enable statistics logging                                    |
-| `--init-sleep-seconds` | int    | `20`                          | Initialization sleep time                                    |
-| `--batch-timeout`      | int    | `5`                           | Batch timeout                                                |
-| `--init-timeout`       | int    | `300`                         | Initialization timeout                                       |
-
-------
-
-#### ⚙️ Stage Configuration Parameters (bagel.yaml)
-
- **Stage 0 - Thinker (LLM Stage)**
-
-| Parameter                        | Value                           | Description              |
-| :------------------------------- | :------------------------------ | :----------------------- |
-| `stage_type`                     | `llm`                           | Stage type               |
-| `devices`                        | `"0"`                           | GPU device ID            |
-| `max_batch_size`                 | `1`                             | Maximum batch size       |
-| `model_stage`                    | `thinker`                       | Model stage identifier   |
-| `model_arch`                     | `BagelForConditionalGeneration` | Model architecture       |
-| `gpu_memory_utilization`         | `0.4`                           | GPU memory utilization   |
-| `tensor_parallel_size`           | `1`                             | Tensor parallel size     |
-| `max_num_batched_tokens`         | `32768`                         | Maximum batched tokens   |
-| `omni_kv_config.need_send_cache` | `true`                          | Whether to send KV cache |
-
-------
-
-**Stage 1 - DiT (Diffusion Stage)**
-
-| Parameter                        | Value       | Description                 |
-| :------------------------------- | :---------- | :-------------------------- |
-| `stage_type`                     | `diffusion` | Stage type                  |
-| `devices`                        | `"0"`       | GPU device ID               |
-| `max_batch_size`                 | `1`         | Maximum batch size          |
-| `model_stage`                    | `dit`       | Model stage identifier      |
-| `gpu_memory_utilization`         | `0.4`       | GPU memory utilization      |
-| `omni_kv_config.need_recv_cache` | `true`      | Whether to receive KV cache |
-| `engine_input_source`            | `[0]`       | Input source from Stage 0   |
-
-------
-
-#### 🔗 Runtime Configuration
-
-| Parameter             | Value   | Description                      |
-| :-------------------- | :------ | :------------------------------- |
-| `window_size`         | `-1`    | Window size (-1 means unlimited) |
-| `max_inflight`        | `1`     | Maximum inflight requests        |
-| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB)   |
-
-## FAQ
-
-- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
-
-```bash
-sudo apt update
-sudo apt install ffmpeg
-```
-
-- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
-
-| Stage               | VRAM                         |
-| :------------------ | :--------------------------- |
-| Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
-| Stage-1 (DiT)       | **26.50 GiB**                |
-| Total               | **~42 GiB + KV Cache**       |

examples/offline_inference/bagel/run_single_prompt.sh

@@ -0,0 +1,6 @@
+prompt="<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"
+
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                                 --prompt_type text \
+                                 --init-sleep-seconds 0 \
+                                 --prompts ${prompt}

examples/offline_inference/image_to_image/image_edit.py

@@ -77,11 +77,11 @@ from pathlib import Path
 import torch
 from PIL import Image
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
 
 def parse_args() -> argparse.Namespace:
@@ -280,32 +280,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
-    parser.add_argument(
-        "--vae_use_slicing",
-        action="store_true",
-        help="Enable VAE slicing for memory optimization.",
-    )
-    parser.add_argument(
-        "--vae_use_tiling",
-        action="store_true",
-        help="Enable VAE tiling for memory optimization.",
-    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
         help="Enable CPU offloading for diffusion models.",
     )
-    parser.add_argument(
-        "--enable-layerwise-offload",
-        action="store_true",
-        help="Enable layerwise (blockwise) offloading on DiT modules.",
-    )
-    parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
-    )
     return parser.parse_args()
 
 
@@ -327,8 +306,12 @@ def main():
     else:
         input_image = input_images
 
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    device = detect_device_type()
+    generator = torch.Generator(device=device).manual_seed(args.seed)
 
+    # Enable VAE memory optimizations on NPU
+    vae_use_slicing = is_npu()
+    vae_use_tiling = is_npu()
     parallel_config = DiffusionParallelConfig(
         ulysses_degree=args.ulysses_degree,
         ring_degree=args.ring_degree,
@@ -361,10 +344,8 @@ def main():
     # Initialize Omni with appropriate pipeline
     omni = Omni(
         model=args.model,
-        enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
-        vae_use_slicing=args.vae_use_slicing,
-        vae_use_tiling=args.vae_use_tiling,
+        vae_use_slicing=vae_use_slicing,
+        vae_use_tiling=vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         parallel_config=parallel_config,
@@ -441,6 +422,7 @@ def main():
 
     if not outputs:
         raise ValueError("No output generated from omni.generate()")
+    logger.info("Outputs: %s", outputs)
 
     # Extract images from OmniRequestOutput
     # omni.generate() returns list[OmniRequestOutput], extract images from request_output[0].images

examples/offline_inference/image_to_image/image_to_image.md

@@ -47,9 +47,4 @@ Key arguments:
 - `--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`](../../../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.

examples/offline_inference/image_to_video/README.md

@@ -54,9 +54,4 @@ Key arguments:
 - `--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`](../../../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.

examples/offline_inference/image_to_video/image_to_video.py

@@ -26,11 +26,10 @@ import numpy as np
 import PIL.Image
 import torch
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
 
 def parse_args() -> argparse.Namespace:
@@ -60,44 +59,11 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
     parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
-    parser.add_argument(
-        "--vae_use_slicing",
-        action="store_true",
-        help="Enable VAE slicing for memory optimization.",
-    )
-    parser.add_argument(
-        "--vae_use_tiling",
-        action="store_true",
-        help="Enable VAE tiling for memory optimization.",
-    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
         help="Enable CPU offloading for diffusion models.",
     )
-    parser.add_argument(
-        "--enable-layerwise-offload",
-        action="store_true",
-        help="Enable layerwise (blockwise) offloading on DiT modules.",
-    )
-    parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
-    )
-    parser.add_argument(
-        "--cfg_parallel_size",
-        type=int,
-        default=1,
-        choices=[1, 2],
-        help="Number of GPUs used for classifier free guidance parallel size.",
-    )
-    parser.add_argument(
-        "--enforce_eager",
-        action="store_true",
-        help="Disable torch.compile and force eager execution.",
-    )
     return parser.parse_args()
 
 
@@ -114,7 +80,8 @@ def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> t
 
 def main():
     args = parse_args()
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    device = detect_device_type()
+    generator = torch.Generator(device=device).manual_seed(args.seed)
 
     # Load input image
     image = PIL.Image.open(args.image).convert("RGB")
@@ -131,38 +98,26 @@ def main():
     # Resize image to target dimensions
     image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
 
+    # Enable VAE memory optimizations on NPU
+    vae_use_slicing = is_npu()
+    vae_use_tiling = is_npu()
+
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
-    parallel_config = DiffusionParallelConfig(
-        cfg_parallel_size=args.cfg_parallel_size,
-    )
+
     omni = Omni(
         model=args.model,
-        enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
-        vae_use_slicing=args.vae_use_slicing,
-        vae_use_tiling=args.vae_use_tiling,
+        vae_use_slicing=vae_use_slicing,
+        vae_use_tiling=vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
         enable_cpu_offload=args.enable_cpu_offload,
-        parallel_config=parallel_config,
-        enforce_eager=args.enforce_eager,
     )
 
     if profiler_enabled:
         print("[Profiler] Starting profiling...")
         omni.start_profile()
 
-    # Print generation configuration
-    print(f"\n{'=' * 60}")
-    print("Generation Configuration:")
-    print(f"  Model: {args.model}")
-    print(f"  Inference steps: {args.num_inference_steps}")
-    print(f"  Frames: {args.num_frames}")
-    print(f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}")
-    print(f"  Video size: {args.width}x{args.height}")
-    print(f"{'=' * 60}\n")
-
     # omni.generate() returns Generator[OmniRequestOutput, None, None]
     frames = omni.generate(
         {

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -410,7 +410,7 @@ def main(args):
         elif stage_outputs.final_output_type == "audio":
             for output in stage_outputs.request_output:
                 request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output["audio"]
+                audio_tensor = output.multimodal_output["audio"]
                 output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
                 sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
                 print(f"Request ID: {request_id}, Saved audio to {output_wav}")

examples/offline_inference/qwen3_omni/end2end.py

@@ -418,7 +418,7 @@ def main(args):
         elif stage_outputs.final_output_type == "audio":
             for output in stage_outputs.request_output:
                 request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output["audio"]
+                audio_tensor = output.multimodal_output["audio"]
                 output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
 
                 # Convert to numpy array and ensure correct format

examples/offline_inference/qwen3_tts/end2end.py

@@ -244,9 +244,9 @@ def main(args):
     for stage_outputs in omni_generator:
         for output in stage_outputs.request_output:
             request_id = output.request_id
-            audio_tensor = output.outputs[0].multimodal_output["audio"]
+            audio_tensor = output.multimodal_output["audio"]
             output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-            audio_samplerate = output.outputs[0].multimodal_output["sr"].item()
+            audio_samplerate = output.multimodal_output["sr"].item()
             # Convert to numpy array and ensure correct format
             audio_numpy = audio_tensor.float().detach().cpu().numpy()
 

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -22,7 +22,7 @@ import torch
 
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.platform_utils import detect_device_type
 
 
 def parse_args() -> argparse.Namespace:
@@ -118,7 +118,8 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410
 
 def main():
     args = parse_args()
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    device = detect_device_type()
+    generator = torch.Generator(device=device).manual_seed(args.seed)
 
     print(f"\n{'=' * 60}")
     print("Stable Audio Open - Text-to-Audio Generation")

examples/offline_inference/text_to_image/README.md

@@ -96,13 +96,8 @@ Key arguments:
 - `--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`](../../../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.
 
 ## Web UI Demo

examples/offline_inference/text_to_image/gradio_demo.py

@@ -7,7 +7,7 @@ import torch
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
 ASPECT_RATIOS: dict[str, tuple[int, int]] = {
     "1:1": (1328, 1328),
@@ -62,8 +62,8 @@ def parse_args() -> argparse.Namespace:
 @lru_cache(maxsize=1)
 def get_omni(model_name: str) -> Omni:
     # Enable VAE memory optimizations on NPU
-    vae_use_slicing = current_omni_platform.is_npu()
-    vae_use_tiling = current_omni_platform.is_npu()
+    vae_use_slicing = is_npu()
+    vae_use_tiling = is_npu()
     return Omni(
         model=model_name,
         vae_use_slicing=vae_use_slicing,
@@ -72,6 +72,7 @@ def get_omni(model_name: str) -> Omni:
 
 
 def build_demo(args: argparse.Namespace) -> gr.Blocks:
+    device = detect_device_type()
     omni = get_omni(args.model)
 
     def run_inference(
@@ -98,7 +99,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
             raise gr.Error("Inference steps must be a positive integer.")
         if num_images not in {1, 2, 3, 4}:
             raise gr.Error("Number of images must be 1, 2, 3, or 4.")
-        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
+        generator = torch.Generator(device=device).manual_seed(seed)
         outputs = omni.generate(
             prompt.strip(),
             OmniDiffusionSamplingParams(

examples/offline_inference/text_to_image/text_to_image.py

@@ -12,7 +12,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
 
 def parse_args() -> argparse.Namespace:
@@ -73,11 +73,6 @@ def parse_args() -> argparse.Namespace:
             "Default: None (no cache acceleration)."
         ),
     )
-    parser.add_argument(
-        "--enable-cache-dit-summary",
-        action="store_true",
-        help="Enable cache-dit summary logging after diffusion forward passes.",
-    )
     parser.add_argument(
         "--ulysses_degree",
         type=int,
@@ -107,39 +102,23 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable CPU offloading for diffusion models.",
     )
-    parser.add_argument(
-        "--enable-layerwise-offload",
-        action="store_true",
-        help="Enable layerwise (blockwise) offloading on DiT modules.",
-    )
-    parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
-    )
     parser.add_argument(
         "--tensor_parallel_size",
         type=int,
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
-    parser.add_argument(
-        "--vae_use_slicing",
-        action="store_true",
-        help="Enable VAE slicing for memory optimization.",
-    )
-    parser.add_argument(
-        "--vae_use_tiling",
-        action="store_true",
-        help="Enable VAE tiling for memory optimization.",
-    )
     return parser.parse_args()
 
 
 def main():
     args = parse_args()
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    device = detect_device_type()
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+
+    # Enable VAE memory optimizations on NPU
+    vae_use_slicing = is_npu()
+    vae_use_tiling = is_npu()
 
     # Configure cache based on backend type
     cache_config = None
@@ -183,13 +162,10 @@ def main():
 
     omni = Omni(
         model=args.model,
-        enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
-        vae_use_slicing=args.vae_use_slicing,
-        vae_use_tiling=args.vae_use_tiling,
+        vae_use_slicing=vae_use_slicing,
+        vae_use_tiling=vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
-        enable_cache_dit_summary=args.enable_cache_dit_summary,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,

examples/offline_inference/text_to_video/text_to_video.md

@@ -26,12 +26,7 @@ Key arguments:
 - `--num_frames`: Number of frames (Wan default is 81).
 - `--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).
-- `--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).
+- `--boundary_ratio`: Boundary split ratio for low/high DiT.
 - `--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`](../../../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.

examples/offline_inference/text_to_video/text_to_video.py

@@ -13,7 +13,7 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.platform_utils import detect_device_type, is_npu
 
 
 def parse_args() -> argparse.Namespace:
@@ -32,64 +32,17 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--width", type=int, default=1280, help="Video width.")
     parser.add_argument("--num_frames", type=int, default=81, help="Number of frames (Wan default is 81).")
     parser.add_argument("--num_inference_steps", type=int, default=40, help="Sampling steps.")
-    parser.add_argument(
-        "--boundary_ratio",
-        type=float,
-        default=0.875,
-        help="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).",
-    )
+    parser.add_argument("--boundary_ratio", type=float, default=0.875, help="Boundary split ratio for low/high DiT.")
     parser.add_argument(
         "--flow_shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
     )
-    parser.add_argument(
-        "--cache_backend",
-        type=str,
-        default=None,
-        choices=["cache_dit"],
-        help=(
-            "Cache backend to use for acceleration. "
-            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer). "
-            "Default: None (no cache acceleration)."
-        ),
-    )
-    parser.add_argument(
-        "--enable-cache-dit-summary",
-        action="store_true",
-        help="Enable cache-dit summary logging after diffusion forward passes.",
-    )
     parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
     parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
-    parser.add_argument(
-        "--vae_use_slicing",
-        action="store_true",
-        help="Enable VAE slicing for memory optimization.",
-    )
-    parser.add_argument(
-        "--vae_use_tiling",
-        action="store_true",
-        help="Enable VAE tiling for memory optimization.",
-    )
-    parser.add_argument(
-        "--enforce_eager",
-        action="store_true",
-        help="Disable torch.compile and force eager execution.",
-    )
     parser.add_argument(
         "--enable-cpu-offload",
         action="store_true",
         help="Enable CPU offloading for diffusion models.",
     )
-    parser.add_argument(
-        "--enable-layerwise-offload",
-        action="store_true",
-        help="Enable layerwise (blockwise) offloading on DiT modules.",
-    )
-    parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
-    )
     parser.add_argument(
         "--ulysses_degree",
         type=int,
@@ -103,44 +56,27 @@ def parse_args() -> argparse.Namespace:
         help="Number of GPUs used for ring sequence parallelism.",
     )
     parser.add_argument(
-        "--cfg_parallel_size",
-        type=int,
-        default=1,
-        choices=[1, 2],
-        help="Number of GPUs used for classifier free guidance parallel size.",
+        "--enforce_eager",
+        action="store_true",
+        help="Disable torch.compile and force eager execution.",
     )
-
     return parser.parse_args()
 
 
 def main():
     args = parse_args()
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    device = detect_device_type()
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+
+    # Enable VAE memory optimizations on NPU
+    vae_use_slicing = is_npu()
+    vae_use_tiling = is_npu()
 
-    # Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
-    cache_config = None
-    if args.cache_backend == "cache_dit":
-        cache_config = {
-            # DBCache parameters [cache-dit only]
-            "Fn_compute_blocks": 1,  # Optimized for single-transformer models
-            "Bn_compute_blocks": 0,  # Number of backward compute blocks
-            "max_warmup_steps": 4,  # Maximum warmup steps (works for few-step models)
-            "max_cached_steps": 20,
-            "residual_diff_threshold": 0.24,  # Higher threshold for more aggressive caching
-            "max_continuous_cached_steps": 3,  # Limit to prevent precision degradation
-            # TaylorSeer parameters [cache-dit only]
-            "enable_taylorseer": False,  # Disabled by default (not suitable for few-step models)
-            "taylorseer_order": 1,  # TaylorSeer polynomial order
-            # SCM (Step Computation Masking) parameters [cache-dit only]
-            "scm_steps_mask_policy": None,  # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
-            "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
-        }
     # Configure parallel settings (only SP is supported for Wan)
     # Note: cfg_parallel and tensor_parallel are not implemented for Wan models
     parallel_config = DiffusionParallelConfig(
         ulysses_degree=args.ulysses_degree,
         ring_degree=args.ring_degree,
-        cfg_parallel_size=args.cfg_parallel_size,
     )
 
     # Check if profiling is requested via environment variable
@@ -148,15 +84,10 @@ def main():
 
     omni = Omni(
         model=args.model,
-        enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
-        vae_use_slicing=args.vae_use_slicing,
-        vae_use_tiling=args.vae_use_tiling,
+        vae_use_slicing=vae_use_slicing,
+        vae_use_tiling=vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
-        cache_backend=args.cache_backend,
-        cache_config=cache_config,
-        enable_cache_dit_summary=args.enable_cache_dit_summary,
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
@@ -172,9 +103,7 @@ def main():
     print(f"  Model: {args.model}")
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Frames: {args.num_frames}")
-    print(
-        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
-    )
+    print(f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}")
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
 

examples/online_serving/bagel/README.md

@@ -1,230 +0,0 @@
-# BAGEL-7B-MoT
-
-## 🛠️ Installation
-
-Please refer to [README.md](../../../README.md)
-
-## Run examples (BAGEL-7B-MoT)
-
-**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
-
-### Launch the Server
-
-```bash
-# Use default configuration
-vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
-```
-
-Or use the convenience script:
-
-```bash
-cd /workspace/vllm-omni/examples/online_serving/bagel
-bash run_server.sh
-```
-
-If you have a custom stage configs file, launch the server with the command below:
-
-```bash
-vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
-```
-
-### Send Multi-modal Request
-
-Get into the bagel folder:
-
-```bash
-cd examples/online_serving/bagel
-```
-
-Send request via Python
-
-```bash
-python openai_chat_client.py --prompt "A cute cat" --modality text2img
-```
-
-The Python client supports the following command-line arguments:
-
-- `--prompt` (or `-p`): Text prompt for generation (default: `A cute cat`)
-- `--output` (or `-o`): Output file path for image results (default: `bagel_output.png`)
-- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
-- `--image-url` (or `-i`): Input image URL or local file path (for img2img/img2text modes)
-- `--modality` (or `-m`): Task modality (default: `text2img`). Options: `text2img`, `img2img`, `img2text`, `text2text`
-- `--height`: Image height in pixels (default: 512)
-- `--width`: Image width in pixels (default: 512)
-- `--steps`: Number of inference steps (default: 25)
-- `--seed`: Random seed (default: 42)
-- `--negative`: Negative prompt for image generation
-
-Example with custom parameters:
-
-```bash
-python openai_chat_client.py \
-    --prompt "A futuristic city" \
-    --modality text2img \
-    --height 768 \
-    --width 768 \
-    --steps 50 \
-    --seed 42 \
-    --negative "blurry, low quality"
-```
-
-## Modality Control
-
-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)
-
-| Modality    | Input        | Output | Description                            |
-| ----------- | ------------ | ------ | -------------------------------------- |
-| `text2img`  | Text         | Image  | Generate images from text prompts      |
-| `img2img`   | Image + Text | Image  | Transform images using text guidance   |
-| `img2text`  | Image + Text | Text   | Generate text descriptions from images |
-| `text2text` | Text         | Text   | Pure text generation                   |
-
-### Text to Image (text2img)
-
-Generate images from text prompts:
-
-**Using Python client**
-
-```bash
-python openai_chat_client.py \
-    --prompt "A beautiful sunset over mountains" \
-    --modality text2img \
-    --output sunset.png \
-    --steps 50
-```
-
-**Using curl**
-
-```bash
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}]}],
-    "modalities": ["image"],
-    "height": 512,
-    "width": 512,
-    "num_inference_steps": 50,
-    "seed": 42
-  }'
-```
-
-
-### Image to Image (img2img)
-
-Transform images based on text prompts:
-
-**Using Python client**
-
-```bash
-python openai_chat_client.py \
-    --prompt "Make the cat stand up" \
-    --modality img2img \
-    --image-url /path/to/input.jpg \
-    --output transformed.png
-```
-
-**Using curl**
-
-```bash
-IMAGE_BASE64=$(base64 -w 0 cat.jpg)
-
-cat <<EOF > payload.json
-{
-    "messages": [{
-      "role": "user",
-      "content": [
-        {"type": "text", "text": "<|im_start|>Make the cat stand up<|im_end|>"},
-        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
-      ]
-    }],
-    "modalities": ["image"],
-    "height": 512,
-    "width": 512,
-    "num_inference_steps": 50,
-    "seed": 42
-}
-EOF
-
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d @payload.json
-
-```
-
-### Image to Text (img2text)
-
-Generate text descriptions from images:
-
-**Using Python client**
-
-```bash
-python openai_chat_client.py \
-    --prompt "Describe this image in detail" \
-    --modality img2text \
-    --image-url /path/to/image.jpg
-```
-
-**Using curl**
-
-```bash
-IMAGE_BASE64=$(base64 -w 0 cat.jpg)
-
-cat <<EOF > payload.json
-{
-  "messages": [{
-    "role": "user",
-    "content": [
-      {"type": "text", "text": "<|im_start|>user\n<|image_pad|>\nDescribe this image in detail<|im_end|>\n<|im_start|>assistant\n"},
-      {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
-    ]
-  }],
-  "modalities": ["text"]
-}
-EOF
-
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d @payload.json
-```
-
-### Text to Text (text2text)
-
-Pure text generation:
-
-**Using Python client**
-
-```bash
-python openai_chat_client.py \
-    --prompt "What is the capital of France?" \
-    --modality text2text
-```
-
-**Using curl**
-
-```bash
-curl http://localhost:8091/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{
-    "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}]
-    "modalities": ["text"]
-  }'
-```
-
-## FAQ
-
-- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
-
-```bash
-sudo apt update
-sudo apt install ffmpeg
-```
-
-- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
-
-| Stage               | VRAM                         |
-| :------------------ | :--------------------------- |
-| Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
-| Stage-1 (DiT)       | **26.50 GiB**                |
-| Total               | **~42 GiB + KV Cache**       |

examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml

@@ -13,7 +13,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.9
       enforce_eager: true

examples/online_serving/qwen3_tts/README.md

@@ -41,14 +41,12 @@ python openai_speech_client.py \
 
 # 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 \
@@ -75,9 +73,6 @@ curl -X POST http://localhost:8000/v1/audio/speech \
         "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

examples/online_serving/qwen3_tts/openai_speech_client.py

@@ -90,7 +90,6 @@ def run_tts_generation(args) -> None:
     if args.x_vector_only:
         payload["x_vector_only_mode"] = True
 
-    print(f"Model: {args.model}")
     print(f"Task type: {args.task_type or 'CustomVoice'}")
     print(f"Text: {args.text}")
     print(f"Voice: {args.voice}")
@@ -111,10 +110,6 @@ def run_tts_generation(args) -> None:
         print(response.text)
         return
 
-    if response.content.decode("utf-8").startswith('{"error"'):
-        print(f"Error: {response.content.decode('utf-8')}")
-        return
-
     # Save audio response
     output_path = args.output or "tts_output.wav"
     with open(output_path, "wb") as f:

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vllm-omni"
-version = "0.14.0"
+version = "0.14.0rc1"
 description = "A framework for efficient model inference with omni-modality models"
 readme = "README.md"
 requires-python = ">=3.10,<3.14"
@@ -39,7 +39,7 @@ dependencies = [
     "cache-dit==1.2.0",
     "tqdm>=4.66.0",
     "torchsde>=0.2.6",  # Required for Stable Audio scheduler
-    "fa3-fwd==0.0.1", # flash attention 3, maintained by @ZJY0516
+    "fa3-fwd", # flash attention 3, maintained by @ZJY0516
     "openai-whisper>=20250625",
     "imageio[ffmpeg]>=2.37.2",
     "onnxruntime>=1.19.0",
@@ -59,8 +59,7 @@ dev = [
     "psutil>=7.2.0",
     "soundfile>=0.13.1",
     "imageio[ffmpeg]>=0.6.0",
-    "opencv-python>=4.12.0.88",
-    "mooncake-transfer-engine==0.3.8.post1"
+    "opencv-python>=4.12.0.88"
 ]
 
 docs = [

tests/benchmarks/test_serve_cli.py

@@ -1,60 +0,0 @@
-import subprocess
-from pathlib import Path
-
-import pytest
-
-from tests.conftest import OmniServer
-
-models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
-stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen3_omni_ci.yaml")]
-
-# Create parameter combinations for model and stage config
-test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
-
-
-@pytest.fixture(scope="module")
-def omni_server(request):
-    """Start vLLM-Omni server as a subprocess with actual model weights.
-    Uses session scope so the server starts only once for the entire test session.
-    Multi-stage initialization can take 10-20+ minutes.
-    """
-    model, stage_config_path = request.param
-
-    print(f"Starting OmniServer with model: {model}")
-    print("This may take 10-20+ minutes for initialization...")
-
-    with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
-        print("OmniServer started successfully")
-        yield server
-        print("OmniServer stopped")
-
-
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_bench_serve_chat(omni_server):
-    command = [
-        "vllm",
-        "bench",
-        "serve",
-        "--omni",
-        "--model",
-        omni_server.model,
-        "--port",
-        str(omni_server.port),
-        "--dataset-name",
-        "random",
-        "--random-input-len",
-        "32",
-        "--random-output-len",
-        "4",
-        "--num-prompts",
-        "5",
-        "--endpoint",
-        "/v1/chat/completions",
-        "--backend",
-        "openai-chat-omni",
-    ]
-    result = subprocess.run(command, capture_output=True, text=True)
-    print(result.stdout)
-    print(result.stderr)
-
-    assert result.returncode == 0, f"Benchmark failed: {result.stderr}"

tests/conftest.py

@@ -10,7 +10,6 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 if "VLLM_TARGET_DEVICE" not in os.environ:
     os.environ["VLLM_TARGET_DEVICE"] = "cpu"
 
-import gc
 import socket
 import subprocess
 import sys
@@ -60,12 +59,13 @@ def clean_gpu_memory_between_tests():
     _run_post_test_cleanup()
 
 
-def _run_pre_test_cleanup(enable_force=False):
-    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
+def _run_pre_test_cleanup():
+    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
         print("GPU cleanup disabled")
         return
 
     print("Pre-test GPU status:")
+    _print_simple_gpu_status()
 
     num_gpus = torch.cuda.device_count()
     if num_gpus > 0:
@@ -74,25 +74,44 @@ def _run_pre_test_cleanup(enable_force=False):
 
             wait_for_gpu_memory_to_clear(
                 devices=list(range(num_gpus)),
-                threshold_ratio=0.05,
+                threshold_ratio=0.1,
             )
         except Exception as e:
             print(f"Pre-test cleanup note: {e}")
 
 
-def _run_post_test_cleanup(enable_force=False):
-    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
-        print("GPU cleanup disabled")
+def _run_post_test_cleanup():
+    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
         return
 
+    import gc
+
     if torch.cuda.is_available():
         gc.collect()
         torch.cuda.empty_cache()
 
         print("Post-test GPU status:")
+        _print_simple_gpu_status()
         _print_gpu_processes()
 
 
+def _print_simple_gpu_status():
+    """Print simple GPU memory status"""
+    if not torch.cuda.is_available():
+        print("  CUDA not available")
+        return
+
+    num_devices = torch.cuda.device_count()
+    for device_id in range(num_devices):
+        try:
+            torch.cuda.set_device(device_id)
+            allocated = torch.cuda.memory_allocated(device_id) / (1024**2)
+            reserved = torch.cuda.memory_reserved(device_id) / (1024**2)
+            print(f"  GPU {device_id}: Allocated: {allocated:.1f}MB, Reserved: {reserved:.1f}MB")
+        except Exception:
+            print(f"  GPU {device_id}: Error reading status")
+
+
 def _print_gpu_processes():
     """Print GPU information including nvidia-smi and system processes"""
 
@@ -604,36 +623,6 @@ def convert_audio_to_text(audio_data):
         return ""
 
 
-def merge_base64_and_convert_to_text(base64_list):
-    """
-    Merge a list of base64 encoded audio data and convert to text.
-    """
-    import whisper
-    from pydub import AudioSegment
-
-    merged_audio = None
-    for base64_data in base64_list:
-        audio_data = base64.b64decode(base64_data.split(",", 1)[-1])
-        seg = AudioSegment.from_file(io.BytesIO(audio_data))
-        if merged_audio is None:
-            merged_audio = seg
-        else:
-            merged_audio += seg
-    output_path = f"./test_{int(time.time())}"
-    merged_audio.export(output_path, format="wav")
-    model = whisper.load_model("base")
-    text = model.transcribe(
-        output_path,
-        temperature=0.0,
-        word_timestamps=True,
-        condition_on_previous_text=False,
-    )["text"]
-    if text:
-        return text
-    else:
-        return ""
-
-
 def modify_stage_config(
     yaml_path: str,
     updates: dict[str, Any],
@@ -882,9 +871,6 @@ class OmniServer:
         *,
         env_dict: dict[str, str] | None = None,
     ) -> None:
-        _run_pre_test_cleanup(enable_force=True)
-        _run_post_test_cleanup(enable_force=True)
-        cleanup_dist_env_and_memory()
         self.model = model
         self.serve_args = serve_args
         self.env_dict = env_dict
@@ -920,7 +906,7 @@ class OmniServer:
         )
 
         # Wait for server to be ready
-        max_wait = 1200  # 20 minutes
+        max_wait = 600  # 10 minutes
         start_time = time.time()
         while time.time() - start_time < max_wait:
             try:
@@ -1000,6 +986,5 @@ class OmniServer:
     def __exit__(self, exc_type, exc_val, exc_tb):
         if self.proc:
             self._kill_process_tree(self.proc.pid)
-        _run_pre_test_cleanup(enable_force=True)
-        _run_post_test_cleanup(enable_force=True)
+        _run_post_test_cleanup()
         cleanup_dist_env_and_memory()

tests/diffusion/attention/test_attention_sp.py

@@ -6,6 +6,7 @@ import tempfile
 
 import pytest
 import torch
+from vllm.platforms import current_platform
 
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.data import (
@@ -18,7 +19,15 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.platform_utils import detect_device_type
+
+device_type = detect_device_type()
+if device_type == "cuda":
+    torch_device = torch.cuda
+elif device_type == "npu":
+    torch_device = torch.npu
+else:
+    raise ValueError(f"Unsupported device type: {device_type} for this test script! Expected GPU or NPU.")
 
 
 def update_environment_variables(envs_dict: dict[str, str]):
@@ -162,7 +171,7 @@ def test_sequence_parallel(
     sequence_parallel_size = ulysses_degree * ring_degree
 
     # Skip if not enough GPUs available
-    available_gpus = current_omni_platform.get_device_count()
+    available_gpus = torch_device.device_count()
     if available_gpus < sequence_parallel_size:
         pytest.skip(f"Test requires {sequence_parallel_size} GPUs but only {available_gpus} available")
 
@@ -329,13 +338,13 @@ def ulysses_attention_on_test_model(
     """Run Ulysses attention test on a test model and save results for comparison."""
     # Use fixed seed for reproducibility across baseline and SP runs
     RANDOM_SEED = 42
-    current_omni_platform.seed_everything(RANDOM_SEED)
+    current_platform.seed_everything(RANDOM_SEED)
 
     mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})"
     print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}")
 
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
+    device = torch.device(f"{device_type}:{local_rank}")
+    torch_device.set_device(device)
     torch.set_default_device(device)
     torch.set_default_dtype(dtype)
 
@@ -414,7 +423,7 @@ def ulysses_attention_on_test_model(
             # Generate and save full input data with fixed seed
             # Reinitialize RNG to ensure reproducibility
             torch.manual_seed(42)
-            current_omni_platform.seed_everything(42)
+            torch_device.manual_seed_all(42)
             full_hidden_states = torch.randn(
                 (batch_size, seq_len, hidden_size),
                 dtype=dtype,

tests/diffusion/distributed/test_cfg_parallel.py

@@ -1,423 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for CFG (Classifier-Free Guidance) parallel functionality.
-
-This test verifies that predict_noise_maybe_with_cfg produces numerically
-equivalent results with and without CFG parallel using fixed random inputs.
-"""
-
-import os
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-from vllm_omni.diffusion.distributed.parallel_state import (
-    destroy_distributed_env,
-    get_classifier_free_guidance_rank,
-    get_classifier_free_guidance_world_size,
-    init_distributed_environment,
-    initialize_model_parallel,
-)
-from vllm_omni.platforms import current_omni_platform
-
-
-def update_environment_variables(envs_dict: dict[str, str]):
-    """Update multiple environment variables."""
-    for k, v in envs_dict.items():
-        os.environ[k] = v
-
-
-class SimpleTransformer(torch.nn.Module):
-    """Simple transformer model for testing with random initialization.
-
-    Contains:
-    - Input projection (conv to hidden_dim)
-    - QKV projection layers
-    - Self-attention layer
-    - Output projection
-    """
-
-    def __init__(self, in_channels: int = 4, hidden_dim: int = 128, num_heads: int = 8):
-        super().__init__()
-        self.hidden_dim = hidden_dim
-        self.num_heads = num_heads
-        self.head_dim = hidden_dim // num_heads
-
-        assert hidden_dim % num_heads == 0, "hidden_dim must be divisible by num_heads"
-
-        # Input projection: (B, C, H, W) -> (B, hidden_dim, H, W)
-        self.input_proj = torch.nn.Conv2d(in_channels, hidden_dim, 1)
-
-        # QKV projection layers
-        self.q_proj = torch.nn.Linear(hidden_dim, hidden_dim)
-        self.k_proj = torch.nn.Linear(hidden_dim, hidden_dim)
-        self.v_proj = torch.nn.Linear(hidden_dim, hidden_dim)
-
-        # Output projection after attention
-        self.out_proj = torch.nn.Linear(hidden_dim, hidden_dim)
-
-        # Final output projection: (B, hidden_dim, H, W) -> (B, C, H, W)
-        self.final_proj = torch.nn.Conv2d(hidden_dim, in_channels, 1)
-
-        # Layer norm
-        self.norm1 = torch.nn.LayerNorm(hidden_dim)
-        self.norm2 = torch.nn.LayerNorm(hidden_dim)
-
-    def forward(self, x: torch.Tensor, **kwargs) -> tuple[torch.Tensor]:
-        """Forward pass with self-attention.
-
-        Args:
-            x: Input tensor of shape (B, C, H, W)
-
-        Returns:
-            Output tensor of shape (B, C, H, W)
-        """
-        B, C, H, W = x.shape
-
-        # Input projection
-        x = self.input_proj(x)  # (B, hidden_dim, H, W)
-
-        # Reshape to sequence: (B, hidden_dim, H, W) -> (B, H*W, hidden_dim)
-        x = x.flatten(2).transpose(1, 2)  # (B, H*W, hidden_dim)
-
-        # Self-attention with residual connection
-        residual = x
-        x = self.norm1(x)
-
-        # QKV projection
-        q = self.q_proj(x)  # (B, H*W, hidden_dim)
-        k = self.k_proj(x)  # (B, H*W, hidden_dim)
-        v = self.v_proj(x)  # (B, H*W, hidden_dim)
-
-        # Reshape for multi-head attention: (B, H*W, hidden_dim) -> (B, num_heads, H*W, head_dim)
-        seq_len = H * W
-        q = q.view(B, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
-        k = k.view(B, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
-        v = v.view(B, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
-
-        # Scaled dot-product attention
-        scale = self.head_dim**-0.5
-        attn_scores = torch.matmul(q, k.transpose(-2, -1)) * scale  # (B, num_heads, H*W, H*W)
-        attn_weights = torch.nn.functional.softmax(attn_scores, dim=-1)
-        attn_output = torch.matmul(attn_weights, v)
-        attn_output = attn_output.transpose(1, 2).contiguous().view(B, seq_len, self.hidden_dim)
-
-        attn_output = self.out_proj(attn_output)
-
-        x = residual + attn_output
-        residual = x
-        x = self.norm2(x)
-        x = residual + x
-        x = x.transpose(1, 2).view(B, self.hidden_dim, H, W)
-
-        out = self.final_proj(x)
-
-        return (out,)
-
-
-class TestCFGPipeline(CFGParallelMixin):
-    """Test pipeline using CFGParallelMixin."""
-
-    def __init__(self, in_channels: int = 4, hidden_dim: int = 128, seed: int = 42):
-        # Set seed BEFORE creating transformer to ensure consistent layer initialization
-        torch.manual_seed(seed)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(seed)
-
-        self.transformer = SimpleTransformer(in_channels, hidden_dim)
-
-        # Re-initialize all parameters with fixed seed for full reproducibility
-        torch.manual_seed(seed)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(seed)
-        for param in self.transformer.parameters():
-            torch.nn.init.normal_(param, mean=0.0, std=0.02)
-
-
-def _test_cfg_parallel_worker(
-    local_rank: int,
-    world_size: int,
-    cfg_parallel_size: int,
-    dtype: torch.dtype,
-    test_config: dict,
-    result_queue: torch.multiprocessing.Queue,
-):
-    """Worker function for CFG parallel test."""
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
-
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "29502",
-        }
-    )
-
-    init_distributed_environment()
-    initialize_model_parallel(cfg_parallel_size=cfg_parallel_size)
-
-    cfg_rank = get_classifier_free_guidance_rank()
-    cfg_world_size = get_classifier_free_guidance_world_size()
-
-    assert cfg_world_size == cfg_parallel_size
-
-    # Create pipeline with same seed to ensure identical model weights across all ranks
-    # Note: model_seed is set inside TestCFGPipeline.__init__
-    pipeline = TestCFGPipeline(
-        in_channels=test_config["channels"],
-        hidden_dim=test_config["hidden_dim"],
-        seed=test_config["model_seed"],
-    )
-    pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
-    pipeline.transformer.eval()  # Set to eval mode for deterministic behavior
-
-    # Create fixed inputs with explicit seed setting for reproducibility
-    # Set both CPU and CUDA seeds to ensure identical inputs across all ranks
-    torch.manual_seed(test_config["input_seed"])
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(test_config["input_seed"])
-
-    batch_size = test_config["batch_size"]
-    channels = test_config["channels"]
-    height = test_config["height"]
-    width = test_config["width"]
-
-    # Positive input
-    positive_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
-
-    # Negative input with different seed
-    torch.manual_seed(test_config["input_seed"] + 1)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
-    negative_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
-
-    # Prepare kwargs for predict_noise_maybe_with_cfg
-    positive_kwargs = {"x": positive_input}
-    negative_kwargs = {"x": negative_input}
-
-    with torch.no_grad():
-        # Call predict_noise_maybe_with_cfg
-        noise_pred = pipeline.predict_noise_maybe_with_cfg(
-            do_true_cfg=True,
-            true_cfg_scale=test_config["cfg_scale"],
-            positive_kwargs=positive_kwargs,
-            negative_kwargs=negative_kwargs,
-            cfg_normalize=test_config["cfg_normalize"],
-        )
-
-    # Only rank 0 has valid output in CFG parallel mode
-    if cfg_rank == 0:
-        assert noise_pred is not None
-        result_queue.put(noise_pred.cpu())
-    else:
-        assert noise_pred is None
-
-    destroy_distributed_env()
-
-
-def _test_cfg_sequential_worker(
-    local_rank: int,
-    world_size: int,
-    dtype: torch.dtype,
-    test_config: dict,
-    result_queue: torch.multiprocessing.Queue,
-):
-    """Worker function for sequential CFG test (baseline)."""
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
-
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "29503",
-        }
-    )
-
-    init_distributed_environment()
-    initialize_model_parallel(cfg_parallel_size=1)  # No CFG parallel
-
-    cfg_world_size = get_classifier_free_guidance_world_size()
-    assert cfg_world_size == 1
-
-    # Create pipeline with same seed to ensure identical model weights as CFG parallel
-    # Note: model_seed is set inside TestCFGPipeline.__init__
-    pipeline = TestCFGPipeline(
-        in_channels=test_config["channels"],
-        hidden_dim=test_config["hidden_dim"],
-        seed=test_config["model_seed"],
-    )
-    pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
-    pipeline.transformer.eval()
-
-    # Create fixed inputs (same seed as CFG parallel to ensure identical inputs)
-    # Set both CPU and CUDA seeds for full reproducibility
-    torch.manual_seed(test_config["input_seed"])
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(test_config["input_seed"])
-
-    batch_size = test_config["batch_size"]
-    channels = test_config["channels"]
-    height = test_config["height"]
-    width = test_config["width"]
-
-    # Positive input
-    positive_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
-
-    # Negative input with different seed
-    torch.manual_seed(test_config["input_seed"] + 1)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
-    negative_input = torch.randn(batch_size, channels, height, width, dtype=dtype, device=device)
-
-    positive_kwargs = {"x": positive_input}
-    negative_kwargs = {"x": negative_input}
-
-    with torch.no_grad():
-        noise_pred = pipeline.predict_noise_maybe_with_cfg(
-            do_true_cfg=True,
-            true_cfg_scale=test_config["cfg_scale"],
-            positive_kwargs=positive_kwargs,
-            negative_kwargs=negative_kwargs,
-            cfg_normalize=test_config["cfg_normalize"],
-        )
-
-    # Sequential CFG always returns output
-    assert noise_pred is not None
-    result_queue.put(noise_pred.cpu())
-
-    destroy_distributed_env()
-
-
-@pytest.mark.parametrize("cfg_parallel_size", [2])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("cfg_normalize", [False, True])
-def test_predict_noise_maybe_with_cfg(cfg_parallel_size: int, dtype: torch.dtype, batch_size: int, cfg_normalize: bool):
-    """
-    Test that predict_noise_maybe_with_cfg produces identical results
-    with and without CFG parallel.
-
-    Args:
-        cfg_parallel_size: Number of GPUs for CFG parallel
-        dtype: Data type for computation
-        batch_size: Batch size for testing
-        cfg_normalize: Whether to normalize CFG output
-    """
-    available_gpus = current_omni_platform.get_device_count()
-    if available_gpus < cfg_parallel_size:
-        pytest.skip(f"Test requires {cfg_parallel_size} GPUs but only {available_gpus} available")
-
-    test_config = {
-        "batch_size": batch_size,
-        "channels": 4,
-        "height": 16,
-        "width": 16,
-        "hidden_dim": 128,
-        "cfg_scale": 7.5,
-        "cfg_normalize": cfg_normalize,
-        "model_seed": 42,  # Fixed seed for model initialization
-        "input_seed": 123,  # Fixed seed for input generation
-    }
-
-    mp_context = torch.multiprocessing.get_context("spawn")
-
-    manager = mp_context.Manager()
-    baseline_queue = manager.Queue()
-    cfg_parallel_queue = manager.Queue()
-
-    # Run baseline (sequential CFG) on single GPU
-    torch.multiprocessing.spawn(
-        _test_cfg_sequential_worker,
-        args=(1, dtype, test_config, baseline_queue),
-        nprocs=1,
-    )
-
-    # Run CFG parallel on multiple GPUs
-    torch.multiprocessing.spawn(
-        _test_cfg_parallel_worker,
-        args=(cfg_parallel_size, cfg_parallel_size, dtype, test_config, cfg_parallel_queue),
-        nprocs=cfg_parallel_size,
-    )
-
-    # Get results from queues
-    baseline_output = baseline_queue.get()
-    cfg_parallel_output = cfg_parallel_queue.get()
-
-    # Verify shapes match
-    assert baseline_output.shape == cfg_parallel_output.shape, (
-        f"Shape mismatch: baseline {baseline_output.shape} vs CFG parallel {cfg_parallel_output.shape}"
-    )
-
-    # Verify numerical equivalence with appropriate tolerances
-    if dtype == torch.float32:
-        rtol, atol = 1e-5, 1e-5
-    elif dtype == torch.bfloat16:
-        rtol, atol = 1e-2, 1e-2
-    else:
-        rtol, atol = 1e-3, 1e-3
-
-    torch.testing.assert_close(
-        cfg_parallel_output,
-        baseline_output,
-        rtol=rtol,
-        atol=atol,
-        msg=(
-            f"CFG parallel output differs from sequential CFG\n"
-            f"  dtype={dtype}, batch_size={batch_size}, cfg_normalize={cfg_normalize}\n"
-            f"  Max diff: {(cfg_parallel_output - baseline_output).abs().max().item():.6e}"
-        ),
-    )
-
-    print(
-        f"✓ Test passed: cfg_size={cfg_parallel_size}, dtype={dtype}, "
-        f"batch_size={batch_size}, cfg_normalize={cfg_normalize}"
-    )
-
-
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-def test_predict_noise_without_cfg(dtype: torch.dtype):
-    """
-    Test predict_noise_maybe_with_cfg when do_true_cfg=False.
-
-    When CFG is disabled, only the positive branch should be computed.
-    This test runs on a single GPU without distributed environment.
-    """
-    available_gpus = current_omni_platform.get_device_count()
-    if available_gpus < 1:
-        pytest.skip("Test requires at least 1 GPU")
-
-    device = torch.device(f"{current_omni_platform.device_type}:0")
-    current_omni_platform.set_device(device)
-
-    # Create pipeline without distributed environment
-    pipeline = TestCFGPipeline(in_channels=4, hidden_dim=128, seed=42)
-    pipeline.transformer = pipeline.transformer.to(device=device, dtype=dtype)
-    pipeline.transformer.eval()
-
-    # Set seed for input generation
-    torch.manual_seed(123)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(123)
-    positive_input = torch.randn(1, 4, 16, 16, dtype=dtype, device=device)
-
-    with torch.no_grad():
-        noise_pred = pipeline.predict_noise_maybe_with_cfg(
-            do_true_cfg=False,  # No CFG
-            true_cfg_scale=7.5,
-            positive_kwargs={"x": positive_input},
-            negative_kwargs=None,
-            cfg_normalize=False,
-        )
-
-    # Should always return output when do_true_cfg=False
-    assert noise_pred is not None
-    assert noise_pred.shape == (1, 4, 16, 16)
-
-    print(f"✓ Test passed: predict_noise without CFG (dtype={dtype})")

tests/diffusion/distributed/test_comm.py

@@ -14,7 +14,15 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.platform_utils import detect_device_type
+
+device_type = detect_device_type()
+if device_type == "cuda":
+    torch_device = torch.cuda
+elif device_type == "npu":
+    torch_device = torch.npu
+else:
+    raise ValueError(f"Unsupported device type: {device_type} for this test script! Expected GPU or NPU.")
 
 
 def update_environment_variables(envs_dict: dict[str, str]):
@@ -41,7 +49,7 @@ def test_4d_identity(
 ):
     """Test that two consecutive all-to-all operations return the original input."""
     # Skip if not enough GPUs available
-    available_gpus = current_omni_platform.get_device_count()
+    available_gpus = torch_device.device_count()
     if available_gpus < world_size:
         pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
 
@@ -77,8 +85,8 @@ def _test_4d_identity_worker(
 ):
     """Worker function for test_4d_identity."""
     # Set device
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
+    device = torch.device(f"{device_type}:{local_rank}")
+    torch_device.set_device(device)
 
     # Set environment variables for distributed training
     update_environment_variables(
@@ -175,7 +183,7 @@ def test_5d_identity(
 ):
     """Test that two consecutive all-to-all operations return the original input."""
     # Skip if not enough GPUs available
-    available_gpus = current_omni_platform.get_device_count()
+    available_gpus = torch_device.device_count()
     if available_gpus < world_size:
         pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
 
@@ -211,8 +219,8 @@ def _test_5d_identity_worker(
 ):
     """Worker function for test_5d_identity."""
     # Set device
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
+    device = torch.device(f"{device_type}:{local_rank}")
+    torch_device.set_device(device)
 
     # Set environment variables for distributed training
     update_environment_variables(
@@ -308,7 +316,7 @@ def test_ring_p2p(
 ):
     """Test Ring P2P communication (send_recv)."""
     # Skip if not enough GPUs available
-    available_gpus = current_omni_platform.get_device_count()
+    available_gpus = torch_device.device_count()
     if available_gpus < world_size:
         pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
 
@@ -331,8 +339,8 @@ def _test_ring_p2p_worker(
     import sys
 
     # Set device
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
+    device = torch.device(f"{device_type}:{local_rank}")
+    torch_device.set_device(device)
 
     # Set env vars
     # Use a different port to avoid conflict with other tests if run in parallel

tests/diffusion/test_gpu_diffusion_worker.py

@@ -2,9 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Unit tests for DiffusionWorker class.
+Unit tests for GPUDiffusionWorker class.
 
-This module tests the DiffusionWorker implementation:
+This module tests the GPUDiffusionWorker implementation:
 - load_weights: Loading model weights
 - sleep: Putting worker into sleep mode (levels 1 and 2)
 - wake_up: Waking worker from sleep mode
@@ -15,7 +15,7 @@ from unittest.mock import Mock, patch
 import pytest
 import torch
 
-from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
+from vllm_omni.diffusion.worker.gpu_diffusion_worker import GPUDiffusionWorker
 
 
 @pytest.fixture
@@ -33,19 +33,18 @@ def mock_od_config():
 
 @pytest.fixture
 def mock_gpu_worker(mock_od_config):
-    """Create a DiffusionWorker with mocked initialization."""
-    with patch.object(DiffusionWorker, "init_device"):
-        worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
+    """Create a GPUDiffusionWorker with mocked initialization."""
+    with patch.object(GPUDiffusionWorker, "init_device"):
+        worker = GPUDiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
         # Mock the model_runner with pipeline
         worker.model_runner = Mock()
         worker.model_runner.pipeline = Mock()
-        worker.device = torch.device("cuda", 0)
         worker._sleep_saved_buffers = {}
         return worker
 
 
-class TestDiffusionWorkerLoadWeights:
-    """Test DiffusionWorker.load_weights method."""
+class TestGPUDiffusionWorkerLoadWeights:
+    """Test GPUDiffusionWorker.load_weights method."""
 
     def test_load_weights_calls_pipeline(self, mock_gpu_worker):
         """Test that load_weights delegates to model_runner.load_weights."""
@@ -76,21 +75,20 @@ class TestDiffusionWorkerLoadWeights:
         assert result == set()
 
 
-class TestDiffusionWorkerSleep:
-    """Test DiffusionWorker.sleep method."""
+class TestGPUDiffusionWorkerSleep:
+    """Test GPUDiffusionWorker.sleep method."""
 
-    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_gpu_worker):
+    def test_sleep_level_1(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
         """Test sleep mode level 1 (offload weights only)."""
         # Setup memory info mocks
-        # Before sleep: 1GB free
-        # After sleep: 3GB free (freed 2GB)
-        mock_platform.get_free_memory.side_effect = [
-            1 * 1024**3,  # Before sleep
-            3 * 1024**3,  # After sleep
+        # Before sleep: 1GB free, 8GB total
+        # After sleep: 3GB free, 8GB total (freed 2GB)
+        mock_mem_info.side_effect = [
+            (1 * 1024**3, 8 * 1024**3),  # Before sleep
+            (3 * 1024**3, 8 * 1024**3),  # After sleep
         ]
-        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         # Setup allocator mock
         mock_allocator = Mock()
@@ -106,16 +104,15 @@ class TestDiffusionWorkerSleep:
         # Verify buffers were NOT saved (level 1 doesn't save buffers)
         assert len(mock_gpu_worker._sleep_saved_buffers) == 0
 
-    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_gpu_worker):
+    def test_sleep_level_2(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
         """Test sleep mode level 2 (offload all, save buffers)."""
         # Setup memory info mocks
-        mock_platform.get_free_memory.side_effect = [
-            1 * 1024**3,  # Before sleep
-            5 * 1024**3,  # After sleep (freed 4GB)
+        mock_mem_info.side_effect = [
+            (1 * 1024**3, 8 * 1024**3),  # Before sleep
+            (5 * 1024**3, 8 * 1024**3),  # After sleep (freed 4GB)
         ]
-        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         # Setup allocator mock
         mock_allocator = Mock()
@@ -144,16 +141,15 @@ class TestDiffusionWorkerSleep:
         assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
         assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
 
-    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_platform, mock_gpu_worker):
+    def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
         """Test that sleep validates memory was actually freed."""
         # Simulate memory increase (should trigger assertion error)
-        mock_platform.get_free_memory.side_effect = [
-            3 * 1024**3,  # Before sleep: 3GB free
-            1 * 1024**3,  # After sleep: 1GB free (negative freed!)
+        mock_mem_info.side_effect = [
+            (3 * 1024**3, 8 * 1024**3),  # Before sleep: 3GB free
+            (1 * 1024**3, 8 * 1024**3),  # After sleep: 1GB free (negative freed!)
         ]
-        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         mock_allocator = Mock()
         mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
@@ -164,8 +160,8 @@ class TestDiffusionWorkerSleep:
             mock_gpu_worker.sleep(level=1)
 
 
-class TestDiffusionWorkerWakeUp:
-    """Test DiffusionWorker.wake_up method."""
+class TestGPUDiffusionWorkerWakeUp:
+    """Test GPUDiffusionWorker.wake_up method."""
 
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
     def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker):

tests/distributed/omni_connectors/test_kv_flow.py

@@ -1,251 +1,194 @@
-import pytest
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
 import torch
 
-from tests.utils import hardware_test
+from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
-    OmniKVCacheConfig,
-    OmniKVTransferManager,
-)
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
 
 
-class MockConnector:
+class MockGPUARModelRunner(GPUARModelRunner):
+    """Subclass to bypass heavy initialization."""
+
+    def __init__(self, kv_caches, input_batch):
+        self.kv_caches = kv_caches
+        self.input_batch = input_batch
+        self.device = "cpu"
+        self.cache_config = MagicMock()
+        self.cache_config.block_size = 16
+        self.cache_config.cache_dtype = "auto"
+        self.logger = MagicMock()
+
+
+class MockBagelPipeline(BagelPipeline):
+    """Subclass to bypass heavy initialization."""
+
     def __init__(self):
-        self.store = {}
-
-    def put(self, from_stage, to_stage, put_key, data):
-        # The manager now passes full key as put_key
-        key = f"{from_stage}->{to_stage}:{put_key}"
-        self.store[key] = data
-        return True, len(str(data)), None  # (success, size, metadata)
-
-    def get(self, from_stage, to_stage, get_key, metadata=None):
-        # The manager now passes full key as get_key
-        key = f"{from_stage}->{to_stage}:{get_key}"
-        if key in self.store:
-            return self.store[key], len(str(self.store[key]))
-        return None
+        self.device = "cpu"
+        self.od_config = MagicMock()
+        self.bagel = MagicMock()
+        self.tokenizer = MagicMock()
+        self.language_model = MagicMock()
+        self.new_token_ids = {}
 
 
-@pytest.fixture
-def mock_connector():
-    return MockConnector()
+class TestKVFlow(unittest.TestCase):
+    def setUp(self):
+        # Common constants
+        self.num_layers = 2
+        self.num_heads = 4
+        self.head_dim = 16
+        self.block_size = 8
+        self.x = 8  # PagedAttention factor
+        self.seq_len = 20
+        self.req_id = "req_test_1"
+
+    def test_sender_extraction_logic(self):
+        """Test extraction logic in GPUARModelRunner."""
+        if GPUARModelRunner is object:
+            self.skipTest("vLLM not installed")
+
+        # 1. Setup KV Cache (List of tuples (K, V))
+        # Shape: [num_blocks, block_size, num_heads, head_dim] matching 4D expectation
+        num_blocks = 10
+        kv_caches = []
+        for _ in range(self.num_layers):
+            k_cache = torch.randn(num_blocks, self.block_size, self.num_heads, self.head_dim)
+            v_cache = torch.randn(num_blocks, self.block_size, self.num_heads, self.head_dim)
+            # Stack K and V to create [2, num_blocks, block_size, n_heads, head_dim]
+            layer_cache = torch.stack([k_cache, v_cache], dim=0)
+            kv_caches.append(layer_cache)
+
+        # 2. Setup Input Batch Mock
+        block_ids = [1, 3, 5]
+        mock_input_batch = MagicMock()
+        mock_input_batch.req_id_to_index = {self.req_id: 0}
+        mock_input_batch.block_table.get_row.return_value = block_ids
+
+        # 3. Instantiate Runner
+        runner = MockGPUARModelRunner(kv_caches, mock_input_batch)
+
+        # 4. Run Extraction
+        # calling _extract_kv_cache directly: (req_id, block_ids, seq_len)
+        result = runner._extract_kv_cache(self.req_id, block_ids, self.seq_len)
+
+        # 5. Verify Result
+        self.assertIsNotNone(result)
+        self.assertEqual(result.request_id, self.req_id)
+
+        # Check keys "key_cache" and "value_cache" exist (length 2)
+        self.assertEqual(len(result.layer_blocks["key_cache"]), 2)
+        self.assertEqual(len(result.layer_blocks["value_cache"]), 2)
+
+        # Check Tensor Shape: [seq_len, num_heads, head_dim]
+        # result.layer_blocks["key_cache"] is a list of tensors
+        expected_shape = (self.seq_len, self.num_heads, self.head_dim)
+        self.assertEqual(result.layer_blocks["key_cache"][0].shape, expected_shape)
+
+        return result  # Return for use in next test
+
+    def test_receiver_injection_logic(self):
+        """Test injection logic in BagelPipeline."""
+        if BagelPipeline is object:
+            self.skipTest("vLLM not installed")
+
+        # 1. Get Data from Sender Test (simulate transfer)
+        # Re-create data manually to be independent
+        key_cache = []
+        value_cache = []
+        expected_shape = (self.seq_len, self.num_heads, self.head_dim)
+        for i in range(self.num_layers):
+            key_cache.append(torch.randn(expected_shape))
+            value_cache.append(torch.randn(expected_shape))
+
+        layer_blocks = {"key_cache": key_cache, "value_cache": value_cache}
+
+        transfer_data = MagicMock()  # Mock KVCacheTransferData
+        transfer_data.layer_blocks = layer_blocks
+        transfer_data.metadata = {"kv_lens": [self.seq_len], "ropes": [0]}
+
+        # 2. Setup Request with Injected Data
+        sp = OmniDiffusionSamplingParams(
+            past_key_values=SimpleNamespace(**layer_blocks),
+            kv_metadata=transfer_data.metadata,
+        )
+
+        req = OmniDiffusionRequest(["test"], sp)
+
+        # 3. Setup Pipeline
+        pipeline = MockBagelPipeline()
+
+        # Mock Bagel's NaiveCache
+        class RealNaiveCache:  # Minimal impl
+            def __init__(self, n):
+                self.key_cache = {i: None for i in range(n)}
+                self.value_cache = {i: None for i in range(n)}
+
+        pipeline.bagel.config.llm_config.num_hidden_layers = self.num_layers
+
+        captured_context = {}
+
+        def mock_prepare_prompts(curr_kvlens, curr_rope, **kwargs):
+            # Capture the state passed to prepare_prompts
+            captured_context["kv_lens"] = curr_kvlens
+            captured_context["ropes"] = curr_rope
+            return {}, [0], [0]
+
+        pipeline.bagel.prepare_prompts = MagicMock(side_effect=mock_prepare_prompts)
+        with patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.NaiveCache") as MockNaiveCacheCls:
+            # Setup the instance returned by the constructor
+            mock_cache_instance = RealNaiveCache(self.num_layers)
+            MockNaiveCacheCls.return_value = mock_cache_instance
+
+        # Verification of Injection Logic (Simulation)
+        current_cache = RealNaiveCache(self.num_layers)
+
+        # --- Logic from Source Code ---
+        injected_kv = req.sampling_params.past_key_values
+        if isinstance(current_cache, RealNaiveCache) and hasattr(injected_kv, "key_cache"):
+            # Assuming injected_kv is SimpleNamespace or object with list attrs
+            for layer_idx in range(len(injected_kv.key_cache)):
+                if injected_kv.key_cache[layer_idx] is not None:
+                    k_tensor = injected_kv.key_cache[layer_idx]
+                    v_tensor = injected_kv.value_cache[layer_idx]
+
+                    if k_tensor.device != pipeline.device:
+                        k_tensor = k_tensor.to(pipeline.device)
+                    if v_tensor.device != pipeline.device:
+                        v_tensor = v_tensor.to(pipeline.device)
+
+                    current_cache.key_cache[layer_idx] = k_tensor
+                    current_cache.value_cache[layer_idx] = v_tensor
+
+        self.assertTrue(torch.allclose(current_cache.key_cache[0], layer_blocks["key_cache"][0]))
+        self.assertTrue(torch.allclose(current_cache.value_cache[1], layer_blocks["value_cache"][1]))
+
+    def test_integration(self):
+        """Simulate the flow from Sender -> Connector (Dict) -> Receiver."""
+        if GPUARModelRunner is object or BagelPipeline is object:
+            self.skipTest("vLLM not installed")
+
+        # 1. Sender (Extraction)
+        runner_test_result = self.test_sender_extraction_logic()  # Get KVCacheTransferData
+
+        # 2. Connector (Serialize/Deserialize Simulation)
+        # KVCacheTransferData has to_dict
+        data_dict = runner_test_result.to_dict()
+
+        # 3. Receiver (Request Setup)
+        sp = OmniDiffusionSamplingParams(
+            past_key_values=data_dict["layer_blocks"],
+            kv_metadata=data_dict["metadata"],
+        )
+        req = OmniDiffusionRequest(["integration_test"], sp)  # noqa: F841
+
+        # 4. Receiver (Injection Simulation)
+        # Use the logic verification again
+        pass
 
 
-@pytest.fixture
-def kv_config():
-    return OmniKVCacheConfig(
-        connector_config={"type": "mock"},
-        from_stage="stage1",
-        to_stage="stage2",
-        stage_id="stage2",  # Acting as receiver for some tests
-        need_recv_cache=True,
-        need_send_cache=True,
-        recv_timeout=1.0,  # Short timeout for tests
-    )
-
-
-@pytest.fixture
-def common_constants():
-    return {
-        "num_layers": 2,
-        "num_heads": 4,
-        "head_dim": 16,
-        "block_size": 8,
-        "seq_len": 20,
-        "req_id": "req_test_1",
-    }
-
-
-@pytest.mark.cache
-@hardware_test(
-    res={"cuda": "L4"},
-    num_cards=2,
-)
-def test_manager_extraction(kv_config, mock_connector, common_constants):
-    """Test extraction and sending logic in OmniKVTransferManager."""
-    num_layers = common_constants["num_layers"]
-    block_size = common_constants["block_size"]
-    num_heads = common_constants["num_heads"]
-    head_dim = common_constants["head_dim"]
-    seq_len = common_constants["seq_len"]
-    req_id = common_constants["req_id"]
-
-    num_blocks = 10
-    kv_caches = []
-    for _ in range(num_layers):
-        k_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
-        v_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
-        # Stack K and V to create [2, num_blocks, block_size, n_heads, head_dim]
-        layer_cache = torch.stack([k_cache, v_cache], dim=0)
-        kv_caches.append(layer_cache)
-
-    block_ids = [1, 3, 5]
-    finished_reqs = {req_id: {"block_ids": block_ids, "seq_len": seq_len}}
-
-    manager = OmniKVTransferManager(kv_config)
-    # Mock the connector factory or injection
-    manager._connector = mock_connector
-
-    processed = manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32")
-
-    assert req_id in processed
-
-    # Check if data was put into connector
-    # Manager builds full key: omni_{from}_to_{to}_kv_cache_{req_id}
-    full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}"
-    expected_key = f"stage1->stage2:{full_request_id}"
-    assert expected_key in mock_connector.store
-
-    data = mock_connector.store[expected_key]
-    assert data["request_id"] == req_id
-    assert "layer_blocks" in data
-    assert len(data["layer_blocks"]["key_cache"]) == num_layers
-
-    # Verify shape of extracted tensor: [seq_len, heads, dim]
-    # Note: Manager detaches and moves to CPU
-    expected_shape = (seq_len, num_heads, head_dim)
-    assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
-
-
-@pytest.mark.cache
-@hardware_test(
-    res={"cuda": "L4"},
-    num_cards=2,
-)
-def test_manager_reception(kv_config, mock_connector, common_constants):
-    """Test reception and injection logic in OmniKVTransferManager."""
-    num_layers = common_constants["num_layers"]
-    block_size = common_constants["block_size"]
-    num_heads = common_constants["num_heads"]
-    head_dim = common_constants["head_dim"]
-    seq_len = common_constants["seq_len"]
-    req_id = common_constants["req_id"]
-
-    expected_shape = (seq_len, num_heads, head_dim)
-    key_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
-    value_cache = [torch.randn(expected_shape) for _ in range(num_layers)]
-
-    layer_blocks = {"key_cache": key_cache, "value_cache": value_cache}
-    metadata = {
-        "block_size": block_size,
-        "num_layers": num_layers,
-        "dtype": "float32",
-        "seq_len": seq_len,
-    }
-
-    data_to_receive = {
-        "request_id": req_id,
-        "layer_blocks": layer_blocks,
-        "metadata": metadata,
-        "block_ids": [],
-    }
-
-    # In setUp, from_stage="stage1", stage_id="stage2". recv_stages=("stage1", "stage2")
-
-    manager = OmniKVTransferManager(kv_config)
-    manager._connector = mock_connector
-
-    # Pre-populate connector with data
-    # Manager builds full key: omni_{from}_to_{to}_kv_cache_{req_id}
-    full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}"
-    store_key = f"stage1->stage2:{full_request_id}"
-    mock_connector.store[store_key] = data_to_receive
-
-    req = OmniDiffusionRequest(
-        prompts=["test_recv"],
-        sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=[req_id],
-    )
-    # req.need_kv_receive = True # Implicitly handled by receive_kv_cache check? No, manager doesn't check it, runner does.
-    # But receive_kv_cache in manager checks request_id. Which we need to fix in manager next.
-    success = manager.receive_kv_cache(req, target_device=torch.device("cpu"))
-
-    assert success
-    assert hasattr(req, "past_key_values")
-    assert hasattr(req, "kv_metadata")
-
-    assert len(req.past_key_values.key_cache) == num_layers
-    assert torch.allclose(req.past_key_values.key_cache[0], key_cache[0])
-    assert req.kv_metadata["seq_len"] == seq_len
-
-
-@pytest.mark.cache
-@hardware_test(
-    res={"cuda": "L4"},
-    num_cards=2,
-)
-def test_integration_flow(common_constants):
-    """Simulate extraction -> connector -> reception."""
-    num_layers = common_constants["num_layers"]
-    block_size = common_constants["block_size"]
-    num_heads = common_constants["num_heads"]
-    head_dim = common_constants["head_dim"]
-    req_id = common_constants["req_id"]
-
-    sender_config = OmniKVCacheConfig(
-        connector_config={"type": "mock"}, from_stage="sender", to_stage="receiver", need_send_cache=True
-    )
-    sender_manager = OmniKVTransferManager(sender_config)
-    connector = MockConnector()
-    sender_manager._connector = connector  # Shared connector
-
-    # Create Data
-    num_blocks = 5
-    kv_caches = []
-    for _ in range(num_layers):
-        layer = torch.randn(2, num_blocks, block_size, num_heads, head_dim)
-        kv_caches.append(layer)
-
-    finished_reqs = {req_id: {"block_ids": [0, 1], "seq_len": 10}}
-
-    # Send
-    sender_manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32")
-
-    receiver_config = OmniKVCacheConfig(
-        connector_config={"type": "mock"},
-        from_stage="sender",
-        stage_id="receiver",
-        need_recv_cache=True,
-        recv_timeout=1.0,
-    )
-    receiver_manager = OmniKVTransferManager(receiver_config)
-    receiver_manager._connector = connector  # Share the same mock connector instance
-
-    req = OmniDiffusionRequest(
-        prompts=["test_integ"],
-        sampling_params=OmniDiffusionSamplingParams(),
-        request_ids=[req_id],
-    )
-
-    # Receive
-    success = receiver_manager.receive_kv_cache(req)
-
-    # Verify
-    assert success
-    assert req.past_key_values is not None
-    assert req.kv_metadata["seq_len"] == 10
-
-
-@pytest.mark.cache
-@hardware_test(
-    res={"cuda": "L4"},
-    num_cards=2,
-)
-def test_manager_extraction_no_connector(kv_config, common_constants):
-    """Test extraction when connector is unavailable (should still return IDs)."""
-    block_size = common_constants["block_size"]
-    req_id = common_constants["req_id"]
-
-    manager = OmniKVTransferManager(kv_config)
-    # Force connector to be None
-    manager._connector = None
-    manager.config.connector_config = None
-    finished_reqs = {req_id: {"block_ids": [1, 2], "seq_len": 10}}
-
-    processed = manager.handle_finished_requests_kv_transfer(
-        finished_reqs, kv_caches=[], block_size=block_size, cache_dtype="float32"
-    )
-
-    assert req_id in processed
+if __name__ == "__main__":
+    unittest.main()

tests/e2e/offline_inference/conftest.py

@@ -10,7 +10,7 @@ import pytest
 from vllm import TextPrompt
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.conftest import _run_post_test_cleanup, _run_pre_test_cleanup
+from tests.conftest import _run_post_test_cleanup
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -51,9 +51,6 @@ class OmniRunner:
             stage_configs_path: Optional path to YAML stage config file
             **kwargs: Additional arguments passed to Omni
         """
-        cleanup_dist_env_and_memory()
-        _run_pre_test_cleanup(enable_force=True)
-        _run_post_test_cleanup(enable_force=True)
         self.model_name = model_name
         self.seed = seed
 
@@ -340,7 +337,7 @@ class OmniRunner:
         self.close()
         del self.omni
         cleanup_dist_env_and_memory()
-        _run_post_test_cleanup(enable_force=True)
+        _run_post_test_cleanup()
 
     def close(self):
         """Close and cleanup the Omni instance."""

tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml

@@ -1,85 +0,0 @@
-# stage config for running BAGEL with Mooncake connector for CI e2e tests.
-# This config is optimized for single GPU tests with Mooncake inter-stage communication.
-
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: thinker
-      model_arch: BagelForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: text
-      distributed_executor_backend: mp
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      tensor_parallel_size: 1
-      omni_kv_config:
-        need_send_cache: true
-        kv_transfer_criteria:
-          type: prefill_finished
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 52
-      detokenize: true
-      repetition_penalty: 1.05
-    output_connectors:
-      to_stage_1: mooncake_connector
-  - stage_id: 1
-    stage_type: diffusion
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: dit
-      gpu_memory_utilization: 0.55
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: image
-      distributed_executor_backend: mp
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      tensor_parallel_size: 1
-      omni_kv_config:
-        need_recv_cache: true
-    engine_input_source: [0]
-    final_output: true
-    final_output_type: image
-    is_comprehension: false
-    default_sampling_params:
-      seed: 52
-    input_connectors:
-      from_stage_0: mooncake_connector
-
-# Top-level runtime config with Mooncake connector
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1
-  connectors:
-    mooncake_connector:
-      name: MooncakeConnector
-      extra:
-        host: "${MOONCAKE_HOST}"
-        metadata_server: "http://${MOONCAKE_HOST}:${MOONCAKE_HTTP_PORT}/metadata"
-        master: "${MOONCAKE_HOST}:${MOONCAKE_RPC_PORT}"
-        segment: 64000000
-        localbuf: 64000000
-        proto: tcp
-  edges:
-    - from: 0
-      to: 1
-      window_size: -1

tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml

@@ -1,83 +0,0 @@
-# stage config for running BAGEL with SharedMemory connector for CI e2e tests.
-# This config is optimized for single GPU tests with SharedMemory inter-stage communication.
-
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: thinker
-      model_arch: BagelForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.35
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: text
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      tensor_parallel_size: 1
-      omni_kv_config:
-        need_send_cache: true
-        kv_transfer_criteria:
-          type: prefill_finished #or special token generated
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 52
-      detokenize: True
-      repetition_penalty: 1.05
-
-  - stage_id: 1
-    stage_type: diffusion
-    runtime:
-      devices: "0"
-      max_batch_size: 1
-    engine_args:
-      model_stage: dit
-      gpu_memory_utilization: 0.55
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: image
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      tensor_parallel_size: 1
-      omni_kv_config:
-        need_recv_cache: true
-    engine_input_source: [0]
-
-    final_output: true
-    final_output_type: image
-    is_comprehension: false
-    default_sampling_params:
-      seed: 52
-
-# Runtime edges
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1
-
-  # Distributed connectors configuration (optional)
-  # More connectors will be supported in the future.
-  connectors:
-    shared_memory_connector:
-      name: SharedMemoryConnector
-      extra:
-        shm_threshold_bytes: 65536 # 64KB threshold
-
-
-  edges:
-    - from: 0
-      to: 1
-      window_size: -1

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

@@ -10,7 +10,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -40,7 +40,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -70,7 +70,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: generation
+      worker_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true

tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml

@@ -11,7 +11,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -41,7 +41,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -71,7 +71,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: generation
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true

tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml

@@ -7,21 +7,21 @@
 stage_args:
   - stage_id: 0
     runtime:
-      devices: "0"
+      devices: "0,1"
       max_batch_size: 1
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
+      gpu_memory_utilization: 0.6
       enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
       hf_config_name: thinker_config
-      tensor_parallel_size: 1
+      tensor_parallel_size: 2
       load_format: dummy
     final_output: true
     final_output_type: text
@@ -42,9 +42,9 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_type: ar
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.6
+       gpu_memory_utilization: 0.3
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
@@ -68,12 +68,12 @@ stage_args:
 
   - stage_id: 2
     runtime:
-      devices: "1"
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: generation
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true

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

@@ -11,7 +11,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -41,7 +41,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       max_model_len: 896
       max_num_batched_tokens: 896
@@ -71,7 +71,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_type: generation
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true

tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -7,21 +7,21 @@
 stage_args:
   - stage_id: 0
     runtime:
-      devices: "0"
+      devices: "0,1"
       max_batch_size: 1
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
+      gpu_memory_utilization: 0.6
       enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
       hf_config_name: thinker_config
-      tensor_parallel_size: 1
+      tensor_parallel_size: 2
       load_format: dummy
     final_output: true
     final_output_type: text
@@ -42,9 +42,9 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_type: ar
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.6
+       gpu_memory_utilization: 0.3
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
@@ -68,12 +68,12 @@ stage_args:
 
   - stage_id: 2
     runtime:
-      devices: "1"
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: generation
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -1,308 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-End-to-end test for Bagel text2img generation.
-
-This test validates that the Bagel model generates images that match
-expected reference pixel values within a ±5 tolerance.
-
-Equivalent to running:
-    python3 examples/offline_inference/bagel/end2end.py \
-        --prompts "A futuristic city skyline at twilight, cyberpunk style" \
-        --modality text2img --step 15
-"""
-
-import os
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
-
-import signal
-import socket
-import subprocess
-import tempfile
-import time
-from pathlib import Path
-from typing import Any
-
-import pytest
-from PIL import Image
-
-from tests.utils import hardware_test
-from vllm_omni.entrypoints.omni import Omni
-
-# Reference pixel data extracted from the known-good output image
-# Each entry contains (x, y) position and expected (R, G, B) values
-# "Generated with seed=52, num_inference_steps=15,
-# prompt='A futuristic city skyline at twilight, cyberpunk style'"
-REFERENCE_PIXELS = [
-    {"position": (100, 100), "rgb": (68, 107, 134)},
-    {"position": (400, 50), "rgb": (95, 139, 166)},
-    {"position": (700, 100), "rgb": (99, 122, 151)},
-    {"position": (150, 400), "rgb": (111, 125, 153)},
-    {"position": (512, 512), "rgb": (97, 107, 131)},
-    {"position": (700, 400), "rgb": (48, 64, 98)},
-    {"position": (100, 700), "rgb": (79, 63, 84)},
-    {"position": (400, 700), "rgb": (40, 58, 79)},
-    {"position": (700, 700), "rgb": (60, 75, 103)},
-    {"position": (256, 256), "rgb": (97, 128, 156)},
-]
-
-# Maximum allowed difference per color channel
-PIXEL_TOLERANCE = 5
-
-# Default test prompt
-DEFAULT_PROMPT = "<|im_start|>A futuristic city skyline at twilight, cyberpunk style<|im_end|>"
-
-
-def _find_free_port() -> int:
-    """Find and return a free ephemeral port by binding to port 0."""
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        s.listen(1)
-        port = s.getsockname()[1]
-    return port
-
-
-def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_steps: int = 15) -> list:
-    """Configure sampling parameters for Bagel text2img generation.
-
-    Args:
-        omni: The Omni instance to get default params from.
-        max_tokens: Maximum tokens for the first stage.
-        num_inference_steps: Number of inference steps for the diffusion stage.
-
-    Returns:
-        Configured sampling params list.
-    """
-    params_list = omni.default_sampling_params_list
-    params_list[0].max_tokens = max_tokens  # type: ignore
-    if len(params_list) > 1:
-        params_list[1].num_inference_steps = num_inference_steps  # type: ignore
-    return params_list
-
-
-def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
-    """Extract the generated image from Omni outputs.
-
-    Args:
-        omni_outputs: List of outputs from omni.generate().
-
-    Returns:
-        The first generated PIL Image, or None if no image found.
-    """
-    for req_output in omni_outputs:
-        if images := getattr(req_output, "images", None):
-            return images[0]
-        if hasattr(req_output, "request_output") and req_output.request_output:
-            for stage_out in req_output.request_output:
-                if hasattr(stage_out, "images") and stage_out.images:
-                    return stage_out.images[0]
-    return None
-
-
-def _validate_pixels(
-    image: Image.Image,
-    reference_pixels: list[dict[str, Any]] = REFERENCE_PIXELS,
-    tolerance: int = PIXEL_TOLERANCE,
-) -> None:
-    """Validate that image pixels match expected reference values.
-
-    Args:
-        image: The PIL Image to validate.
-        reference_pixels: List of dicts with 'position' (x, y) and 'rgb' (R, G, B).
-        tolerance: Maximum allowed difference per color channel.
-
-    Raises:
-        AssertionError: If any pixel differs beyond tolerance.
-    """
-    for ref in reference_pixels:
-        x, y = ref["position"]
-        expected = ref["rgb"]
-        actual = image.getpixel((x, y))[:3]
-        assert all(abs(a - e) <= tolerance for a, e in zip(actual, expected)), (
-            f"Pixel mismatch at ({x}, {y}): expected {expected}, got {actual}"
-        )
-
-
-def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Image:
-    """Generate an image using Bagel model with configured parameters.
-
-    Args:
-        omni: The Omni instance to use for generation.
-        prompt: The text prompt for image generation.
-
-    Returns:
-        The generated PIL Image.
-
-    Raises:
-        AssertionError: If no image is generated or size is incorrect.
-    """
-    params_list = _configure_sampling_params(omni)
-
-    omni_outputs = list(
-        omni.generate(
-            prompts=[{"prompt": prompt, "modalities": ["image"]}],
-            sampling_params_list=params_list,
-        )
-    )
-
-    generated_image = _extract_generated_image(omni_outputs)
-    assert generated_image is not None, "No images generated"
-    assert generated_image.size == (1024, 1024), f"Expected 1024x1024, got {generated_image.size}"
-
-    return generated_image
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "H100"})
-def test_bagel_text2img_shared_memory_connector():
-    """Test Bagel text2img with shared memory connector."""
-    config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
-    omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300)
-
-    try:
-        generated_image = _generate_bagel_image(omni)
-        _validate_pixels(generated_image)
-    finally:
-        omni.close()
-
-
-def _wait_for_port(host: str, port: int, timeout: int = 30) -> bool:
-    """Wait for a port to become available.
-
-    Args:
-        host: The host address.
-        port: The port number.
-        timeout: Maximum seconds to wait.
-
-    Returns:
-        True if port becomes available, False otherwise.
-    """
-    for _ in range(timeout):
-        try:
-            with socket.create_connection((host, port), timeout=1):
-                return True
-        except (TimeoutError, ConnectionRefusedError):
-            time.sleep(1)
-    return False
-
-
-def _cleanup_mooncake_processes(timeout_secs: int = 5) -> None:
-    """Clean up any existing mooncake_master processes.
-
-    Args:
-        timeout_secs: Maximum seconds to wait for graceful termination.
-    """
-    subprocess.run(
-        ["pkill", "-f", "mooncake_master"],
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.DEVNULL,
-    )
-    start_time = time.time()
-    while time.time() - start_time < timeout_secs:
-        result = subprocess.run(
-            ["pgrep", "-f", "mooncake_master"],
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.DEVNULL,
-        )
-        if result.returncode != 0:
-            break
-        time.sleep(0.5)
-    else:
-        subprocess.run(
-            ["pkill", "-9", "-f", "mooncake_master"],
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.DEVNULL,
-        )
-    time.sleep(1)
-
-
-def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
-    """Load Mooncake config from YAML and substitute placeholders.
-
-    Args:
-        host: Mooncake host address.
-        rpc_port: RPC port for Mooncake master.
-        http_port: HTTP metadata server port.
-
-    Returns:
-        Path to the temporary config file with substituted values.
-    """
-    config_path = str(Path(__file__).parent / "stage_configs" / "bagel_mooncake_ci.yaml")
-    with open(config_path) as f:
-        config_content = f.read()
-
-    # Substitute placeholders
-    config_content = config_content.replace("${MOONCAKE_HOST}", host)
-    config_content = config_content.replace("${MOONCAKE_RPC_PORT}", str(rpc_port))
-    config_content = config_content.replace("${MOONCAKE_HTTP_PORT}", str(http_port))
-
-    # Write to temp file
-    temp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
-    temp_file.write(config_content)
-    temp_file.close()
-    return temp_file.name
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "H100"})
-def test_bagel_text2img_mooncake_connector():
-    """Test Bagel text2img with Mooncake connector for inter-stage communication."""
-    MOONCAKE_HOST = "127.0.0.1"
-    MOONCAKE_RPC_PORT = _find_free_port()
-    MOONCAKE_HTTP_PORT = _find_free_port()
-    MOONCAKE_METRICS_PORT = _find_free_port()
-
-    mooncake_master_proc = None
-    temp_config_file = None
-    omni = None
-
-    try:
-        _cleanup_mooncake_processes()
-
-        # Start mooncake_master
-        mooncake_master_proc = subprocess.Popen(
-            [
-                "mooncake_master",
-                f"--rpc_port={MOONCAKE_RPC_PORT}",
-                "--enable_http_metadata_server=true",
-                "--http_metadata_server_host=0.0.0.0",
-                f"--http_metadata_server_port={MOONCAKE_HTTP_PORT}",
-                f"--metrics_port={MOONCAKE_METRICS_PORT}",
-            ],
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.DEVNULL,
-            preexec_fn=os.setsid,
-        )
-
-        assert _wait_for_port(MOONCAKE_HOST, MOONCAKE_RPC_PORT), "mooncake_master failed to start"
-
-        # Create temp config and initialize Omni
-        temp_config_file = _load_mooncake_config(
-            host=MOONCAKE_HOST,
-            rpc_port=MOONCAKE_RPC_PORT,
-            http_port=MOONCAKE_HTTP_PORT,
-        )
-
-        omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file, stage_init_timeout=300)
-
-        generated_image = _generate_bagel_image(omni)
-        _validate_pixels(generated_image)
-
-    finally:
-        if omni:
-            omni.close()
-        if temp_config_file:
-            try:
-                os.unlink(temp_config_file)
-            except OSError:
-                pass
-        if mooncake_master_proc:
-            try:
-                os.killpg(os.getpgid(mooncake_master_proc.pid), signal.SIGKILL)
-            except OSError:
-                pass

tests/e2e/offline_inference/test_diffusion_cpu_offload.py

@@ -7,7 +7,7 @@ from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
 from tests.utils import GPUMemoryMonitor
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.platform_utils import is_npu, is_rocm
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -20,7 +20,7 @@ models = ["riverclouds/qwen_image_random"]
 
 
 def inference(model_name: str, offload: bool = True):
-    current_omni_platform.empty_cache()
+    torch.cuda.empty_cache()
     device_index = torch.cuda.current_device()
     monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
@@ -45,7 +45,7 @@ def inference(model_name: str, offload: bool = True):
     return peak
 
 
-@pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported")
+@pytest.mark.skipif(is_npu() or is_rocm(), reason="Hardware not supported")
 @pytest.mark.parametrize("model_name", models)
 def test_cpu_offload_diffusion_model(model_name: str):
     try:

tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

@@ -1,110 +0,0 @@
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-
-from tests.utils import GPUMemoryMonitor
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.platforms import current_omni_platform
-
-# ruff: noqa: E402
-REPO_ROOT = Path(__file__).resolve().parents[2]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-from vllm_omni import Omni
-
-# Models to test and expected saved memory in MB, correspondingly
-MODELS_SAVED_MEMORY_MB = {"riverclouds/qwen_image_random": 4500}
-
-
-def run_inference(
-    model_name: str,
-    layerwise_offload: bool = False,
-    num_gpu_layers: int = 1,
-    num_inference_steps: int = 3,
-) -> float:
-    # For now, only support on GPU, so apply torch.cuda operations here
-    # NPU / ROCm platforms are expected to be detected and skipped this test function
-    torch.cuda.empty_cache()
-    device_index = torch.cuda.current_device()
-    monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02)
-    monitor.start()
-
-    m = Omni(
-        model=model_name,
-        enable_layerwise_offload=layerwise_offload,
-        layerwise_num_gpu_layers=num_gpu_layers,
-        boundary_ratio=0.875,
-        flow_shift=5.0,
-    )
-
-    torch.cuda.reset_peak_memory_stats(device=device_index)
-
-    # Refer to tests/e2e/offline_inference/test_t2v_model.py
-    # Use minimal settings for testing
-    height = 480
-    width = 640
-    num_frames = 5
-
-    m.generate(
-        "A cat sitting on a table",
-        OmniDiffusionSamplingParams(
-            height=height,
-            width=width,
-            generator=torch.Generator("cuda").manual_seed(42),
-            guidance_scale=1.0,
-            num_inference_steps=num_inference_steps,
-            num_frames=num_frames,
-        ),
-    )
-
-    peak = monitor.peak_used_mb
-    monitor.stop()
-
-    return peak
-
-
-@pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported")
-@pytest.mark.parametrize("model_name", MODELS_SAVED_MEMORY_MB.keys())
-def test_layerwise_offload_diffusion_model(model_name: str):
-    """Test that layerwise offloading reduces GPU memory usage.
-
-    This test verifies that layerwise offloading significantly reduces peak
-    GPU memory usage compared to loading the entire model on GPU. The layerwise
-    offloader keeps only a single transformer block on GPU at a time, with
-    prefetching for compute-memory overlap.
-    """
-    try:
-        # Run without layerwise offloading (baseline)
-        no_offload_peak_memory = run_inference(model_name, layerwise_offload=False)
-        cleanup_dist_env_and_memory()
-
-        # Run with layerwise offloading (1 layer on device)
-        layerwise_offload_peak_memory = run_inference(model_name, layerwise_offload=True, num_gpu_layers=1)
-        cleanup_dist_env_and_memory()
-
-        # Run with 2 layers on device
-        layerwise_offload_two_layers_peak = run_inference(model_name, layerwise_offload=True, num_gpu_layers=2)
-    except Exception:
-        pytest.fail("Inference failed")
-
-    print(f"Layerwise offload peak memory (1 GPU layer): {layerwise_offload_peak_memory} MB")
-    print(f"Layerwise offload peak memory (2 GPU layers): {layerwise_offload_two_layers_peak} MB")
-    print(f"No offload peak memory: {no_offload_peak_memory} MB")
-
-    # Verify that layerwise offloading significantly reduces memory usage
-    # Passes only if the actual savings exceeds the expected savings
-    assert layerwise_offload_peak_memory + MODELS_SAVED_MEMORY_MB[model_name] < no_offload_peak_memory, (
-        f"Layerwise offload peak memory {layerwise_offload_peak_memory} MB "
-        f"should be significantly less than no offload peak memory {no_offload_peak_memory} MB"
-    )
-
-    # Verify that 2 GPU layers uses more memory than 1 GPU layer
-    # But not excessively more (should be a reasonable increase)
-    assert layerwise_offload_peak_memory < layerwise_offload_two_layers_peak, (
-        f"1 GPU layer peak {layerwise_offload_peak_memory} MB should be < "
-        f"2 GPU layers peak {layerwise_offload_two_layers_peak} MB"
-    )

tests/e2e/offline_inference/test_diffusion_lora.py

@@ -9,6 +9,7 @@ from safetensors.torch import save_file
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.utils.platform_utils import is_npu
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -24,6 +25,10 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 # model to reduce runtime and avoid extra downloads.
 models = ["Tongyi-MAI/Z-Image-Turbo"]
 
+# NPU still can't run Tongyi-MAI/Z-Image-Turbo properly.
+if is_npu():
+    pytest.skip("Tongyi-MAI/Z-Image-Turbo is not supported on NPU yet.", allow_module_level=True)
+
 
 @pytest.mark.parametrize("model_name", models)
 def test_diffusion_model(model_name: str, tmp_path: Path):

tests/e2e/offline_inference/test_qwen2_5_omni.py

@@ -13,7 +13,7 @@ from vllm.assets.video import VideoAsset
 from vllm.envs import VLLM_USE_MODELSCOPE
 from vllm.multimodal.image import convert_image_mode
 
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils import is_npu, is_rocm
 
 from .conftest import OmniRunner
 from .utils import create_new_process_for_each_test
@@ -21,9 +21,9 @@ from .utils import create_new_process_for_each_test
 models = ["Qwen/Qwen2.5-Omni-3B"]
 
 # CI stage config optimized for 24GB GPU (L4/RTX3090) or NPU
-if current_omni_platform.is_npu():
+if is_npu():
     stage_config = str(Path(__file__).parent / "stage_configs" / "npu" / "qwen2_5_omni_ci.yaml")
-elif current_omni_platform.is_rocm():
+elif is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_config = str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
 else:
@@ -88,7 +88,7 @@ def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: t
         assert len(audio_output.request_output) > 0
 
         # Verify audio tensor exists and has content
-        audio_tensor = audio_output.request_output[0].outputs[0].multimodal_output["audio"]
+        audio_tensor = audio_output.request_output[0].multimodal_output["audio"]
         assert audio_tensor is not None
         assert audio_tensor.numel() > 0
 

tests/e2e/offline_inference/test_qwen3_omni.py

@@ -14,14 +14,14 @@ from pathlib import Path
 import pytest
 from vllm.assets.video import VideoAsset
 
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils import is_rocm
 
 from .conftest import OmniRunner
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325
-if current_omni_platform.is_rocm():
+if is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
 else:
@@ -75,6 +75,6 @@ def test_video_to_audio(omni_runner: type[OmniRunner], test_config) -> None:
         assert len(audio_output.request_output) > 0
 
         # Verify audio tensor exists and has content
-        audio_tensor = audio_output.request_output[0].outputs[0].multimodal_output["audio"]
+        audio_tensor = audio_output.request_output[0].multimodal_output["audio"]
         assert audio_tensor is not None
         assert audio_tensor.numel() > 0

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -2,9 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-System test for Sequence Parallel (SP) backends: Ulysses and Ring attention.
+System test for Ulysses sequence parallel backend.
 
-Tests verify that SP inference produces correct outputs compared to baseline.
+This test verifies that Ulysses-SP (DeepSpeed Ulysses Sequence Parallel) works
+correctly with diffusion models. It uses minimal settings to keep test time
+short for CI.
 """
 
 import gc
@@ -12,7 +14,6 @@ import os
 import sys
 import time
 from pathlib import Path
-from typing import NamedTuple
 
 import numpy as np
 import pytest
@@ -29,24 +30,27 @@ if str(REPO_ROOT) not in sys.path:
 
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.diffusion.distributed.parallel_state import device_count
+from vllm_omni.diffusion.envs import get_device_name
+
+models = ["riverclouds/qwen_image_random"]
 
-# Test configuration
-MODELS = ["riverclouds/qwen_image_random"]
 PROMPT = "a photo of a cat sitting on a laptop keyboard"
-DEFAULT_HEIGHT = 256
-DEFAULT_WIDTH = 256
-DEFAULT_SEED = 42
-DEFAULT_STEPS = 4
-DIFF_MEAN_THRESHOLD = 2e-2
-DIFF_MAX_THRESHOLD = 2e-1
 
 
-class InferenceResult(NamedTuple):
-    """Result of an inference run."""
+def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
+    """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3]."""
+    arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
+    return torch.from_numpy(arr)
 
-    images: list[Image.Image]
-    elapsed_ms: float
+
+def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
+    """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1]."""
+    ta = _pil_to_float_rgb_tensor(a)
+    tb = _pil_to_float_rgb_tensor(b)
+    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
+    abs_diff = torch.abs(ta - tb)
+    return abs_diff.mean().item(), abs_diff.max().item()
 
 
 def _cleanup_distributed():
@@ -58,223 +62,205 @@ def _cleanup_distributed():
         os.environ.pop(key, None)
 
     gc.collect()
-    if current_omni_platform.is_available():
-        current_omni_platform.empty_cache()
-        current_omni_platform.synchronize()
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
 
     time.sleep(5)
 
 
-def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
-    """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1]."""
-    ta = torch.from_numpy(np.asarray(a.convert("RGB"), dtype=np.float32) / 255.0)
-    tb = torch.from_numpy(np.asarray(b.convert("RGB"), dtype=np.float32) / 255.0)
-    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
-    abs_diff = torch.abs(ta - tb)
-    return abs_diff.mean().item(), abs_diff.max().item()
-
-
-def _run_inference(
-    model_name: str,
-    dtype: torch.dtype,
-    attn_backend: str,
-    ulysses_degree: int = 1,
-    ring_degree: int = 1,
-    height: int = DEFAULT_HEIGHT,
-    width: int = DEFAULT_WIDTH,
-    seed: int = DEFAULT_SEED,
-    warmup: bool = True,
-) -> InferenceResult:
-    """Run inference with specified configuration.
-
-    Args:
-        warmup: If True, run one warmup iteration before the timed run.
-    """
-    parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
-    omni = Omni(
+def _run_baseline(model_name: str, dtype: torch.dtype, attn_backend: str, height: int, width: int, seed: int):
+    """Run baseline inference (no SP)."""
+    baseline_parallel_config = DiffusionParallelConfig(ulysses_degree=1, ring_degree=1)
+    baseline = Omni(
         model=model_name,
-        parallel_config=parallel_config,
+        parallel_config=baseline_parallel_config,
         dtype=dtype,
         attention_backend=attn_backend,
     )
 
     try:
-        # Warmup run (not timed)
-        if warmup:
-            _ = omni.generate(
-                PROMPT,
-                OmniDiffusionSamplingParams(
-                    height=height,
-                    width=width,
-                    num_inference_steps=DEFAULT_STEPS,
-                    guidance_scale=0.0,
-                    generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed + 1000),
-                    num_outputs_per_prompt=1,
-                ),
-            )
-
-        # Timed run
-        start = time.time()
-        outputs = omni.generate(
+        outputs = baseline.generate(
             PROMPT,
             OmniDiffusionSamplingParams(
                 height=height,
                 width=width,
-                num_inference_steps=DEFAULT_STEPS,
+                num_inference_steps=4,
                 guidance_scale=0.0,
-                generator=torch.Generator(current_omni_platform.device_type).manual_seed(seed),
+                generator=torch.Generator(get_device_name()).manual_seed(seed),
                 num_outputs_per_prompt=1,
             ),
         )
-        elapsed_ms = (time.time() - start) * 1000
-
-        return InferenceResult(
-            images=outputs[0].request_output[0].images,
-            elapsed_ms=elapsed_ms,
-        )
+        return outputs[0].request_output[0].images
     finally:
-        omni.close()
+        baseline.close()
         _cleanup_distributed()
 
 
-# =============================================================================
-# Correctness & Performance Tests
-# =============================================================================
+def _run_sp(
+    model_name: str,
+    dtype: torch.dtype,
+    attn_backend: str,
+    height: int,
+    width: int,
+    seed: int,
+    ulysses_degree: int,
+    ring_degree: int,
+):
+    """Run SP inference."""
+    sp_parallel_config = DiffusionParallelConfig(ulysses_degree=ulysses_degree, ring_degree=ring_degree)
+    sp = Omni(
+        model=model_name,
+        parallel_config=sp_parallel_config,
+        dtype=dtype,
+        attention_backend=attn_backend,
+    )
 
-# SP configurations: (ulysses_degree, ring_degree, height, width, warmup, is_perf_test)
-# - warmup: whether to run warmup for this SP config
-# - is_perf_test: whether this is a performance test (show speedup metrics)
-SP_CONFIGS = [
-    (2, 1, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True),  # Ulysses-2 - performance test
-    (1, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True),  # Ring-2 - performance test
-    (2, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, False, False),  # Hybrid - correctness only
-    (4, 1, 272, 272, False, False),  # Ulysses-4 - shape and correctness
-]
-
-
-def _get_sp_mode(ulysses_degree: int, ring_degree: int) -> str:
-    """Get SP mode name for logging."""
-    if ulysses_degree > 1 and ring_degree == 1:
-        return f"ulysses-{ulysses_degree}"
-    elif ring_degree > 1 and ulysses_degree == 1:
-        return f"ring-{ring_degree}"
-    else:
-        return f"hybrid-{ulysses_degree}x{ring_degree}"
-
-
-@pytest.mark.parametrize("model_name", MODELS)
-def test_sp_correctness(model_name: str):
-    """Test that SP inference produces correct outputs and measure performance.
-
-    Runs baseline once per unique (height, width), then tests all SP configs.
-
-    Note: Run with `pytest -v -s` to see detailed output.
-    """
-    device_count = current_omni_platform.get_device_count()
-
-    # Cache baseline results by (height, width)
-    # Key: (height, width), Value: (result, warmup_used)
-    baseline_cache: dict[tuple[int, int], InferenceResult] = {}
-
-    # Collect results for summary
-    results: list[dict] = []
-
-    print("\n" + "=" * 70)
-    print(f"Sequence Parallel Test - Model: {model_name}")
-    print(f"Available GPUs: {device_count}")
-    print("=" * 70)
-
-    for ulysses_degree, ring_degree, height, width, sp_warmup, is_perf_test in SP_CONFIGS:
-        sp_size = ulysses_degree * ring_degree
-        sp_mode = _get_sp_mode(ulysses_degree, ring_degree)
-
-        if device_count < sp_size:
-            print(f"\n[{sp_mode}] SKIPPED (requires {sp_size} GPUs)")
-            continue
-
-        # Determine baseline warmup: only for default size (performance tests)
-        cache_key = (height, width)
-        baseline_warmup = height == DEFAULT_HEIGHT and width == DEFAULT_WIDTH
-
-        # Get or compute baseline for this (height, width)
-        if cache_key not in baseline_cache:
-            print(f"\n--- Running baseline {height}x{width} (warmup={baseline_warmup}) ---")
-            baseline = _run_inference(
-                model_name,
-                torch.bfloat16,
-                "sdpa",
+    try:
+        outputs = sp.generate(
+            PROMPT,
+            OmniDiffusionSamplingParams(
                 height=height,
                 width=width,
-                warmup=baseline_warmup,
-            )
-            assert len(baseline.images) == 1
-            baseline_cache[cache_key] = baseline
-            print(f"[baseline] {height}x{width}: {baseline.elapsed_ms:.0f}ms")
-        else:
-            baseline = baseline_cache[cache_key]
-
-        # Run SP
-        print(f"\n--- Running {sp_mode} (warmup={sp_warmup}) ---")
-        sp_result = _run_inference(
-            model_name,
-            torch.bfloat16,
-            "sdpa",
-            ulysses_degree=ulysses_degree,
-            ring_degree=ring_degree,
-            height=height,
-            width=width,
-            warmup=sp_warmup,
+                num_inference_steps=4,
+                guidance_scale=0.0,
+                generator=torch.Generator(get_device_name()).manual_seed(seed),
+                num_outputs_per_prompt=1,
+            ),
         )
-        assert len(sp_result.images) == 1
+        return outputs[0].request_output[0].images
+    finally:
+        sp.close()
+        _cleanup_distributed()
 
-        # Compare outputs (correctness)
-        mean_diff, max_diff = _diff_metrics(baseline.images[0], sp_result.images[0])
 
-        # Build result entry
-        result = {
-            "mode": sp_mode,
-            "sp_size": sp_size,
-            "height": height,
-            "width": width,
-            "baseline_ms": baseline.elapsed_ms,
-            "sp_ms": sp_result.elapsed_ms,
-            "mean_diff": mean_diff,
-            "max_diff": max_diff,
-            "is_perf_test": is_perf_test,
-        }
-        results.append(result)
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_baseline_only(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test baseline inference only (no SP)."""
+    height = 256
+    width = 256
+    seed = 42
 
-        # Output based on test type
-        if is_perf_test:
-            speedup = baseline.elapsed_ms / sp_result.elapsed_ms if sp_result.elapsed_ms > 0 else 0
-            result["speedup"] = speedup
-            print(
-                f"[{sp_mode}] {sp_size} GPUs | "
-                f"baseline: {baseline.elapsed_ms:.0f}ms, sp: {sp_result.elapsed_ms:.0f}ms, "
-                f"speedup: {speedup:.2f}x"
-            )
-        else:
-            print(f"[{sp_mode}] {sp_size} GPUs | sp: {sp_result.elapsed_ms:.0f}ms (correctness only)")
+    images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
 
-        print(f"[{sp_mode}] diff: mean={mean_diff:.6e}, max={max_diff:.6e}")
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height
 
-        # Assert correctness
-        assert mean_diff <= DIFF_MEAN_THRESHOLD and max_diff <= DIFF_MAX_THRESHOLD, (
-            f"[{sp_mode}] SP output differs from baseline: mean={mean_diff:.6e}, max={max_diff:.6e}"
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sp_ulysses2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test SP inference only (ulysses=2)."""
+    if device_count() < 2:
+        pytest.skip(f"Test requires 2 GPUs but only {device_count()} available")
+
+    height = 256
+    width = 256
+    seed = 42
+
+    images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=2, ring_degree=1)
+
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sp_ring2_only(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test SP inference only (ring=2)."""
+    if device_count() < 2:
+        pytest.skip(f"Test requires 2 GPUs but only {device_count()} available")
+
+    height = 256
+    width = 256
+    seed = 42
+
+    images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree=1, ring_degree=2)
+
+    assert images is not None
+    assert len(images) == 1
+    assert images[0].width == width
+    assert images[0].height == height
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("ulysses_degree", [1, 2])
+@pytest.mark.parametrize("ring_degree", [1, 2])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sequence_parallel(
+    model_name: str,
+    ulysses_degree: int,
+    ring_degree: int,
+    dtype: torch.dtype,
+    attn_backend: str,
+):
+    """Compare baseline (ulysses_degree=1) vs SP (ulysses_degree>1) outputs."""
+    if ulysses_degree <= 1 and ring_degree <= 1:
+        pytest.skip(
+            "This test compares ulysses_degree * ring_degree = 1 vs ulysses_degree * ring_degree > 1; "
+            "provide ulysses_degree or ring_degree>1."
         )
 
-    # Summary
-    print("\n" + "=" * 70)
-    print("SUMMARY")
-    print("=" * 70)
-    print(f"{'Mode':<15} {'GPUs':<6} {'Size':<10} {'Baseline':<12} {'SP':<12} {'Speedup':<10} {'Status'}")
-    print("-" * 70)
-    for r in results:
-        speedup_str = f"{r['speedup']:.2f}x" if r.get("speedup") else "N/A"
-        baseline_str = f"{r['baseline_ms']:.0f}ms" if r["is_perf_test"] else "N/A"
-        status = "PASS" if r["mean_diff"] <= DIFF_MEAN_THRESHOLD else "FAIL"
-        print(
-            f"{r['mode']:<15} {r['sp_size']:<6} {r['height']}x{r['width']:<5} "
-            f"{baseline_str:<12} {r['sp_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}"
-        )
-    print("=" * 70)
+    sp_size = ulysses_degree * ring_degree
+    if device_count() < sp_size:
+        pytest.skip(f"Test requires {sp_size} GPUs but only {device_count()} available")
+
+    height = 256
+    width = 256
+    seed = 42
+
+    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+    assert baseline_images is not None and len(baseline_images) == 1
+
+    sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree)
+    assert sp_images is not None and len(sp_images) == 1
+
+    mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
+
+    mean_threshold = 2e-2
+    max_threshold = 2e-1
+
+    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
+        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
+    )
+
+
+@pytest.mark.parametrize("model_name", models)
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("attn_backend", ["sdpa"])
+def test_sequence_parallel_ulysses4(model_name: str, dtype: torch.dtype, attn_backend: str):
+    """Test SP with ulysses_degree=4."""
+    ulysses_degree = 4
+    ring_degree = 1
+
+    if device_count() < ulysses_degree * ring_degree:
+        pytest.skip(f"Test requires {ulysses_degree * ring_degree} GPUs but only {device_count()} available")
+
+    height = 272
+    width = 272
+    seed = 42
+
+    baseline_images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+    assert baseline_images is not None and len(baseline_images) == 1
+
+    sp_images = _run_sp(model_name, dtype, attn_backend, height, width, seed, ulysses_degree, ring_degree)
+    assert sp_images is not None and len(sp_images) == 1
+
+    mean_abs_diff, max_abs_diff = _diff_metrics(baseline_images[0], sp_images[0])
+
+    mean_threshold = 2e-2
+    max_threshold = 2e-1
+
+    assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, (
+        f"Image diff exceeded threshold: mean={mean_abs_diff:.6e}, max={max_abs_diff:.6e}"
+    )

tests/e2e/offline_inference/test_t2i_model.py

@@ -7,7 +7,7 @@ import torch
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.platform_utils import is_npu, is_rocm
 
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -21,11 +21,12 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 
 models = ["Tongyi-MAI/Z-Image-Turbo", "riverclouds/qwen_image_random"]
 
+# NPU still can't run Tongyi-MAI/Z-Image-Turbo properly
 # Modelscope can't find riverclouds/qwen_image_random
 # TODO: When NPU support is ready, remove this branch.
-if current_omni_platform.is_npu():
-    models = ["Tongyi-MAI/Z-Image-Turbo", "Qwen/Qwen-Image"]
-elif current_omni_platform.is_rocm():
+if is_npu():
+    models = ["Qwen/Qwen-Image"]
+elif is_rocm():
     # TODO: When ROCm support is ready, remove this branch.
     # vLLM V0.11.0 has issues running riverclouds/qwen_image_random
     # on ROCm

tests/e2e/offline_inference/test_zimage_tensor_parallel.py

@@ -23,7 +23,7 @@ from tests.utils import GPUMemoryMonitor
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.platform_utils import is_npu, is_rocm
 
 # os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -129,7 +129,7 @@ def _run_zimage_generate(
 
 @pytest.mark.integration
 def test_zimage_tensor_parallel_tp2(tmp_path: Path):
-    if current_omni_platform.is_npu() or current_omni_platform.is_rocm():
+    if is_npu() or is_rocm():
         pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.")
     if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
         pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.")

tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml

@@ -6,24 +6,22 @@
 # The following config has been verified on 2x H100-80G GPUs.
 stage_args:
   - stage_id: 0
-    stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
-      devices: "0"
+      devices: "0,1"
       max_batch_size: 5
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
+      gpu_memory_utilization: 0.6
       enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
-      max_num_batched_tokens: 32768
       hf_config_name: thinker_config
-      tensor_parallel_size: 1
+      tensor_parallel_size: 2
       load_format: dummy
     final_output: true
     final_output_type: text
@@ -38,16 +36,15 @@ stage_args:
       repetition_penalty: 1.05
 
   - stage_id: 1
-    stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
        devices: "1"
        max_batch_size: 5
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_type: ar
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.6
+       gpu_memory_utilization: 0.3
        enforce_eager: false
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
@@ -70,14 +67,13 @@ stage_args:
       stop_token_ids: [2150]
 
   - stage_id: 2
-    stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
-      devices: "1"
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: generation
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true

tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml

@@ -7,7 +7,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.6
       enforce_eager: true

tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml

@@ -2,25 +2,27 @@
 # Stage 0: Thinker (multimodal understanding + text generation)
 # Stage 1: Talker (text embeddings → 16-layer RVQ codec codes)
 # Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
+
 # The following config has been verified on 2x H100-80G GPUs.
 stage_args:
   - stage_id: 0
     runtime:
-      devices: "0"
+      devices: "0,1"
       max_batch_size: 5
     engine_args:
       model_stage: thinker
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
+      gpu_memory_utilization: 0.6
       enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent  # Output hidden states for talker
       distributed_executor_backend: "mp"
       enable_prefix_caching: false
       hf_config_name: thinker_config
-      tensor_parallel_size: 1
+      tensor_parallel_size: 2
+      load_format: dummy
     final_output: true
     final_output_type: text
     is_comprehension: true
@@ -40,9 +42,9 @@ stage_args:
     engine_args:
        model_stage: talker
        model_arch: Qwen3OmniMoeForConditionalGeneration
-       worker_type: ar
+       worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
        scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.6
+       gpu_memory_utilization: 0.3
        enforce_eager: true
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
@@ -50,6 +52,7 @@ stage_args:
        enable_prefix_caching: false
        distributed_executor_backend: "mp"
        hf_config_name: talker_config
+       load_format: dummy
     engine_input_source: [0]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
     # final_output: true
@@ -57,7 +60,7 @@ stage_args:
     default_sampling_params:
       temperature: 0.9
       top_k: 50
-      max_tokens: 1000
+      max_tokens: 100
       seed: 42
       detokenize: False
       repetition_penalty: 1.05
@@ -65,12 +68,12 @@ stage_args:
 
   - stage_id: 2
     runtime:
-      devices: "1"
+      devices: "0"
       max_batch_size: 1
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: generation
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
@@ -80,6 +83,7 @@ stage_args:
       distributed_executor_backend: "mp"
       max_num_batched_tokens: 1000000
       hf_config_name: thinker_config
+      load_format: dummy
       async_scheduling: false
     engine_input_source: [1]
     custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
@@ -89,7 +93,7 @@ stage_args:
       temperature: 0.0
       top_p: 1.0
       top_k: -1
-      max_tokens: 2000
+      max_tokens: 200
       seed: 42
       detokenize: True
       repetition_penalty: 1.1

tests/e2e/online_serving/test_i2i_multi_image_input.py

@@ -10,25 +10,19 @@ import signal
 import socket
 import subprocess
 import sys
-import threading
 import time
 from io import BytesIO
-from typing import Any
 
 import openai
 import pytest
-import requests
 from PIL import Image
 from vllm.assets.image import ImageAsset
 from vllm.utils.network_utils import get_open_port
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-# Increase timeout for downloading assets from S3 (default 5s is too short for CI)
-os.environ.setdefault("VLLM_IMAGE_FETCH_TIMEOUT", "60")
 
 models = ["Qwen/Qwen-Image-Edit-2509"]
 test_params = models
-t2i_models = ["Tongyi-MAI/Z-Image-Turbo"]
 
 
 class OmniServer:
@@ -180,94 +174,35 @@ def _decode_data_url_to_image_bytes(data_url: str) -> bytes:
 
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_i2i_multi_image_input_qwen_image_edit_2509(
+    client: openai.OpenAI,
     omni_server,
     base64_encoded_images: list[str],
 ) -> None:
-    """Test multi-image input editing via OpenAI API with concurrent requests."""
+    """Test multi-image input editing via OpenAI API."""
     image_data_urls = [f"data:image/png;base64,{img}" for img in base64_encoded_images]
     messages = dummy_messages_from_image_data(image_data_urls)
 
-    barrier = threading.Barrier(2)
-    results: list[tuple[int, int]] = []
-
-    def _call_chat(width: int, height: int) -> None:
-        client = openai.OpenAI(
-            base_url=f"http://{omni_server.host}:{omni_server.port}/v1",
-            api_key="EMPTY",
-        )
-        barrier.wait()
-        chat_completion = client.chat.completions.create(
-            model=omni_server.model,
-            messages=messages,
-            extra_body={
-                "height": height,
-                "width": width,
-                "num_inference_steps": 2,
-                "guidance_scale": 0.0,
-                "seed": 42,
-            },
-        )
-
-        assert len(chat_completion.choices) == 1
-        choice = chat_completion.choices[0]
-        assert choice.finish_reason == "stop"
-        assert choice.message.role == "assistant"
-
-        image_data_url = _extract_image_data_url(choice.message.content)
-        image_bytes = _decode_data_url_to_image_bytes(image_data_url)
-        img = Image.open(BytesIO(image_bytes))
-        img.load()
-        results.append(img.size)
-
-    threads = [
-        threading.Thread(target=_call_chat, args=(1248, 832)),
-        threading.Thread(target=_call_chat, args=(1024, 768)),
-    ]
-    for t in threads:
-        t.start()
-    for t in threads:
-        t.join()
-
-    # TODO @ZJY
-    # assert (1248, 832) in results
-    # assert (1024, 768) in results
-
-
-@pytest.mark.parametrize("omni_server", t2i_models, indirect=True)
-def test_t2i_concurrent_requests_different_sizes(omni_server) -> None:
-    """Test /v1/images/generations concurrent requests with different sizes."""
-    base_url = f"http://{omni_server.host}:{omni_server.port}"
-    url = f"{base_url}/v1/images/generations"
-
-    barrier = threading.Barrier(2)
-    results: list[tuple[int, int]] = []
-
-    def _call_generate(size: str) -> None:
-        payload: dict[str, Any] = {
-            "prompt": "cute cat playing with a ball",
-            "n": 1,
-            "size": size,
-            "response_format": "b64_json",
+    height = 832
+    width = 1248
+    chat_completion = client.chat.completions.create(
+        model=omni_server.model,
+        messages=messages,
+        extra_body={
+            "height": height,
+            "width": width,
             "num_inference_steps": 2,
-        }
-        barrier.wait()
-        response = requests.post(url, json=payload, timeout=120)
-        assert response.status_code == 200
-        data = response.json()
-        image_b64 = data["data"][0]["b64_json"]
-        image_bytes = base64.b64decode(image_b64)
-        img = Image.open(BytesIO(image_bytes))
-        img.load()
-        results.append(img.size)
+            "guidance_scale": 0.0,
+            "seed": 42,
+        },
+    )
 
-    threads = [
-        threading.Thread(target=_call_generate, args=("512x512",)),
-        threading.Thread(target=_call_generate, args=("768x512",)),
-    ]
-    for t in threads:
-        t.start()
-    for t in threads:
-        t.join()
+    assert len(chat_completion.choices) == 1
+    choice = chat_completion.choices[0]
+    assert choice.finish_reason == "stop"
+    assert choice.message.role == "assistant"
 
-    assert (512, 512) in results
-    assert (768, 512) in results
+    image_data_url = _extract_image_data_url(choice.message.content)
+    image_bytes = _decode_data_url_to_image_bytes(image_data_url)
+    img = Image.open(BytesIO(image_bytes))
+    img.load()
+    assert img.size == (width, height)

tests/e2e/online_serving/test_images_generations_lora.py

@@ -12,6 +12,10 @@ This validates:
 import base64
 import json
 import os
+import signal
+import subprocess
+import sys
+import time
 from io import BytesIO
 from pathlib import Path
 
@@ -21,8 +25,9 @@ import requests
 import torch
 from PIL import Image
 from safetensors.torch import save_file
+from vllm.utils.network_utils import get_open_port
 
-from tests.conftest import OmniServer
+from vllm_omni.utils.platform_utils import is_npu
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
@@ -34,8 +39,91 @@ SIZE = "256x256"
 SEED = 42
 
 
+class OmniServer:
+    """Omniserver for vLLM-Omni tests."""
+
+    def __init__(
+        self,
+        model: str,
+        serve_args: list[str],
+        *,
+        env_dict: dict[str, str] | None = None,
+    ) -> None:
+        self.model = model
+        self.serve_args = serve_args
+        self.env_dict = env_dict
+        self.proc: subprocess.Popen | None = None
+        self.host = "127.0.0.1"
+        self.port = get_open_port()
+
+    def _start_server(self) -> None:
+        env = os.environ.copy()
+        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+        if self.env_dict is not None:
+            env.update(self.env_dict)
+
+        cmd = [
+            sys.executable,
+            "-m",
+            "vllm_omni.entrypoints.cli.main",
+            "serve",
+            self.model,
+            "--omni",
+            "--host",
+            self.host,
+            "--port",
+            str(self.port),
+        ] + self.serve_args
+
+        print(f"Launching OmniServer with: {' '.join(cmd)}")
+        self.proc = subprocess.Popen(
+            cmd,
+            env=env,
+            cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),  # vllm-omni root
+            start_new_session=True,
+        )
+
+        # Wait for server to be ready.
+        max_wait = 1200
+        url = f"http://{self.host}:{self.port}/v1/models"
+        start_time = time.time()
+        while time.time() - start_time < max_wait:
+            try:
+                resp = requests.get(url, headers={"Authorization": "Bearer EMPTY"}, timeout=10)
+                if resp.status_code == 200:
+                    print(f"Server ready on {self.host}:{self.port}")
+                    return
+            except Exception:
+                pass
+            time.sleep(2)
+
+        raise RuntimeError(f"Server failed to become ready within {max_wait} seconds")
+
+    def __enter__(self):
+        self._start_server()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.proc is None:
+            return
+        try:
+            os.killpg(self.proc.pid, signal.SIGTERM)
+        except ProcessLookupError:
+            pass
+        try:
+            self.proc.wait(timeout=30)
+        except subprocess.TimeoutExpired:
+            try:
+                os.killpg(self.proc.pid, signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+            self.proc.wait()
+
+
 @pytest.fixture(scope="module")
 def omni_server():
+    if is_npu():
+        pytest.skip("Tongyi-MAI/Z-Image-Turbo is not supported on NPU yet.")
     with OmniServer(MODEL, ["--num-gpus", "1"]) as server:
         yield server
 
@@ -99,30 +187,17 @@ def _image_blue_tail_slice(img: Image.Image) -> np.ndarray:
     return tail
 
 
-def _slice_diff_stats(actual: np.ndarray, expected: np.ndarray) -> tuple[float, float]:
-    diff = np.abs(actual - expected)
-    return float(diff.max()), float(diff.mean())
-
-
-def _assert_slice_close(
-    actual: np.ndarray,
-    expected: np.ndarray,
-    *,
-    label: str,
-    base_max: float,
-    base_mean: float,
-) -> None:
+def _assert_slice_close(actual: np.ndarray, expected: np.ndarray, *, label: str) -> None:
     assert actual.shape == (3, 3)
     assert expected.shape == (3, 3)
-    max_diff, mean_diff = _slice_diff_stats(actual, expected)
+    diff = np.abs(actual - expected)
+    max_diff = float(diff.max())
+    mean_diff = float(diff.mean())
     # NOTE: Different attention backends / torch.compile can introduce small
     # floating-point drift that shows up as a few LSBs in uint8 pixels. Keep
     # the reset check tolerant but bounded to avoid flaky CI.
-    max_thresh = max(10.0, base_max + 4.0)
-    mean_thresh = max(6.0, base_mean + 4.0)
-    assert max_diff <= max_thresh and mean_diff <= mean_thresh, (
-        f"{label} slice mismatch (max={max_diff:.1f} > {max_thresh:.1f} or "
-        f"mean={mean_diff:.1f} > {mean_thresh:.1f}): {actual.tolist()}"
+    assert max_diff <= 5.0 and mean_diff <= 3.0, (
+        f"{label} slice mismatch (max={max_diff:.1f}, mean={mean_diff:.1f}): {actual.tolist()}"
     )
 
 
@@ -148,9 +223,6 @@ def test_images_generations_per_request_lora_switching(omni_server: OmniServer,
     # Base generation.
     base_img = _post_images(omni_server, _basic_payload())
     base_slice = _image_blue_tail_slice(base_img)
-    base_ref_img = _post_images(omni_server, _basic_payload())
-    base_ref_slice = _image_blue_tail_slice(base_ref_img)
-    base_ref_max, base_ref_mean = _slice_diff_stats(base_ref_slice, base_slice)
 
     # Adapter A: apply delta to V slice only.
     lora_a_dir = tmp_path / "zimage_lora_a"
@@ -177,17 +249,11 @@ def test_images_generations_per_request_lora_switching(omni_server: OmniServer,
     # Ensure switching back to no-LoRA restores the base output.
     base_img_2 = _post_images(omni_server, _basic_payload())
     base_slice_2 = _image_blue_tail_slice(base_img_2)
-    _, base_reset_mean = _slice_diff_stats(base_slice_2, base_slice)
-    _assert_slice_close(
-        base_slice_2,
-        base_slice,
-        label="base_after_reset",
-        base_max=base_ref_max,
-        base_mean=base_ref_mean,
-    )
+    _assert_slice_close(base_slice_2, base_slice, label="base_after_reset")
+    base_reset = float(np.abs(base_slice_2 - base_slice).mean())
 
     # Ensure LoRA effects are clearly above the baseline drift.
-    min_delta = max(base_reset_mean + 1.0, 1.5)
+    min_delta = base_reset + 0.5
     assert a_vs_base > min_delta, f"lora_a_vs_base drift too small: {a_vs_base} <= {min_delta}"
     assert b_vs_base > min_delta, f"lora_b_vs_base drift too small: {b_vs_base} <= {min_delta}"
     assert b_vs_a > min_delta, f"lora_b_vs_lora_a drift too small: {b_vs_a} <= {min_delta}"

tests/e2e/online_serving/test_qwen3_omni.py

@@ -25,10 +25,9 @@ from tests.conftest import (
     generate_synthetic_audio,
     generate_synthetic_image,
     generate_synthetic_video,
-    merge_base64_and_convert_to_text,
     modify_stage_config,
 )
-from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils import is_rocm
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 
@@ -58,7 +57,7 @@ def get_chunk_config():
 
 CHUNK_CONFIG_PATH = get_chunk_config()
 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325
-if current_omni_platform.is_rocm():
+if is_rocm():
     # ROCm stage config optimized for MI325 GPU
     stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
 else:
@@ -134,7 +133,7 @@ def dummy_messages_from_video_data(
 
 def get_prompt(prompt_type="text_only"):
     prompts = {
-        "text_only": "What is the capital of China? Answer in 20 words.",
+        "text_only": "What is the capital of China?",
         "mix": "What is recited in the audio? What is in this image? Describe the video briefly.",
     }
     return prompts.get(prompt_type, prompts["text_only"])
@@ -145,6 +144,7 @@ def get_max_batch_size(size_type="few"):
     return batch_sizes.get(size_type, 5)
 
 
+@pytest.mark.skipif(is_rocm(), reason="Test skipped on AMD environment due to known output issues")
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None:
     """
@@ -155,6 +155,11 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
     Input Setting: stream=True
     Datasets: single request
     """
+    # TODO：This skip will be removed when the chunk scenario supports multimodal input.
+    param = request.node.callspec.params.get("omni_server")
+
+    if param[1] == CHUNK_CONFIG_PATH:
+        pytest.skip("The current chunk scenario does not support multimodal.")
 
     # Test single completion
     e2e_list = list()
@@ -174,7 +179,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
     chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, stream=True)
 
     text_content = ""
-    audio_data = []
+    audio_data = None
     for chunk in chat_completion:
         for choice in chunk.choices:
             if hasattr(choice, "delta"):
@@ -185,7 +190,11 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
             modality = getattr(chunk, "modality", None)
 
             if modality == "audio" and content:
-                audio_data.append(content)
+                # Audio chunk - content
+                if audio_data is None:
+                    audio_data = content
+                else:
+                    audio_data += content
             elif modality == "text" and content:
                 # Text chunk - accumulate text content
                 text_content += content if content else ""
@@ -207,7 +216,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
     ), "The output does not contain any of the keywords."
 
     # Verify text output same as audio output
-    audio_content = merge_base64_and_convert_to_text(audio_data)
+    audio_content = convert_audio_to_text(audio_data)
     print(f"text content is: {text_content}")
     print(f"audio content is: {audio_content}")
     similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
@@ -215,6 +224,7 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N
     assert similarity > 0.9, "The audio content is not same as the text"
 
 
+@pytest.mark.skipif(is_rocm(), reason="Test skipped on AMD environment due to known output issues")
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_text_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
     """

tests/e2e/stage_configs/qwen3_omni_ci.yaml

@@ -7,23 +7,22 @@
 stage_args:
 - stage_id: 0
   runtime:
-    devices: "0"
+    devices: "0,1"
     max_batch_size: 5
   engine_args:
     model_stage: thinker
     model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_type: ar
+    worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
     scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    gpu_memory_utilization: 0.9
+    gpu_memory_utilization: 0.6
     enforce_eager: false
     trust_remote_code: true
     engine_output_type: latent  # Output hidden states for talker
     distributed_executor_backend: "mp"
     max_num_batched_tokens: 32768
-    max_model_len: 32768
     enable_prefix_caching: false
     hf_config_name: thinker_config
-    tensor_parallel_size: 1
+    tensor_parallel_size: 2
   final_output: true
   final_output_type: text
   is_comprehension: true
@@ -44,15 +43,14 @@ stage_args:
   engine_args:
     model_stage: talker
     model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_type: ar
+    worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
     scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    gpu_memory_utilization: 0.6
+    gpu_memory_utilization: 0.3
     enforce_eager: false
     trust_remote_code: true
     engine_output_type: latent  # Output codec codes for code2wav
     enable_prefix_caching: false
     max_num_batched_tokens: 32768
-    max_model_len: 32768
     distributed_executor_backend: "mp"
     hf_config_name: talker_config
   engine_input_source: [0]
@@ -68,12 +66,12 @@ stage_args:
 
 - stage_id: 2
   runtime:
-    devices: "1"
+    devices: "0"
     max_batch_size: 1
   engine_args:
     model_stage: code2wav
     model_arch: Qwen3OmniMoeForConditionalGeneration
-    worker_type: generation
+    worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
     scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
     enforce_eager: true
     trust_remote_code: true

tests/entrypoints/openai_api/test_image_server.py

@@ -9,7 +9,6 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py.
 
 import base64
 import io
-from argparse import Namespace
 from unittest.mock import AsyncMock, Mock
 
 import pytest
@@ -116,11 +115,9 @@ class FakeAsyncOmni:
         self.stage_list = ["llm", "diffusion"]
         self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
         self.captured_sampling_params_list = None
-        self.captured_prompt = None
 
     async def generate(self, prompt, request_id, sampling_params_list):
         self.captured_sampling_params_list = sampling_params_list
-        self.captured_prompt = prompt
         images = [Image.new("RGB", (64, 64), color="green")]
         yield MockGenerationResult(images)
 
@@ -136,8 +133,6 @@ def mock_async_diffusion():
         # Return n PIL images wrapped in result object
         print("!!!!!!!!!!!!!!!!!!!!! kwargs", kwargs)
         n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
-        mock.captured_sampling_params_list = kwargs["sampling_params_list"]
-        mock.captured_prompt = kwargs["prompt"]
         images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
         return MockGenerationResult(images)
 
@@ -160,10 +155,6 @@ def test_client(mock_async_diffusion):
     app.state.diffusion_engine = mock_async_diffusion  # Also set for health endpoint
     app.state.stage_configs = [{"stage_type": "diffusion"}]
     app.state.diffusion_model_name = "Qwen/Qwen-Image"  # For models endpoint
-    app.state.args = Namespace(
-        default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
-        max_generated_image_size=4096,  # 64*64
-    )
 
     return TestClient(app)
 
@@ -180,10 +171,7 @@ def async_omni_test_client():
 
     app.state.engine_client = FakeAsyncOmni()
     app.state.stage_configs = [{"stage_type": "llm"}, {"stage_type": "diffusion"}]
-    app.state.args = Namespace(
-        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
-        max_generated_image_size=4096,  # 64*64
-    )
+
     return TestClient(app)
 
 
@@ -538,279 +526,3 @@ def test_model_field_omitted_works(test_client):
         },
     )
     assert response.status_code == 200
-
-
-def make_test_image_bytes(size=(64, 64)) -> bytes:
-    img = Image.new(
-        "RGB",
-        size,
-    )
-    buf = io.BytesIO()
-    img.save(buf, format="PNG")
-    return buf.getvalue()
-
-
-def test_image_edit_images_processing(async_omni_test_client):
-    img_bytes_1 = make_test_image_bytes((16, 16))
-    img_bytes_2 = make_test_image_bytes((32, 32))
-
-    # uploadfile with image key
-    response = async_omni_test_client.post(
-        "/v1/images/edits",
-        files=[
-            ("image", img_bytes_1),
-            ("image", img_bytes_2),
-        ],
-        data={"prompt": "hello world."},
-    )
-    assert response.status_code == 200
-    engine = async_omni_test_client.app.state.engine_client
-    captured_prompt = engine.captured_prompt
-    processed_images = captured_prompt["multi_modal_data"]["image"]
-    assert len(processed_images) == 2
-    assert isinstance(processed_images[0], Image.Image)
-    assert isinstance(processed_images[1], Image.Image)
-    assert processed_images[0].size == (16, 16)
-    assert processed_images[1].size == (32, 32)
-
-    # uploadfile with image[] key
-    response = async_omni_test_client.post(
-        "/v1/images/edits",
-        files=[
-            ("image[]", img_bytes_2),
-            ("image[]", img_bytes_1),
-        ],
-        data={"prompt": "hello world."},
-    )
-
-    assert response.status_code == 200
-    engine = async_omni_test_client.app.state.engine_client
-    captured_prompt = engine.captured_prompt
-    processed_images = captured_prompt["multi_modal_data"]["image"]
-    assert len(processed_images) == 2
-    assert isinstance(processed_images[0], Image.Image)
-    assert isinstance(processed_images[1], Image.Image)
-    assert processed_images[0].size == (32, 32)
-    assert processed_images[1].size == (16, 16)
-
-    # base64 url
-    buf1 = io.BytesIO()
-    img1 = Image.new("RGB", (16, 16))
-    img1.save(buf1, format="PNG")
-    b64_1 = "data:image/png;base64," + base64.b64encode(buf1.getvalue()).decode()
-
-    buf2 = io.BytesIO()
-    img2 = Image.new("RGB", (24, 24))
-    img2.save(buf2, format="PNG")
-    b64_2 = "data:image/png;base64," + base64.b64encode(buf2.getvalue()).decode()
-
-    response = async_omni_test_client.post(
-        "/v1/images/edits",
-        data={
-            "prompt": "hello from base64",
-            "url": [b64_1, b64_2],
-        },
-    )
-    assert response.status_code == 200
-    processed_images = engine.captured_prompt["multi_modal_data"]["image"]
-    assert len(processed_images) == 2
-    assert isinstance(processed_images[0], Image.Image)
-    assert isinstance(processed_images[1], Image.Image)
-    assert processed_images[0].size == (16, 16)
-    assert processed_images[1].size == (24, 24)
-
-
-def test_image_edit_parameter_pass(async_omni_test_client):
-    img_bytes_1 = make_test_image_bytes((16, 16))
-
-    # uploadfile with image key
-    response = async_omni_test_client.post(
-        "/v1/images/edits",
-        files=[("image", img_bytes_1)],
-        data={
-            "prompt": "hello world.",
-            "size": "16x24",
-            "output_format": "jpeg",
-            "num_inference_steps": 20,
-            "guidance_scale": 8.0,
-            "seed": 1234,
-            "negative_prompt": "negative",
-            "n": 2,
-        },
-    )
-    assert response.status_code == 200
-    engine = async_omni_test_client.app.state.engine_client
-    captured_prompt = engine.captured_prompt
-    captured_sampling_params = engine.captured_sampling_params_list[-1]
-
-    assert captured_prompt["prompt"] == "hello world."
-    assert captured_prompt["negative_prompt"] == "negative"
-    assert captured_sampling_params.num_inference_steps == 20
-    assert captured_sampling_params.guidance_scale == 8.0
-    assert captured_sampling_params.seed == 1234
-    assert captured_sampling_params.num_outputs_per_prompt == 2
-    assert captured_sampling_params.width == 16
-    assert captured_sampling_params.height == 24
-
-    data = response.json()
-    # All images should be valid
-    for img_data in data["data"]:
-        assert "b64_json" in img_data
-        img_bytes = base64.b64decode(img_data["b64_json"])
-        img = Image.open(io.BytesIO(img_bytes))
-        assert img.format.lower() == "jpeg"
-        assert data["output_format"] == "jpeg"
-        assert data["size"] == "16x24"
-
-
-def test_image_edit_parameter_default(async_omni_test_client):
-    img_bytes_1 = make_test_image_bytes((24, 16))
-
-    # uploadfile with image key
-    response = async_omni_test_client.post(
-        "/v1/images/edits",
-        files=[("image", img_bytes_1)],
-        data={
-            "prompt": "hello world.",
-            "size": "auto",
-        },
-    )
-    assert response.status_code == 200
-    engine = async_omni_test_client.app.state.engine_client
-    captured_sampling_params = engine.captured_sampling_params_list[-1]
-
-    assert captured_sampling_params.width == 24
-    assert captured_sampling_params.height == 16
-    assert captured_sampling_params.num_outputs_per_prompt == 1
-    assert captured_sampling_params.num_inference_steps == 4
-    assert captured_sampling_params.guidance_scale == 7.5
-
-    response = async_omni_test_client.post(
-        "/v1/images/edits",
-        files=[("image", img_bytes_1)],
-        data={
-            "prompt": "hello world.",
-            "size": "96x96",
-        },
-    )
-    assert response.status_code == 400
-
-
-def test_image_edit_parameter_default_single_stage(test_client):
-    img_bytes_1 = make_test_image_bytes((24, 16))
-
-    # uploadfile with image key
-    response = test_client.post(
-        "/v1/images/edits",
-        files=[("image", img_bytes_1)],
-        data={
-            "prompt": "hello world.",
-        },
-    )
-    assert response.status_code == 200
-    engine = test_client.app.state.engine_client
-    captured_sampling_params = engine.captured_sampling_params_list[0]
-
-    assert captured_sampling_params.width == 24
-    assert captured_sampling_params.height == 16
-    assert captured_sampling_params.num_outputs_per_prompt == 1
-    assert captured_sampling_params.num_inference_steps == 4
-    assert captured_sampling_params.guidance_scale == 7.5
-
-    response = test_client.post(
-        "/v1/images/edits",
-        files=[("image", img_bytes_1)],
-        data={
-            "prompt": "hello world.",
-            "size": "96x96",
-        },
-    )
-    assert response.status_code == 400
-
-
-def test_image_edit_compression_jpeg(test_client):
-    img_bytes_1 = make_test_image_bytes((16, 16))
-    # uploadfile with image key
-    response = test_client.post(
-        "/v1/images/edits",
-        files=[("image", img_bytes_1)],
-        data={"prompt": "hello world.", "output_format": "jpeg", "output_compression": 100},
-    )
-    assert response.status_code == 200
-    data = response.json()
-    img_bytes_100 = base64.b64decode(data["data"][0]["b64_json"])
-    img = Image.open(io.BytesIO(img_bytes_100))
-    assert img.format.lower() == "jpeg"
-
-    response = test_client.post(
-        "/v1/images/edits",
-        files=[("image", img_bytes_1)],
-        data={
-            "prompt": "hello world.",
-            "output_format": "jpeg",
-            "output_compression": 50,
-        },
-    )
-    assert response.status_code == 200
-    data = response.json()
-    img_bytes_50 = base64.b64decode(data["data"][0]["b64_json"])
-
-    response = test_client.post(
-        "/v1/images/edits",
-        files=[("image", img_bytes_1)],
-        data={
-            "prompt": "hello world.",
-            "output_format": "jpeg",
-            "output_compression": 10,
-        },
-    )
-    assert response.status_code == 200
-    data = response.json()
-    img_bytes_10 = base64.b64decode(data["data"][0]["b64_json"])
-
-    assert len(img_bytes_10) < len(img_bytes_50)
-    assert len(img_bytes_50) < len(img_bytes_100)
-
-
-def test_image_edit_compression_png(async_omni_test_client):
-    img_bytes_1 = make_test_image_bytes((16, 16))
-    # uploadfile with image key
-    response = async_omni_test_client.post(
-        "/v1/images/edits",
-        files=[("image", img_bytes_1)],
-        data={"prompt": "hello world.", "output_format": "PNG", "output_compression": 100},
-    )
-    assert response.status_code == 200
-    data = response.json()
-    img_bytes_100 = base64.b64decode(data["data"][0]["b64_json"])
-    img = Image.open(io.BytesIO(img_bytes_100))
-    assert img.format.lower() == "png"
-
-    response = async_omni_test_client.post(
-        "/v1/images/edits",
-        files=[("image", img_bytes_1)],
-        data={
-            "prompt": "hello world.",
-            "output_format": "PNG",
-            "output_compression": 50,
-        },
-    )
-    assert response.status_code == 200
-    data = response.json()
-    img_bytes_50 = base64.b64decode(data["data"][0]["b64_json"])
-
-    response = async_omni_test_client.post(
-        "/v1/images/edits",
-        files=[("image", img_bytes_1)],
-        data={
-            "prompt": "hello world.",
-            "output_format": "PNG",
-            "output_compression": 10,
-        },
-    )
-    assert response.status_code == 200
-    data = response.json()
-    img_bytes_10 = base64.b64decode(data["data"][0]["b64_json"])
-
-    assert len(img_bytes_10) < len(img_bytes_50)
-    assert len(img_bytes_50) < len(img_bytes_100)

tests/entrypoints/openai_api/test_serving_speech.py

@@ -201,13 +201,6 @@ def test_app():
     app = FastAPI()
     app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None)
 
-    # Add list_voices endpoint
-    async def list_voices():
-        speakers = sorted(speech_server.supported_speakers) if speech_server.supported_speakers else []
-        return {"voices": speakers}
-
-    app.add_api_route("/v1/audio/voices", list_voices, methods=["GET"])
-
     return app
 
 
@@ -275,11 +268,6 @@ class TestSpeechAPI:
         assert isinstance(audio_obj, CreateAudio)
         assert audio_obj.speed == 2.5
 
-    def test_list_voices_endpoint(self, client):
-        response = client.get("/v1/audio/voices")
-        assert response.status_code == 200
-        assert "voices" in response.json()
-
 
 class TestTTSMethods:
     """Unit tests for TTS validation and parameter building."""
@@ -323,9 +311,9 @@ class TestTTSMethods:
         req = OpenAICreateSpeechRequest(input="Hello", language="InvalidLang")
         assert "Invalid language" in speech_server._validate_tts_request(req)
 
-        # When no speakers loaded, any voice is accepted (unconstrained)
+        # Invalid speaker
         req = OpenAICreateSpeechRequest(input="Hello", voice="Invalid")
-        assert speech_server._validate_tts_request(req) is None
+        assert "Invalid speaker" in speech_server._validate_tts_request(req)
 
         # Valid request
         req = OpenAICreateSpeechRequest(input="Hello", voice="Vivian")
@@ -354,26 +342,3 @@ class TestTTSMethods:
         assert params["speaker"] == ["Ryan"]
         assert params["language"] == ["English"]
         assert params["task_type"] == ["CustomVoice"]
-
-    def test_load_supported_speakers(self):
-        """Test _load_supported_speakers."""
-        mock_engine_client = MagicMock()
-        mock_engine_client.errored = False
-        mock_engine_client.stage_list = None
-
-        # Mock talker_config with mixed-case speaker names
-        mock_talker_config = MagicMock()
-        mock_talker_config.spk_id = {"Ryan": 0, "Vivian": 1, "Aiden": 2}
-        mock_engine_client.model_config.hf_config.talker_config = mock_talker_config
-
-        mock_models = MagicMock()
-        mock_models.is_base_model.return_value = True
-
-        server = OmniOpenAIServingSpeech(
-            engine_client=mock_engine_client,
-            models=mock_models,
-            request_logger=MagicMock(),
-        )
-
-        # Verify speakers are normalized to lowercase
-        assert server.supported_speakers == {"ryan", "vivian", "aiden"}

tests/entrypoints/test_stage_utils.py

@@ -1,6 +1,5 @@
 import os
 import sys
-from unittest.mock import MagicMock
 
 import pytest
 
@@ -43,14 +42,6 @@ def _make_dummy_torch(call_log):
     return _Torch
 
 
-def _make_mock_platform(device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"):
-    """Create a mock platform for testing."""
-    mock_platform = MagicMock()
-    mock_platform.device_type = device_type
-    mock_platform.device_control_env_var = env_var
-    return mock_platform
-
-
 @pytest.mark.usefixtures("clean_gpu_memory_between_tests")
 def test_set_stage_devices_respects_logical_ids(monkeypatch):
     # Preserve an existing logical mapping and ensure devices "0,1" map through it.
@@ -58,51 +49,9 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch):
     call_log: list[int] = []
     dummy_torch = _make_dummy_torch(call_log)
     monkeypatch.setitem(sys.modules, "torch", dummy_torch)
-
-    # Mock the platform at the source module where it's defined
-    mock_platform = _make_mock_platform(device_type="cuda", env_var="CUDA_VISIBLE_DEVICES")
-    monkeypatch.setattr(
-        "vllm_omni.platforms.current_omni_platform",
-        mock_platform,
-    )
+    monkeypatch.setattr("vllm_omni.utils.detect_device_type", lambda: "cuda")
+    monkeypatch.setattr("vllm_omni.utils.get_device_control_env_var", lambda: "CUDA_VISIBLE_DEVICES")
 
     set_stage_devices(stage_id=0, devices="0,1")
 
     assert os.environ["CUDA_VISIBLE_DEVICES"] == "6,7"
-
-
-@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
-def test_set_stage_devices_npu_platform(monkeypatch):
-    """Test that set_stage_devices works correctly for NPU platform."""
-    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "4,5")
-    call_log: list[int] = []
-
-    # Create NPU mock torch
-    class _Npu:
-        @staticmethod
-        def is_available():
-            return True
-
-        @staticmethod
-        def set_device(idx):
-            call_log.append(idx)
-
-        @staticmethod
-        def device_count():
-            return 2
-
-    class _NpuTorch:
-        npu = _Npu
-
-    monkeypatch.setitem(sys.modules, "torch", _NpuTorch)
-
-    # Mock NPU platform at the source module where it's defined
-    mock_platform = _make_mock_platform(device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES")
-    monkeypatch.setattr(
-        "vllm_omni.platforms.current_omni_platform",
-        mock_platform,
-    )
-
-    set_stage_devices(stage_id=0, devices="0,1")
-
-    assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "4,5"

tests/utils.py

@@ -77,36 +77,11 @@ def wait_for_gpu_memory_to_clear(
     threshold_ratio: float | None = None,
     timeout_s: float = 120,
 ) -> None:
-    import gc
-
     assert threshold_bytes is not None or threshold_ratio is not None
     # Use nvml instead of pytorch to reduce measurement error from torch cuda
     # context.
     devices = get_physical_device_indices(devices)
     start_time = time.time()
-
-    # Print waiting start information
-    device_list = ", ".join(str(d) for d in devices)
-    if threshold_bytes is not None:
-        threshold_str = f"{threshold_bytes / 2**30:.2f} GiB"
-        condition_str = f"Memory usage ≤ {threshold_str}"
-    else:
-        threshold_percent = threshold_ratio * 100
-        threshold_str = f"{threshold_percent:.1f}%"
-        condition_str = f"Memory usage ratio ≤ {threshold_str}"
-
-    print(f"[GPU Memory Monitor] Waiting for GPU {device_list} to free memory, Condition: {condition_str}")
-
-    # Define the is_free function based on threshold type
-    if threshold_bytes is not None:
-
-        def is_free(used, total):
-            return used <= threshold_bytes / 2**30
-    else:
-
-        def is_free(used, total):
-            return used / total <= threshold_ratio
-
     while True:
         output: dict[int, str] = {}
         output_raw: dict[int, tuple[float, float]] = {}
@@ -122,44 +97,33 @@ def wait_for_gpu_memory_to_clear(
                 gb_used = mem_info.used / 2**30
                 gb_total = mem_info.total / 2**30
             output_raw[device] = (gb_used, gb_total)
-            # Format to more readable form
-            usage_percent = (gb_used / gb_total) * 100 if gb_total > 0 else 0
-            output[device] = f"{gb_used:.1f}GiB/{gb_total:.1f}GiB ({usage_percent:.1f}%)"
+            output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
 
-        # Optimized GPU memory status print
-        print("[GPU Memory Status] Current usage:")
-        for device_id, mem_info in output.items():
-            print(f"  GPU {device_id}: {mem_info}")
+        print("gpu memory used/total (GiB): ", end="")
+        for k, v in output.items():
+            print(f"{k}={v}; ", end="")
+        print("")
+
+        if threshold_bytes is not None:
+
+            def is_free(used, total):
+                return used <= threshold_bytes / 2**30  # noqa E731
+
+            threshold = f"{threshold_bytes / 2**30} GiB"
+        else:
+
+            def is_free(used, total):
+                return used / total <= threshold_ratio  # noqa E731
+
+            threshold = f"{threshold_ratio:.2f}"
 
-        # Calculate waiting duration
         dur_s = time.time() - start_time
-        elapsed_minutes = dur_s / 60
-
-        # Check if all devices meet the condition
         if all(is_free(used, total) for used, total in output_raw.values()):
-            # Optimized completion message
-            print(f"[GPU Memory Freed] Devices {device_list} meet memory condition")
-            print(f"   Condition: {condition_str}")
-            print(f"   Wait time: {dur_s:.1f} seconds ({elapsed_minutes:.1f} minutes)")
-            print("   Final status:")
-            for device_id, mem_info in output.items():
-                print(f"     GPU {device_id}: {mem_info}")
+            print(f"Done waiting for free GPU memory on devices {devices=} ({threshold=}) {dur_s=:.02f}")
             break
 
-        # Check timeout
         if dur_s >= timeout_s:
-            raise ValueError(
-                f"[GPU Memory Timeout] Devices {device_list} still don't meet memory condition after {dur_s:.1f} seconds\n"
-                f"Condition: {condition_str}\n"
-                f"Current status:\n" + "\n".join(f"  GPU {device}: {output[device]}" for device in devices)
-            )
-
-        # Add waiting hint (optional)
-        if dur_s > 10 and int(dur_s) % 10 == 0:  # Show hint every 10 seconds
-            print(f"Waiting... Already waited {dur_s:.1f} seconds ({elapsed_minutes:.1f} minutes)")
-
-        gc.collect()
-        torch.cuda.empty_cache()
+            raise ValueError(f"Memory of devices {devices=} not free after {dur_s=:.02f} ({threshold=})")
 
         time.sleep(5)
 

tests/worker/test_omni_gpu_model_runner.py

@@ -1,5 +1,4 @@
 from contextlib import contextmanager
-from types import SimpleNamespace
 
 import torch
 
@@ -87,12 +86,6 @@ def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
 
     runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
 
-    def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
-        batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
-        return (False, batch_desc, None, None, None)
-
-    monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
-
     # Initialize per-request embeds (batch-major inside talker_mtp_inputs_embeds)
     runner.talker_mtp_inputs_embeds.gpu[0] = torch.tensor([1.0, 2.0, 3.0, 4.0])
     runner.talker_mtp_inputs_embeds.gpu[1] = torch.tensor([10.0, 20.0, 30.0, 40.0])

vllm_omni/benchmarks/data_modules/random_multi_modal_dataset.py

@@ -1,152 +0,0 @@
-import base64
-import io
-import logging
-from collections.abc import Mapping
-from typing import Any
-
-import numpy as np
-import soundfile as sf
-import torch
-from vllm.benchmarks.datasets import RandomMultiModalDataset, process_image, process_video
-
-logger = logging.getLogger(__name__)
-
-
-def process_audio(audio: Any) -> Mapping[str, Any]:
-    """
-    Process a single audio input and return a multimedia content dictionary.
-
-    Supports the following input types:
-
-    1. Dictionary with raw audio bytes: - Expects a dict with a 'bytes' key
-       containing raw audio data.
-
-    2. String input: - Treats the string as a URL or local file path.  -
-       Prepends "file://" if the string doesn't start with "http://" or
-       "file://".  - Returns a dictionary with the audio URL.
-
-    Raises:
-        ValueError: If the input is not a supported type.
-    """
-    if isinstance(audio, dict) and "bytes" in audio:
-        audio_bytes = audio["bytes"]
-        audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
-        return {
-            "type": "audio_url",
-            "audio_url": {"url": f"data:audio/mpeg;base64,{audio_base64}"},
-        }
-    if isinstance(audio, str):
-        audio_url = audio if audio.startswith(("http://", "https://", "file://")) else f"file://{audio}"
-        return {"type": "audio_url", "audio_url": {"url": audio_url}}
-
-    raise ValueError(
-        f"Invalid audio input {audio}. Must be a string of local path/remote url, "
-        f"or a dictionary with raw audio bytes in the form of `{{'bytes': raw_audio_bytes}}`."
-    )
-
-
-# -----------------------------------------------------------------------------
-# MultiModalDataset Implementation
-# -----------------------------------------------------------------------------
-class OmniRandomMultiModalDataset(RandomMultiModalDataset):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def generate_synthetic_audio(
-        self,
-        duration: int,  # seconds
-        num_channels: int,  # 1:Mono，2:Stereo 5:5.1 surround sound
-    ) -> dict[str, Any]:
-        """Generate synthetic audio with random values.
-        Default use 48000Hz.
-        """
-        sample_rate = 48000
-        num_samples = int(sample_rate * duration)
-        audio_data = self._rng.uniform(-0.5, 0.5, (num_samples, num_channels))
-        audio_data = np.clip(audio_data, -1.0, 1.0)
-        audio_tensor = torch.FloatTensor(audio_data.T)
-        audio_np = audio_tensor.numpy()
-
-        buffer = io.BytesIO()
-
-        sf.write(buffer, audio_np.T, sample_rate, format="wav")
-
-        buffer.seek(0)
-        audio_bytes = buffer.read()
-        buffer.close()
-        return {
-            "bytes": audio_bytes,
-        }
-
-    def generate_mm_item(
-        self,
-        mm_item_config: tuple[int, int, int],
-    ) -> Mapping[str, Any]:
-        """
-        Create synthetic images and videos and
-        apply process_image/process_video respectively.
-        This follows the OpenAI API chat completions
-        https://github.com/openai/openai-python
-        """
-
-        if self.map_config_to_modality(mm_item_config) == "image":
-            return process_image(self.generate_synthetic_image(mm_item_config[1], mm_item_config[0]))
-        elif self.map_config_to_modality(mm_item_config) == "video":
-            return process_video(self.generate_synthetic_video(mm_item_config[1], mm_item_config[0], mm_item_config[2]))
-        elif self.map_config_to_modality(mm_item_config) == "audio":
-            return process_audio(self.generate_synthetic_audio(mm_item_config[1], mm_item_config[2]))
-        else:
-            raise ValueError(f"Invalid multimodal item configuration: {mm_item_config}")
-
-    def generate_synthetic_video(self, width: int, height: int, num_frames: int) -> Any:
-        """Generate synthetic video with random values."""
-        import imageio
-
-        video_data = self._rng.integers(
-            0,
-            256,
-            (num_frames, height, width, 3),
-            dtype=np.uint8,
-        )
-        buffer = io.BytesIO()
-        writer_kwargs = {
-            "format": "mp4",
-            "fps": 30,
-            "codec": "libx264",
-            "quality": 7,
-            "pixelformat": "yuv420p",
-            "macro_block_size": 16,
-            "ffmpeg_params": [
-                "-preset",
-                "medium",
-                "-crf",
-                "23",
-                "-movflags",
-                "+faststart",
-                "-pix_fmt",
-                "yuv420p",
-                "-vf",
-                f"scale={width}:{height}",
-            ],
-        }
-
-        with imageio.get_writer(buffer, **writer_kwargs) as writer:
-            for frame_idx in range(num_frames):
-                writer.append_data(video_data[frame_idx])
-        buffer.seek(0)
-        video_bytes = buffer.read()
-
-        return {
-            "bytes": video_bytes,
-        }
-
-    def map_config_to_modality(self, config: tuple[int, int, int]) -> str:
-        """Map the configuration to the modality."""
-        if config[0] == 0:
-            return "audio"
-        elif config[-1] == 1:
-            return "image"
-        elif config[-1] > 1:
-            return "video"
-        else:
-            raise ValueError(f"Invalid multimodal item configuration: {config}")

vllm_omni/benchmarks/metrics/metrics.py

@@ -1,330 +0,0 @@
-import warnings
-from dataclasses import dataclass
-
-import numpy as np
-from transformers import PreTrainedTokenizerBase
-from vllm.benchmarks.datasets import SampleRequest
-from vllm.benchmarks.lib.endpoint_request_func import RequestFuncOutput
-from vllm.benchmarks.serve import MILLISECONDS_TO_SECONDS_CONVERSION, TERM_PLOTLIB_AVAILABLE, BenchmarkMetrics, TaskType
-
-
-@dataclass
-class MultiModalsBenchmarkMetrics(BenchmarkMetrics):
-    mean_audio_ttfp_ms: float = 0.0
-    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_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
-
-
-def print_metrics(
-    task_type,
-    selected_percentile_metrics,
-    max_concurrency,
-    request_rate,
-    benchmark_duration,
-    goodput_config_dict,
-    metrics: MultiModalsBenchmarkMetrics,
-):
-    print("{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
-    print("{:<40} {:<10}".format("Successful requests:", metrics.completed))
-    print("{:<40} {:<10}".format("Failed requests:", metrics.failed))
-    if max_concurrency is not None:
-        print("{:<40} {:<10}".format("Maximum request concurrency:", max_concurrency))
-    if request_rate != float("inf"):
-        print("{:<40} {:<10.2f}".format("Request rate configured (RPS):", request_rate))
-    print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration))
-    print("{:<40} {:<10.2f}".format("Request throughput (req/s):", metrics.request_throughput))
-    if goodput_config_dict:
-        print("{:<40} {:<10.2f}".format("Request goodput (req/s):", metrics.request_goodput))
-    if isinstance(metrics, MultiModalsBenchmarkMetrics):
-        print("{:<40} {:<10.2f}".format("Peak concurrent requests:", metrics.max_concurrent_requests))
-    if task_type != TaskType.GENERATION or "e2el" in selected_percentile_metrics:
-        process_one_metric("e2el", metrics)
-    print_text_metrics(task_type, selected_percentile_metrics, metrics)
-    if task_type == TaskType.GENERATION:
-        print_audio_metrics(selected_percentile_metrics, metrics)
-    print("=" * 50)
-
-
-def print_text_metrics(task_type, selected_percentile_metrics, metrics: MultiModalsBenchmarkMetrics):
-    print("{s:{c}^{n}}".format(s=" Text Result ", n=50, c="="))
-    print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input))
-    if isinstance(metrics, MultiModalsBenchmarkMetrics):
-        print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output))
-        print("{:<40} {:<10.2f}".format("Output token throughput (tok/s):", metrics.output_throughput))
-        print("{:<40} {:<10.2f}".format("Peak output token throughput (tok/s):", metrics.max_output_tokens_per_s))
-        print("{:<40} {:<10.2f}".format("Peak concurrent requests:", metrics.max_concurrent_requests))
-    print("{:<40} {:<10.2f}".format("Total Token throughput (tok/s):", metrics.total_token_throughput))
-
-    if task_type == TaskType.GENERATION:
-        for metric in selected_percentile_metrics:
-            if metric == "e2el":
-                continue
-            if not metric.startswith("audio"):
-                process_one_metric(metric, metrics)
-
-
-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}".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:
-        if metric.startswith("audio"):
-            process_one_metric(metric, metrics)
-
-
-def process_one_metric(
-    metric_attribute_name: str,
-    metrics: MultiModalsBenchmarkMetrics,
-):
-    metric_header_map = {
-        "ttft": "Time to First Token",
-        "tpot": "Time per Output Token (excl. 1st token)",
-        "itl": "Inter-token Latency",
-        "e2el": "End-to-end Latency",
-        "audio_ttfp": "Time to First Packet",
-        "audio_rtf": "Real Time Factor",
-    }
-
-    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"
-
-    suffix = "" if is_audio_rtf else "_ms"
-    unit_suffix = "" if is_audio_rtf else " (ms)"
-
-    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}")
-
-    median_attr_name = f"median_{metric_attribute_name}{suffix}"
-    median_value = getattr(metrics, median_attr_name, 0.0)
-    print(f"{f'Median {metric_attribute_name.upper()}{unit_suffix}:':<40} {median_value:<10.2f}")
-
-    percentiles_attr_name = f"percentiles_{metric_attribute_name}{suffix}"
-    percentiles = getattr(metrics, percentiles_attr_name, [])
-
-    for percentile, value in percentiles:
-        p_str = str(int(percentile)) if percentile.is_integer() else str(percentile)
-        label = f"P{p_str} {metric_attribute_name.upper()}{unit_suffix}:"
-        print(f"{label:<40} {value:<10.2f}")
-
-
-def calculate_metrics(
-    input_requests: list[SampleRequest],
-    outputs: list[RequestFuncOutput],
-    dur_s: float,
-    tokenizer: PreTrainedTokenizerBase,
-    selected_percentiles: list[float],
-    goodput_config_dict: dict[str, float],
-    task_type,
-    selected_percentile_metrics,
-    max_concurrency,
-    request_rate,
-    benchmark_duration,
-) -> tuple[BenchmarkMetrics, list[int]]:
-    """Calculate the metrics for the benchmark.
-
-    Args:
-        input_requests: The input requests.
-        outputs: The outputs of the requests.
-        dur_s: The duration of the benchmark.
-        tokenizer: The tokenizer to use.
-        selected_percentiles: The percentiles to select.
-        goodput_config_dict: The goodput configuration.
-
-    Returns:
-        A tuple of the benchmark metrics and the actual output lengths.
-    """
-    actual_output_lens: list[int] = []
-    total_input = 0
-    completed = 0
-    good_completed = 0
-    itls: list[float] = []
-    tpots: list[float] = []
-    all_tpots: list[float] = []
-    ttfts: list[float] = []
-    e2els: list[float] = []
-    audio_ttfps: list[float] = []
-    audio_rtfs: list[float] = []
-    audio_duration: list[float] = []
-    audio_frames: list[int] = []
-    for i in range(len(outputs)):
-        if outputs[i].success:
-            output_len = outputs[i].output_tokens
-
-            if not output_len:
-                # We use the tokenizer to count the number of output tokens
-                # for some serving backends instead of looking at
-                # len(outputs[i].itl) since multiple output tokens may be
-                # bundled together
-                # Note : this may inflate the output token count slightly
-                output_len = len(tokenizer(outputs[i].generated_text, add_special_tokens=False).input_ids)
-            actual_output_lens.append(output_len)
-            total_input += input_requests[i].prompt_len
-            tpot = 0
-            if output_len > 1:
-                latency_minus_ttft = outputs[i].latency - outputs[i].ttft
-                tpot = latency_minus_ttft / (output_len - 1)
-                tpots.append(tpot)
-            # Note: if output_len <= 1, we regard tpot as 0 for goodput
-            all_tpots.append(tpot)
-            itls += outputs[i].itl
-            ttfts.append(outputs[i].ttft)
-            audio_ttfps.append(getattr(outputs[i], "audio_ttfp", 0.0))
-            audio_rtfs.append(getattr(outputs[i], "audio_rtf", 0.0))
-            audio_duration.append(getattr(outputs[i], "audio_duration", 0.0))
-            audio_frames.append(getattr(outputs[i], "audio_frames", 0.0))
-            e2els.append(outputs[i].latency)
-            completed += 1
-        else:
-            actual_output_lens.append(0)
-
-    if goodput_config_dict:
-        valid_metrics = []
-        slo_values = []
-
-        if "ttft" in goodput_config_dict:
-            valid_metrics.append(ttfts)
-            slo_values.append(goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
-        if "audio_ttft" in goodput_config_dict:
-            valid_metrics.append(audio_ttfps)
-            slo_values.append(goodput_config_dict["audio_ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION)
-        if "tpot" in goodput_config_dict:
-            valid_metrics.append(all_tpots)
-            slo_values.append(goodput_config_dict["tpot"] / MILLISECONDS_TO_SECONDS_CONVERSION)
-        if "e2el" in goodput_config_dict:
-            valid_metrics.append(e2els)
-            slo_values.append(goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION)
-
-        for req_metric in zip(*valid_metrics):
-            is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)])
-            if is_good_req:
-                good_completed += 1
-
-    if completed == 0:
-        warnings.warn(
-            "All requests failed. This is likely due to a misconfiguration on the benchmark arguments.",
-            stacklevel=2,
-        )
-
-    # Calculate max output tokens per second metric
-    max_output_tokens_per_s = 0.0
-    max_concurrent_requests = 0
-
-    # Find the time range across all successful requests
-    successful_outputs = [output for output in outputs if output.success]
-    failed_outputs = [output for output in outputs if not output.success]
-    if successful_outputs:
-        min_start_time = min(output.start_time for output in successful_outputs)
-        max_end_time = max(output.start_time + output.latency for output in successful_outputs)
-
-        # Create second buckets (ceiling to ensure we capture all time)
-        duration_seconds = int(np.ceil(max_end_time - min_start_time)) + 1
-        tokens_per_second = np.zeros(duration_seconds)
-        concurrent_requests_per_second = np.zeros(duration_seconds)
-
-        for i, output in enumerate(successful_outputs):
-            # Calculate token generation timestamp using
-            # start_time, ttft, and itl
-            token_times = [output.start_time + output.ttft]
-            current_time = token_times[0]
-            for itl_value in output.itl:
-                current_time += itl_value
-                token_times.append(current_time)
-
-            # Add tokens to second buckets
-            for token_time in token_times:
-                second_bucket = int(token_time - min_start_time)
-                if 0 <= second_bucket < duration_seconds:
-                    tokens_per_second[second_bucket] += 1
-
-            # Track concurrent requests for each second this request was active
-            request_start_second = int(output.start_time - min_start_time)
-            request_end_second = int((output.start_time + output.latency) - min_start_time)
-
-            for second in range(request_start_second, request_end_second + 1):
-                concurrent_requests_per_second[second] += 1
-
-        # Find the maximum tokens per second and corresponding
-        # concurrent requests
-        if len(tokens_per_second) > 0:
-            max_output_tokens_per_s = float(np.max(tokens_per_second))
-            max_concurrent_requests = int(np.max(concurrent_requests_per_second))
-
-        if TERM_PLOTLIB_AVAILABLE:
-            import termplotlib as tpl
-
-            fig = tpl.figure()
-            fig.plot(
-                np.arange(len(tokens_per_second)),
-                tokens_per_second,
-                title="Output tokens per second",
-            )
-            fig.plot(
-                np.arange(len(concurrent_requests_per_second)),
-                concurrent_requests_per_second,
-                title="Concurrent requests per second",
-            )
-            fig.show()
-        else:
-            print("tip: install termplotlib and gnuplot to plot the metrics")
-
-    metrics = MultiModalsBenchmarkMetrics(
-        completed=completed,
-        failed=len(failed_outputs),
-        total_input=total_input,
-        total_output=sum(actual_output_lens),
-        request_throughput=completed / dur_s,
-        request_goodput=good_completed / dur_s,
-        output_throughput=sum(actual_output_lens) / dur_s,
-        total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s,
-        mean_ttft_ms=np.mean(ttfts or 0) * 1000,  # ttfts is empty if streaming is not supported by the endpoint
-        std_ttft_ms=np.std(ttfts or 0) * 1000,
-        median_ttft_ms=np.median(ttfts or 0) * 1000,
-        percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles],
-        mean_audio_ttfp_ms=np.mean(audio_ttfps or 0) * 1000,
-        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),
-        total_audio_frames=sum(audio_frames),
-        audio_throughput=sum(audio_duration) / dur_s,
-        mean_audio_rtf=np.mean(audio_rtfs or 0),
-        std_audio_rtf=np.std(audio_rtfs or 0),
-        median_audio_rtf=np.median(audio_rtfs or 0),
-        percentiles_audio_rtf=[(p, np.percentile(audio_rtfs or 0, p)) for p in selected_percentiles],
-        mean_tpot_ms=np.mean(tpots or 0) * 1000,
-        std_tpot_ms=np.std(tpots or 0) * 1000,
-        median_tpot_ms=np.median(tpots or 0) * 1000,
-        percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles],
-        mean_itl_ms=np.mean(itls or 0) * 1000,
-        std_itl_ms=np.std(itls or 0) * 1000,
-        median_itl_ms=np.median(itls or 0) * 1000,
-        percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles],
-        mean_e2el_ms=np.mean(e2els or 0) * 1000,
-        std_e2el_ms=np.std(e2els or 0) * 1000,
-        median_e2el_ms=np.median(e2els or 0) * 1000,
-        percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000) for p in selected_percentiles],
-        max_output_tokens_per_s=max_output_tokens_per_s,
-        max_concurrent_requests=max_concurrent_requests,
-    )
-    print_metrics(
-        task_type,
-        selected_percentile_metrics,
-        max_concurrency,
-        request_rate,
-        benchmark_duration,
-        goodput_config_dict,
-        metrics,
-    )
-    return metrics, actual_output_lens

vllm_omni/benchmarks/patch/patch.py

@@ -1,538 +0,0 @@
-import asyncio
-import base64
-import contextlib
-import io
-import json
-import os
-import random
-import sys
-import time
-import traceback
-from collections.abc import Iterable
-from dataclasses import dataclass
-from datetime import datetime
-from typing import Literal
-
-import aiohttp
-from pydub import AudioSegment
-from tqdm.asyncio import tqdm
-from transformers import PreTrainedTokenizerBase
-from vllm.benchmarks import datasets
-from vllm.benchmarks.datasets import SampleRequest
-from vllm.benchmarks.lib.endpoint_request_func import (
-    ASYNC_REQUEST_FUNCS,
-    OPENAI_COMPATIBLE_BACKENDS,
-    RequestFuncInput,
-    RequestFuncOutput,
-    StreamedResponseHandler,
-    _get_chat_content,
-    _update_headers_common,
-    _update_payload_common,
-    _validate_api_url,
-)
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset
-
-get_samples_old = datasets.get_samples
-
-
-def get_samples(args, tokenizer):
-    if args.backend not in ["openai-chat-omni"]:
-        raise ValueError("benchmark is only supported on 'openai-chat-omni' backend.")
-    if args.dataset_name == "random-mm":
-        dataset = OmniRandomMultiModalDataset(random_seed=args.seed, dataset_path=args.dataset_path)
-        input_requests = dataset.sample(
-            tokenizer=tokenizer,
-            num_requests=args.num_prompts,
-            prefix_len=args.random_prefix_len,
-            range_ratio=args.random_range_ratio,
-            input_len=args.random_input_len,
-            output_len=args.random_output_len,
-            base_items_per_request=args.random_mm_base_items_per_request,
-            limit_mm_per_prompt=args.random_mm_limit_mm_per_prompt,
-            num_mm_items_range_ratio=args.random_mm_num_mm_items_range_ratio,
-            bucket_config=args.random_mm_bucket_config,
-            request_id_prefix=args.request_id_prefix,
-            no_oversample=args.no_oversample,
-        )
-        return input_requests
-    else:
-        return get_samples_old(args, tokenizer)
-
-
-datasets.get_samples = get_samples
-
-
-@dataclass
-class MixRequestFuncOutput(RequestFuncOutput):
-    audio_ttfp: float = 0.0
-    audio_duration: float = 0.0
-    audio_frames: int = 0
-    audio_rtf: float = 0.0
-
-
-async def async_request_openai_chat_omni_completions(
-    request_func_input: RequestFuncInput,
-    session: aiohttp.ClientSession,
-    pbar: tqdm | None = None,
-    mm_position: Literal["first", "last"] = "last",
-) -> MixRequestFuncOutput:
-    api_url = request_func_input.api_url
-    _validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions")
-
-    content = _get_chat_content(request_func_input, mm_position=mm_position)
-
-    payload = {
-        "model": request_func_input.model_name if request_func_input.model_name else request_func_input.model,
-        "messages": [
-            {"role": "user", "content": content},
-        ],
-        "temperature": 0.0,
-        "max_tokens": request_func_input.output_len,
-        "stream": True,
-        "stream_options": {
-            "include_usage": True,
-        },
-    }
-    _update_payload_common(payload, request_func_input)
-
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
-    }
-    _update_headers_common(headers, request_func_input)
-
-    output = MixRequestFuncOutput()
-    output.prompt_len = request_func_input.prompt_len
-
-    generated_text = ""
-    generated_audio = None
-    ttft = 0.0
-    st = time.perf_counter()
-    output.start_time = st
-    most_recent_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:
-                handler = StreamedResponseHandler()
-                async for chunk_bytes in response.content.iter_any():
-                    chunk_bytes = chunk_bytes.strip()
-                    if not chunk_bytes:
-                        continue
-
-                    messages = handler.add_chunk(chunk_bytes)
-                    for message in messages:
-                        # NOTE: SSE comments (often used as pings) start with
-                        # a colon. These are not JSON data payload and should
-                        # be skipped.
-                        if message.startswith(":"):
-                            continue
-
-                        chunk = message.removeprefix("data: ")
-
-                        if chunk != "[DONE]":
-                            timestamp = time.perf_counter()
-                            data = json.loads(chunk)
-                            if choices := data.get("choices"):
-                                modality = data.get("modality")
-                                content = choices[0]["delta"].get("content")
-                                if modality == "text":
-                                    # First token
-                                    if ttft == 0.0:
-                                        ttft = timestamp - st
-                                        output.ttft = ttft
-                                    else:
-                                        output.itl.append(timestamp - most_recent_timestamp)
-                                    generated_text += content or ""
-                                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
-                                    if content != "":
-                                        audio_bytes = base64.b64decode(content)
-                                        seg = AudioSegment.from_file(io.BytesIO(audio_bytes))
-                                        if seg is not None:
-                                            if generated_audio is None:
-                                                generated_audio = seg
-                                            else:
-                                                generated_audio = generated_audio + seg
-
-                            elif usage := data.get("usage"):
-                                output.output_tokens = usage.get("completion_tokens")
-                            most_recent_timestamp = timestamp
-
-                output.generated_text = generated_text
-                if generated_audio is not None:
-                    output.audio_duration = len(generated_audio) / 1000.0
-                    frame_width = generated_audio.frame_width
-                    if frame_width > 0:
-                        output.audio_frames = len(generated_audio.raw_data) // frame_width
-                    else:
-                        output.audio_frames = 0
-                        logger.warning("Audio frame width is zero")
-                    audio_duration = output.audio_duration
-                    if audio_duration > 0:
-                        output.audio_rtf = audio_generate_time / output.audio_duration
-                    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
-    except Exception:
-        output.success = False
-        exc_info = sys.exc_info()
-        output.error = "".join(traceback.format_exception(*exc_info))
-        logger.error(f"ERROR: send request failed, reason is: {output.error}")
-
-    if pbar:
-        pbar.update(1)
-    return output
-
-
-ASYNC_REQUEST_FUNCS["openai-chat-omni"] = async_request_openai_chat_omni_completions
-if "openai-chat-omni" not in OPENAI_COMPATIBLE_BACKENDS:
-    OPENAI_COMPATIBLE_BACKENDS.append("openai-chat-omni")
-
-# ruff: noqa: E402
-# Prevent import order from causing patch failures
-from vllm.benchmarks import serve
-from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, get_request, wait_for_endpoint
-
-from vllm_omni.benchmarks.metrics.metrics import MultiModalsBenchmarkMetrics, calculate_metrics
-
-# ruff: noqa: E402
-
-benchmark_old = serve.benchmark
-
-
-async def benchmark(
-    task_type: TaskType,
-    endpoint_type: str,
-    api_url: str,
-    base_url: str,
-    model_id: str,
-    model_name: str,
-    tokenizer: PreTrainedTokenizerBase,
-    input_requests: list[SampleRequest],
-    logprobs: int | None,
-    request_rate: float,
-    burstiness: float,
-    disable_tqdm: bool,
-    num_warmups: int,
-    profile: bool,
-    selected_percentile_metrics: list[str],
-    selected_percentiles: list[float],
-    ignore_eos: bool,
-    goodput_config_dict: dict[str, float],
-    max_concurrency: int | None,
-    lora_modules: Iterable[str] | None,
-    extra_headers: dict | None,
-    extra_body: dict | None,
-    ramp_up_strategy: Literal["linear", "exponential"] | None = None,
-    ramp_up_start_rps: int | None = None,
-    ramp_up_end_rps: int | None = None,
-    ready_check_timeout_sec: int = 600,
-):
-    try:
-        request_func = ASYNC_REQUEST_FUNCS[endpoint_type]
-    except KeyError:
-        raise ValueError(f"Unknown backend: {endpoint_type}") from None
-
-    # Reuses connections across requests to reduce TLS handshake overhead.
-    connector = aiohttp.TCPConnector(
-        limit=max_concurrency or 0,
-        limit_per_host=max_concurrency or 0,
-        ttl_dns_cache=300,
-        use_dns_cache=True,
-        keepalive_timeout=60,
-        enable_cleanup_closed=True,
-        force_close=False,
-        ssl=("https://" in api_url),
-    )
-
-    session = aiohttp.ClientSession(
-        connector=connector,
-        trust_env=True,
-        timeout=aiohttp.ClientTimeout(total=6 * 60 * 60),
-    )
-
-    print("Starting initial single prompt test run...")
-    test_prompt, test_prompt_len, test_output_len, test_mm_content = (
-        input_requests[0].prompt,
-        input_requests[0].prompt_len,
-        input_requests[0].expected_output_len,
-        input_requests[0].multi_modal_data,
-    )
-
-    assert (
-        test_mm_content is None
-        or isinstance(test_mm_content, dict)
-        or (isinstance(test_mm_content, list) and all(isinstance(item, dict) for item in test_mm_content))
-    ), "multi_modal_data must be a dict or list[dict]"
-    test_input = RequestFuncInput(
-        model=model_id,
-        model_name=model_name,
-        prompt=test_prompt,
-        api_url=api_url,
-        prompt_len=test_prompt_len,
-        output_len=test_output_len,
-        logprobs=logprobs,
-        multi_modal_content=test_mm_content,
-        ignore_eos=ignore_eos,
-        extra_headers=extra_headers,
-        extra_body=extra_body,
-    )
-
-    if ready_check_timeout_sec > 0:
-        test_output = await wait_for_endpoint(
-            request_func,
-            test_input,
-            session,
-            timeout_seconds=ready_check_timeout_sec,
-        )
-        if not test_output.success:
-            raise ValueError(
-                "Initial test run failed - Please make sure benchmark "
-                "arguments are correctly specified. "
-                f"Error: {test_output.error}"
-            )
-        else:
-            print("Initial test run completed.")
-    else:
-        print("Skipping endpoint ready check.")
-
-    if num_warmups > 0:
-        print(f"Warming up with {num_warmups} requests...")
-        warmup_pbar = None if disable_tqdm else tqdm(total=num_warmups)
-        warmup_semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else contextlib.nullcontext()
-        warmup_tasks = []
-
-        async def warmup_limited_request_func():
-            async with warmup_semaphore:
-                return await request_func(request_func_input=test_input, session=session, pbar=warmup_pbar)
-
-        for _ in range(num_warmups):
-            request_task = asyncio.create_task(warmup_limited_request_func())
-            warmup_tasks.append(request_task)
-        _ = await asyncio.gather(*warmup_tasks)
-
-        if warmup_pbar is not None:
-            warmup_pbar.close()
-        print("Warmup run completed.")
-
-    print("Starting main benchmark run...")
-
-    if lora_modules:
-        # For each input request, choose a LoRA module at random.
-        lora_modules = iter([random.choice(lora_modules) for _ in range(len(input_requests))])
-
-    if profile:
-        print("Starting profiler...")
-        profile_input = RequestFuncInput(
-            model=model_id,
-            model_name=model_name,
-            prompt=test_prompt,
-            api_url=base_url + "/start_profile",
-            prompt_len=test_prompt_len,
-            output_len=test_output_len,
-            logprobs=logprobs,
-            multi_modal_content=test_mm_content,
-            ignore_eos=ignore_eos,
-            extra_headers=extra_headers,
-            extra_body=extra_body,
-        )
-        profile_output = await request_func(request_func_input=profile_input, session=session)
-        if profile_output.success:
-            print("Profiler started")
-
-    distribution = "Poisson process" if burstiness == 1.0 else "Gamma distribution"
-
-    if ramp_up_strategy is not None:
-        print(f"Traffic ramp-up strategy: {ramp_up_strategy}.")
-        print(
-            f"Will increase RPS from {ramp_up_start_rps} to {ramp_up_end_rps} RPS over the duration of the benchmark."
-        )
-    else:
-        print(f"Traffic request rate: {request_rate}")
-
-    print(f"Burstiness factor: {burstiness} ({distribution})")
-    print(f"Maximum request concurrency: {max_concurrency}")
-
-    pbar = None if disable_tqdm else tqdm(total=len(input_requests))
-
-    semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else contextlib.nullcontext()
-
-    async def limited_request_func(request_func_input, session, pbar):
-        async with semaphore:
-            return await request_func(request_func_input=request_func_input, session=session, pbar=pbar)
-
-    benchmark_start_time = time.perf_counter()
-    tasks: list[asyncio.Task] = []
-
-    rps_change_events = []
-    last_int_rps = -1
-    if ramp_up_strategy is not None and ramp_up_start_rps is not None:
-        last_int_rps = ramp_up_start_rps
-        rps_change_events.append(
-            {
-                "rps": last_int_rps,
-                "timestamp": datetime.now().isoformat(),
-            }
-        )
-
-    async for request, current_request_rate in get_request(
-        input_requests,
-        request_rate,
-        burstiness,
-        ramp_up_strategy,
-        ramp_up_start_rps,
-        ramp_up_end_rps,
-    ):
-        if ramp_up_strategy is not None:
-            current_int_rps = int(current_request_rate)
-            if current_int_rps > last_int_rps:
-                timestamp = datetime.now().isoformat()
-                for rps_val in range(last_int_rps + 1, current_int_rps + 1):
-                    rps_change_events.append({"rps": rps_val, "timestamp": timestamp})
-                last_int_rps = current_int_rps
-        prompt, prompt_len, output_len, mm_content, request_id = (
-            request.prompt,
-            request.prompt_len,
-            request.expected_output_len,
-            request.multi_modal_data,
-            request.request_id,
-        )
-        req_model_id, req_model_name = model_id, model_name
-        if lora_modules:
-            req_lora_module = next(lora_modules)
-            req_model_id, req_model_name = req_lora_module, req_lora_module
-
-        request_func_input = RequestFuncInput(
-            model=req_model_id,
-            model_name=req_model_name,
-            prompt=prompt,
-            api_url=api_url,
-            prompt_len=prompt_len,
-            output_len=output_len,
-            logprobs=logprobs,
-            multi_modal_content=mm_content,
-            ignore_eos=ignore_eos,
-            extra_headers=extra_headers,
-            extra_body=extra_body,
-            request_id=request_id,
-        )
-        tasks.append(
-            asyncio.create_task(limited_request_func(request_func_input=request_func_input, session=session, pbar=pbar))
-        )
-    outputs: list[MixRequestFuncOutput] = await asyncio.gather(*tasks)
-
-    if pbar is not None:
-        pbar.close()
-
-    benchmark_duration = time.perf_counter() - benchmark_start_time
-
-    if task_type == TaskType.GENERATION:
-        metrics, actual_output_lens = calculate_metrics(
-            input_requests=input_requests,
-            outputs=outputs,
-            dur_s=benchmark_duration,
-            tokenizer=tokenizer,
-            selected_percentiles=selected_percentiles,
-            goodput_config_dict=goodput_config_dict,
-            task_type=task_type,
-            selected_percentile_metrics=selected_percentile_metrics,
-            max_concurrency=max_concurrency,
-            request_rate=request_rate,
-            benchmark_duration=benchmark_duration,
-        )
-    else:
-        metrics = calculate_metrics_for_embeddings(
-            outputs=outputs,
-            dur_s=benchmark_duration,
-            selected_percentiles=selected_percentiles,
-        )
-        actual_output_lens = 0
-
-    if isinstance(metrics, MultiModalsBenchmarkMetrics):
-        result = {
-            "duration": benchmark_duration,
-            "completed": metrics.completed,
-            "failed": metrics.failed,
-            "total_input_tokens": metrics.total_input,
-            "total_output_tokens": metrics.total_output,
-            "request_throughput": metrics.request_throughput,
-            "request_goodput": metrics.request_goodput if goodput_config_dict else None,
-            "output_throughput": metrics.output_throughput,
-            "total_token_throughput": metrics.total_token_throughput,
-            "input_lens": [output.prompt_len for output in outputs],
-            "output_lens": actual_output_lens,
-            "ttfts": [output.ttft for output in outputs],
-            "itls": [output.itl for output in outputs],
-            "generated_texts": [output.generated_text for output in outputs],
-            "errors": [output.error for output in outputs],
-            "max_output_tokens_per_s": metrics.max_output_tokens_per_s,
-            "max_concurrent_requests": metrics.max_concurrent_requests,
-        }
-    else:
-        result = {
-            "duration": benchmark_duration,
-            "completed": metrics.completed,
-            "total_input_tokens": metrics.total_input,
-            "request_throughput": metrics.request_throughput,
-            "total_token_throughput": metrics.total_token_throughput,
-            "input_lens": [output.prompt_len for output in outputs],
-            "errors": [output.error for output in outputs],
-        }
-
-    if rps_change_events:
-        result["rps_change_events"] = rps_change_events
-
-    def process_one_metric(
-        # E.g., "ttft"
-        metric_attribute_name: str,
-    ):
-        # This function prints and adds statistics of the specified
-        # metric.
-        if metric_attribute_name not in selected_percentile_metrics:
-            return
-        is_audio_rtf = metric_attribute_name == "audio_rtf"
-
-        suffix = "" if is_audio_rtf else "_ms"
-        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
-
-    if task_type == TaskType.GENERATION:
-        for metric in selected_percentile_metrics:
-            process_one_metric(metric)
-    else:
-        process_one_metric("e2el")
-
-    if profile:
-        print("Stopping profiler...")
-        profile_input = RequestFuncInput(
-            model=model_id,
-            prompt=test_prompt,
-            api_url=base_url + "/stop_profile",
-            prompt_len=test_prompt_len,
-            output_len=test_output_len,
-            logprobs=logprobs,
-        )
-        profile_output = await request_func(request_func_input=profile_input, session=session)
-        if profile_output.success:
-            print("Profiler stopped")
-
-    await session.close()
-    return result
-
-
-serve.benchmark = benchmark

vllm_omni/benchmarks/serve.py

@@ -1,9 +0,0 @@
-import argparse
-import asyncio
-from typing import Any
-
-from vllm.benchmarks.serve import main_async
-
-
-def main(args: argparse.Namespace) -> dict[str, Any]:
-    return asyncio.run(main_async(args))

vllm_omni/config/model.py

@@ -110,7 +110,6 @@ class OmniModelConfig(ModelConfig):
         mm_processor_cache_gb: float | None,
         mm_processor_cache_type: MMCacheType | None,
         mm_shm_cache_max_object_size_mb: int | None,
-        mm_encoder_only: bool | None,
         mm_encoder_tp_mode: MMEncoderTPMode | None,
         mm_encoder_attn_backend: AttentionBackendEnum | str | None,
         interleave_mm_strings: bool | None,
@@ -182,15 +181,6 @@ class OmniModelConfig(ModelConfig):
         )
         self.model_arch_config = self.get_model_arch_config()
 
-        if self.convert == "mm_encoder_only":
-            logger.warning_once(
-                "`--convert mm_encoder_only` is deprecated and "
-                "will be removed in v0.15. "
-                "Please use --mm-encoder-only` instead."
-            )
-            mm_encoder_only = True
-            self.convert = "none"
-
         architectures = self.architectures
         registry = self.registry
         is_generative_model = registry.is_text_generation_model(architectures, self)
@@ -272,7 +262,6 @@ class OmniModelConfig(ModelConfig):
                 mm_processor_cache_gb=mm_processor_cache_gb,
                 mm_processor_cache_type=mm_processor_cache_type,
                 mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
-                mm_encoder_only=mm_encoder_only,
                 mm_encoder_tp_mode=mm_encoder_tp_mode,
                 mm_encoder_attn_backend=mm_encoder_attn_backend,
                 interleave_mm_strings=interleave_mm_strings,

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -6,20 +6,20 @@ from dataclasses import asdict, dataclass
 from time import time
 from typing import Any
 
+import numpy as np
 from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.distributed.kv_events import KVEventBatch
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
-from vllm.v1.core.sched.utils import remove_all
+from vllm.v1.core.sched.utils import check_stop, remove_all
 from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
 from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
-from vllm_omni.core.sched.output import OmniSchedulerOutput
 from vllm_omni.distributed.omni_connectors.adapter import get_chunk, put_chunk
 from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
 from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
@@ -185,19 +185,13 @@ class OmniARScheduler(VLLMScheduler):
                 get_chunk(self.omni_connector, scheduler_output)
 
             # Add information about requests needing KV cache transfer
-            finished_reqs = self.get_finished_requests_needing_kv_transfer()
+            scheduler_output.finished_requests_needing_kv_transfer = self.get_finished_requests_needing_kv_transfer()
         except Exception:
             # If anything goes wrong, leave the original output unchanged
             init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
-            finished_reqs = {}
+            scheduler_output.finished_requests_needing_kv_transfer = {}
 
-        # Wrap in omni scheduler output to carry transfer metadata.
-        base_fields = SchedulerOutput.__dataclass_fields__.keys()
-        base_data = {name: getattr(scheduler_output, name) for name in base_fields}
-        return OmniSchedulerOutput(
-            **base_data,
-            finished_requests_needing_kv_transfer=finished_reqs,
-        )
+        return scheduler_output
 
     def update_from_output(
         self,
@@ -284,32 +278,47 @@ class OmniARScheduler(VLLMScheduler):
             pooler_output = pooler_outputs[req_index] if pooler_outputs else None
             kv_transfer_params = None
             status_before_stop = request.status
-            finish_reason = None
-            routed_experts = None
 
             # Check for stop and update request status.
             if new_token_ids:
                 new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
-            elif request.pooling_params and pooler_output is not None:
-                # Pooling stops as soon as there is output.
-                request.status = RequestStatus.FINISHED_STOPPED
-                stopped = True
 
-            # If criteria returns True, it means we must STOP the request.
-            # If criteria returns False, it might have triggered a background
-            # transfer (e.g. prefill finished / special token) but continues decoding.
+            # If criteria returns True, it means we must STOP the request (e.g. legacy behavior, not used now)
+            # If criteria returns False, it might have triggered a background transfer
+            # (e.g. prefill finished / special token) but continues decoding
             if not stopped and self._process_kv_transfer_trigger(request, new_token_ids):
                 stopped = True
 
+            # Stop checking for pooler models.
+            pooler_output = None
+            if pooler_outputs:
+                pooler_output = pooler_outputs[req_index]
+                if request.output_token_ids:
+                    stopped = check_stop(request, self.max_model_len)
+            routed_experts = None
             if stopped:
-                routed_experts = self._get_routed_experts(request)
+                # [Omni] Handle routed experts if enabled
+                if self.vllm_config.model_config.enable_return_routed_experts:
+                    kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)
+                    block_ids = kv_blocks.get_block_ids()[0]
+                    num_tokens = request.num_tokens - 1
 
-                # Capture finish_reason BEFORE _handle_stopped_request, which may
-                # reset the status to WAITING for streaming requests that continue.
-                finish_reason = request.get_finished_reason()
-                finished = self._handle_stopped_request(request)
-                if finished:
-                    kv_transfer_params = self._free_request(request)
+                    # compute slot mapping
+                    block_ids_array = np.array(block_ids, dtype=np.int32)
+                    num_blocks = len(block_ids)
+                    block_size = self.block_size
+
+                    # generate block offsets
+                    block_offsets = np.arange(0, block_size)
+
+                    # compute slot mapping: slot = block_id * block_size + offset
+                    slot_mapping = (
+                        block_offsets.reshape((1, block_size)) + block_ids_array.reshape((num_blocks, 1)) * block_size
+                    ).flatten()[:num_tokens]
+
+                    routed_experts = self.routed_experts_reader.get_routed_experts(indices=slot_mapping)
+
+                kv_transfer_params = self._free_request(request)
                 if status_before_stop == RequestStatus.RUNNING:
                     stopped_running_reqs.add(request)
                 else:
@@ -336,13 +345,13 @@ class OmniARScheduler(VLLMScheduler):
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
+            if new_token_ids or pooler_output is not None or kv_transfer_params:
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
                     EngineCoreOutput(
                         request_id=req_id,
                         new_token_ids=new_token_ids,
-                        finish_reason=finish_reason,
+                        finish_reason=request.get_finished_reason(),
                         new_logprobs=new_logprobs,
                         new_prompt_logprobs_tensors=prompt_logprobs_tensors,
                         pooling_output=pooler_output,

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -1,9 +1,7 @@
 import time
 from collections import defaultdict
 
-from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.distributed.kv_events import KVEventBatch
-from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.output import SchedulerOutput
@@ -11,7 +9,6 @@ from vllm.v1.core.sched.request_queue import create_request_queue
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
-from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.request import Request, RequestStatus
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
@@ -59,21 +56,10 @@ class OmniGenerationScheduler(VLLMScheduler):
         # Temporary queue: preserve waiting order, do not disturb non-diffusion requests
         skipped_waiting_requests = create_request_queue(self.policy)
         req_index = 0
-        # OMNI: Track requests that are already finished (e.g., marked by connector)
-        # These should be removed from running and not scheduled
-        already_finished_reqs: set[Request] = set()
         while req_index < len(self.running) and token_budget > 0:
             request = self.running[req_index]
             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:
-                already_finished_reqs.add(request)
-                req_index += 1
-                continue
-
             num_computed_tokens = request.num_computed_tokens
             required_tokens = max(len(request.prompt_token_ids) - num_computed_tokens, 1)
             num_new_tokens = min(required_tokens, token_budget)
@@ -97,24 +83,12 @@ class OmniGenerationScheduler(VLLMScheduler):
             scheduled_running_reqs.append(request)
             req_index += 1
 
-        # OMNI: Remove already finished requests from running queue
-        if already_finished_reqs:
-            self.running = remove_all(self.running, already_finished_reqs)
-
         # Fast path selection and scheduling (treat all as diffusion requests,
         # independent of pooling_params)
         while self.waiting and token_budget > 0 and len(self.running) < self.max_num_running_reqs:
             request = self.waiting.peek_request()
             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:
-                # Pop the finished request from waiting queue and don't schedule it
-                self.waiting.pop_request()
-                continue
-
             # Uniformly treat as diffusion. A feature flag can be added later
             # via config or request tag.
 
@@ -283,16 +257,9 @@ class OmniGenerationScheduler(VLLMScheduler):
         num_nans_in_logits = model_runner_output.num_nans_in_logits
         kv_connector_output = model_runner_output.kv_connector_output
 
-        cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats
-        perf_stats: PerfStats | None = None
-        if self.perf_metrics and self.perf_metrics.is_enabled():
-            perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output)
-
         outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
         spec_decoding_stats: SpecDecodingStats | None = None
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
+        kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
         # Merge connector-side stats (align with v0.14.0)
         if kv_connector_stats and self.connector:
             kv_stats = self.connector.get_kv_connector_stats()
@@ -341,14 +308,12 @@ class OmniGenerationScheduler(VLLMScheduler):
                     num_accepted_tokens=num_accepted,
                 )
 
-            stopped = False
             new_logprobs = None
             new_token_ids = generated_token_ids
             kv_transfer_params = None
-            pooler_output = pooler_outputs[req_index] if pooler_outputs else None
-            status_before_stop = request.status
-            finish_reason = None
-            routed_experts = None
+            pooler_output = None
+            if pooler_outputs:
+                pooler_output = pooler_outputs[req_index]
 
             # Diffusion request: completes in one step; mark finished and free resources
             if request.status == RequestStatus.FINISHED_STOPPED or (
@@ -358,18 +323,8 @@ class OmniGenerationScheduler(VLLMScheduler):
                 # Optional: set a stop_reason for front-end clarity
                 # (does not affect protocol)
                 request.stop_reason = request.stop_reason  # or "generation_done"
-                stopped = True
-
-            if stopped:
-                routed_experts = self._get_routed_experts(request)
-                finish_reason = request.get_finished_reason()
-                finished = self._handle_stopped_request(request)
-                if finished:
-                    kv_transfer_params = self._free_request(request)
-                if status_before_stop == RequestStatus.RUNNING:
-                    stopped_running_reqs.add(request)
-                else:
-                    stopped_preempted_reqs.add(request)
+                kv_transfer_params = self._free_request(request)
+                stopped_running_reqs.add(request)
 
             # Extract sample logprobs if needed.
             if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
@@ -389,13 +344,13 @@ class OmniGenerationScheduler(VLLMScheduler):
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
+            if new_token_ids or pooler_output is not None or kv_transfer_params:
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
                     EngineCoreOutput(
                         request_id=req_id,
                         new_token_ids=new_token_ids,
-                        finish_reason=finish_reason,
+                        finish_reason=request.get_finished_reason(),
                         new_logprobs=new_logprobs,
                         new_prompt_logprobs_tensors=prompt_logprobs_tensors,
                         pooling_output=pooler_output,
@@ -404,7 +359,6 @@ class OmniGenerationScheduler(VLLMScheduler):
                         kv_transfer_params=kv_transfer_params,
                         trace_headers=request.trace_headers,
                         num_cached_tokens=request.num_cached_tokens,
-                        routed_experts=routed_experts,
                         num_nans_in_logits=request.num_nans_in_logits,
                     )
                 )
@@ -452,7 +406,7 @@ class OmniGenerationScheduler(VLLMScheduler):
                     engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
             finished_req_ids.clear()
 
-        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats)) is not None:
+        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
             # Return stats to only one of the front-ends.
             if (eco := next(iter(engine_core_outputs.values()), None)) is None:
                 # We must return the stats even if there are no request

vllm_omni/core/sched/output.py

@@ -1,6 +1,6 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 
-from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
+from vllm.v1.core.sched.output import CachedRequestData, NewRequestData
 from vllm.v1.request import Request
 
 from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
@@ -66,11 +66,4 @@ class OmniCachedRequestData(CachedRequestData):
         prompt_token_ids: Mapping from request ID to list of prompt token IDs
     """
 
-    prompt_token_ids: dict[str, list[int]]
-
-
-@dataclass
-class OmniSchedulerOutput(SchedulerOutput):
-    """Scheduler output with omni-specific transfer metadata."""
-
-    finished_requests_needing_kv_transfer: dict[str, dict] = field(default_factory=dict)
+    prompt_token_ids: dict[int, list[int]]

vllm_omni/diffusion/attention/backends/abstract.py

@@ -7,18 +7,12 @@ from typing import Generic, TypeVar
 
 import torch
 
-from vllm_omni.platforms import current_omni_platform
-
 
 class AttentionBackend(ABC):
     """Abstract class for diffusion attention backends."""
 
     accept_output_buffer: bool = False
 
-    @classmethod
-    def supports_attention_mask(cls) -> bool:
-        return False
-
     @staticmethod
     @abstractmethod
     def get_name() -> str:
@@ -83,58 +77,12 @@ class AttentionImpl(ABC, Generic[T]):
     ) -> None:
         raise NotImplementedError
 
+    @abstractmethod
     def forward(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
-        attn_metadata: T | None = None,
-    ) -> torch.Tensor:
-        """Dispatch to platform-specific forward implementation."""
-        if current_omni_platform.is_rocm():
-            return self.forward_hip(query, key, value, attn_metadata)
-        elif current_omni_platform.is_cuda():
-            return self.forward_cuda(query, key, value, attn_metadata)
-        elif current_omni_platform.is_npu():
-            return self.forward_npu(query, key, value, attn_metadata)
-        elif current_omni_platform.is_xpu():
-            return self.forward_xpu(query, key, value, attn_metadata)
-        else:
-            raise NotImplementedError(f"No forward implementation for platform: {current_omni_platform}")
-
-    def forward_cuda(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: T | None = None,
+        attn_metadata: T,
     ) -> torch.Tensor:
         raise NotImplementedError
-
-    def forward_npu(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: T | None = None,
-    ) -> torch.Tensor:
-        raise NotImplementedError
-
-    def forward_xpu(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: T | None = None,
-    ) -> torch.Tensor:
-        raise NotImplementedError
-
-    def forward_hip(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: T | None = None,
-    ) -> torch.Tensor:
-        # By default, HIP ops are compatible with CUDA ops.
-        return self.forward_cuda(query, key, value, attn_metadata)

vllm_omni/diffusion/attention/backends/ascend_attn.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from importlib.util import find_spec
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+)
+
+logger = init_logger(__name__)
+
+
+class AscendAttentionBackend(AttentionBackend):
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> list[int]:
+        return [x for x in range(1024)]  # todo
+
+    @staticmethod
+    def get_name() -> str:
+        return "ASCEND"
+
+    @staticmethod
+    def get_impl_cls() -> type["AscendAttentionBackendImpl"]:
+        return AscendAttentionBackendImpl
+
+
+class AscendAttentionBackendImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+
+    def forward_native(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
+        attention_mask = attn_metadata.attn_mask if attn_metadata else None
+
+        output = torch.nn.functional.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=self.causal,
+            scale=self.softmax_scale,
+        )
+        out = output.permute(0, 2, 1, 3)
+        return out
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        if find_spec("mindiesd"):
+            from mindiesd import attention_forward
+
+            attention_mask = attn_metadata.attn_mask if attn_metadata else None
+            output = attention_forward(
+                query,
+                key,
+                value,
+                attn_mask=attention_mask,
+                opt_mode="manual",
+                op_type="fused_attn_score",
+                layout="BNSD",
+            )
+            return output
+        else:
+            return self.forward_native(query, key, value, attn_metadata)

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -9,6 +9,13 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionImpl,
     AttentionMetadata,
 )
+from vllm_omni.diffusion.attention.backends.utils.fa import (
+    _pad_input,
+    _unpad_input,
+    _upad_input,
+    flash_attn_func,
+    flash_attn_varlen_func,
+)
 
 logger = init_logger(__name__)
 
@@ -16,10 +23,6 @@ logger = init_logger(__name__)
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
 
-    @classmethod
-    def supports_attention_mask(cls) -> bool:
-        return True
-
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
         return [64, 96, 128, 192, 256]
@@ -48,32 +51,25 @@ class FlashAttentionImpl(AttentionImpl):
         self.causal = causal
         self.softmax_scale = softmax_scale
 
-    def forward_cuda(
+    def forward(
         self,
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
-        """CUDA/ROCm flash attention implementation."""
-        # Import flash attention functions with fallback chain from utils/fa.py
-        # FA3 (fa3_fwd_interface) -> FA3 (flash_attn_interface) -> FA2 (flash_attn)
-        from vllm_omni.diffusion.attention.backends.utils.fa import (
-            HAS_FLASH_ATTN,
-            _pad_input,
-            _unpad_input,
-            _upad_input,
-            flash_attn_func,
-            flash_attn_varlen_func,
-        )
+        """
+        Flash attention implementation.
 
-        if not HAS_FLASH_ATTN:
-            raise ImportError(
-                "FlashAttentionBackend requires Flash Attention. "
-                "Please install one of: fa3-fwd, flash-attention, or flash-attn. "
-                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
-            )
+        Args:
+            query: (batch_size, seq_len, num_heads, head_dim)
+            key: (batch_size, seq_len, num_heads, head_dim)
+            value: (batch_size, seq_len, num_heads, head_dim)
+            attn_metadata: AttentionMetadata. Attention mask is supported as attn_metadata.attn_mask
 
+        Returns:
+            (batch_size, seq_len, num_heads, head_dim)
+        """
         query_length = query.size(1)
         attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
         #  Contains at least one padding token in the sequence
@@ -102,44 +98,11 @@ class FlashAttentionImpl(AttentionImpl):
             out = _pad_input(out_unpad, indices_q, query.size(0), query_length)
 
         else:
-            out = flash_attn_func(
+            out: torch.Tensor = flash_attn_func(
                 query,
                 key,
                 value,
                 causal=self.causal,
                 softmax_scale=self.softmax_scale,
             )
-            # FA3 may return (out, lse) tuple, FA2 returns just out
-            if isinstance(out, tuple):
-                out = out[0]
         return out
-
-    def forward_npu(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata = None,
-    ) -> torch.Tensor:
-        """NPU attention implementation using mindiesd."""
-        try:
-            from mindiesd import attention_forward
-        except ImportError:
-            raise ImportError(
-                "FlashAttentionBackend NPU implementation requires MindIE-SD. "
-                "Please install MindIE-SD to enable NPU attention support. "
-                "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
-                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
-            )
-
-        attention_mask = attn_metadata.attn_mask if attn_metadata else None
-        output = attention_forward(
-            query,
-            key,
-            value,
-            attn_mask=attention_mask,
-            opt_mode="manual",
-            op_type="fused_attn_score",
-            layout="BNSD",
-        )
-        return output

vllm_omni/diffusion/attention/backends/registry.py

@@ -1,152 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Diffusion attention backend registry.
-
-This module provides an enum-based registry for diffusion attention backends,
-similar to vLLM's AttentionBackendEnum. Each backend registers its class path,
-and platforms can override or extend backends using register_backend().
-"""
-
-from collections.abc import Callable
-from enum import Enum, EnumMeta
-from typing import TYPE_CHECKING
-
-from vllm.logger import init_logger
-from vllm.utils.import_utils import resolve_obj_by_qualname
-
-if TYPE_CHECKING:
-    from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend
-
-logger = init_logger(__name__)
-
-
-class _DiffusionBackendEnumMeta(EnumMeta):
-    """Metaclass for DiffusionAttentionBackendEnum to provide better error messages."""
-
-    def __getitem__(cls, name: str) -> "DiffusionAttentionBackendEnum":
-        """Get backend by name with helpful error messages."""
-        try:
-            return super().__getitem__(name)  # type: ignore[return-value]
-        except KeyError:
-            members = list(cls.__members__.keys())
-            valid_backends = ", ".join(members)
-            raise ValueError(
-                f"Unknown diffusion attention backend: '{name}'. Valid options are: {valid_backends}"
-            ) from None
-
-
-class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta):
-    """Enumeration of all supported diffusion attention backends.
-
-    The enum value is the default class path, but this can be overridden
-    at runtime using register_backend().
-
-    To get the actual backend class (respecting overrides), use:
-        backend.get_class()
-
-    Example:
-        # Get backend class
-        backend = DiffusionAttentionBackendEnum.FLASH_ATTN
-        backend_cls = backend.get_class()
-
-        # Register custom backend
-        @register_diffusion_backend(DiffusionAttentionBackendEnum.CUSTOM)
-        class MyCustomBackend:
-            ...
-    """
-
-    # Common backends (available on most platforms)
-    FLASH_ATTN = "vllm_omni.diffusion.attention.backends.flash_attn.FlashAttentionBackend"
-    TORCH_SDPA = "vllm_omni.diffusion.attention.backends.sdpa.SDPABackend"
-    SAGE_ATTN = "vllm_omni.diffusion.attention.backends.sage_attn.SageAttentionBackend"
-
-    def get_path(self, include_classname: bool = True) -> str:
-        """Get the class path for this backend (respects overrides).
-
-        Returns:
-            The fully qualified class path string
-
-        Raises:
-            ValueError: If backend has empty path and is not registered
-        """
-        path = _DIFFUSION_ATTN_OVERRIDES.get(self, self.value)
-        if not path:
-            raise ValueError(
-                f"Backend {self.name} must be registered before use. "
-                f"Use register_diffusion_backend(DiffusionAttentionBackendEnum.{self.name}, "
-                f"'your.module.YourClass')"
-            )
-        if not include_classname:
-            path = path.rsplit(".", 1)[0]
-        return path
-
-    def get_class(self) -> "type[AttentionBackend]":
-        """Get the backend class (respects overrides).
-
-        Returns:
-            The backend class
-
-        Raises:
-            ImportError: If the backend class cannot be imported
-            ValueError: If backend has empty path and is not registered
-        """
-        return resolve_obj_by_qualname(self.get_path())
-
-    def is_overridden(self) -> bool:
-        """Check if this backend has been overridden.
-
-        Returns:
-            True if the backend has a registered override
-        """
-        return self in _DIFFUSION_ATTN_OVERRIDES
-
-    def clear_override(self) -> None:
-        """Clear any override for this backend, reverting to the default."""
-        _DIFFUSION_ATTN_OVERRIDES.pop(self, None)
-
-
-# Override registry
-_DIFFUSION_ATTN_OVERRIDES: dict[DiffusionAttentionBackendEnum, str] = {}
-
-
-def register_diffusion_backend(
-    backend: DiffusionAttentionBackendEnum,
-    class_path: str | None = None,
-) -> Callable[[type], type]:
-    """Register or override a diffusion backend implementation.
-
-    Args:
-        backend: The DiffusionAttentionBackendEnum member to register
-        class_path: Optional class path. If not provided and used as
-            decorator, will be auto-generated from the class.
-
-    Returns:
-        Decorator function if class_path is None, otherwise a no-op
-
-    Examples:
-        # Override an existing backend
-        @register_diffusion_backend(DiffusionAttentionBackendEnum.FLASH_ATTN)
-        class MyCustomFlashAttn:
-            ...
-
-        # Override an existing backend (e.g., ASCEND_ATTN)
-        @register_diffusion_backend(DiffusionAttentionBackendEnum.ASCEND_ATTN)
-        class CustomAscendAttentionBackend:
-            ...
-
-        # Direct registration
-        register_diffusion_backend(
-            DiffusionAttentionBackendEnum.CUSTOM,
-            "my.module.MyCustomBackend"
-        )
-    """
-
-    def decorator(cls: type) -> type:
-        _DIFFUSION_ATTN_OVERRIDES[backend] = f"{cls.__module__}.{cls.__qualname__}"
-        return cls
-
-    if class_path is not None:
-        _DIFFUSION_ATTN_OVERRIDES[backend] = class_path
-        return lambda x: x
-
-    return decorator

vllm_omni/diffusion/attention/backends/ring/ring_globals.py

@@ -3,53 +3,30 @@
 # Copyright (c) 2024, Jiarui Fang.
 # Adapted from https://github.com/feifeibear/long-context-attention
 
-# test if flash_attn (FA2) is available
+
+# test if flash_attn is available
 try:
     import flash_attn  # noqa: F401
-    from flash_attn.flash_attn_interface import _flash_attn_forward  # noqa: F401
+    from flash_attn.flash_attn_interface import _flash_attn_backward, _flash_attn_forward  # noqa: F401
 
     HAS_FLASH_ATTN = True
-except (ImportError, ModuleNotFoundError):
+except ImportError:
     HAS_FLASH_ATTN = False
 
-# FA3 detection: try multiple sources (forward only, no backward needed for inference)
-# Source 1: flash_attn_interface (from flash-attention source build)
-# Source 2: fa3_fwd_interface (from fa3-fwd PyPI package, supports Ampere/Ada/Hopper)
-# Note: FA3 high-level API may or may not return softmax_lse depending on version.
-#       For Ring Attention which requires LSE, we fall back to low-level API if needed.
-HAS_FA3 = False
-fa3_fwd_func = None  # Low-level forward function (_flash_attn_forward)
-fa3_attn_func = None  # High-level attention function (flash_attn_func)
-
-# Try flash_attn_interface first (from flash-attention source build)
 try:
-    from flash_attn_interface import _flash_attn_forward as fa3_fwd_func  # noqa: F401
-    from flash_attn_interface import flash_attn_func as fa3_attn_func  # noqa: F401
+    from flash_attn_interface import _flash_attn_backward as flash_attn_func_hopper_backward  # noqa: F401
+    from flash_attn_interface import _flash_attn_forward as flash_attn_forward_hopper  # noqa: F401
+    from flash_attn_interface import flash_attn_func as flash3_attn_func  # noqa: F401
 
-    HAS_FA3 = True
-except (ImportError, ModuleNotFoundError):
-    pass
-
-# Fallback: try fa3_fwd_interface (PyPI package, supports Ampere/Ada/Hopper)
-if not HAS_FA3:
-    try:
-        from fa3_fwd_interface import _flash_attn_forward as fa3_fwd_func  # noqa: F401
-        from fa3_fwd_interface import flash_attn_func as fa3_attn_func  # noqa: F401
-
-        HAS_FA3 = True
-    except (ImportError, ModuleNotFoundError):
-        pass
-
-# Legacy aliases for backward compatibility
-HAS_FLASH_ATTN_HOPPER = HAS_FA3
-flash_attn_forward_hopper = fa3_fwd_func
-flash3_attn_func = fa3_attn_func
+    HAS_FLASH_ATTN_HOPPER = True
+except ImportError:
+    HAS_FLASH_ATTN_HOPPER = False
 
 try:
     from flashinfer.prefill import single_prefill_with_kv_cache  # noqa: F401
 
     HAS_FLASHINFER = True
-except (ImportError, ModuleNotFoundError):
+except ImportError:
     HAS_FLASHINFER = False
 
 try:
@@ -57,26 +34,26 @@ try:
     from aiter import flash_attn_func as flash_attn_func_aiter  # noqa: F401
 
     HAS_AITER = True
-except (ImportError, ModuleNotFoundError):
+except ImportError:
     HAS_AITER = False
 
 try:
     import sageattention  # noqa: F401
 
     HAS_SAGE_ATTENTION = True
-except (ImportError, ModuleNotFoundError):
+except ImportError:
     HAS_SAGE_ATTENTION = False
 
 try:
     import spas_sage_attn  # noqa: F401
 
     HAS_SPARSE_SAGE_ATTENTION = True
-except (ImportError, ModuleNotFoundError):
+except ImportError:
     HAS_SPARSE_SAGE_ATTENTION = False
 
 try:
     import torch_npu  # noqa: F401
 
     HAS_NPU = True
-except (ImportError, ModuleNotFoundError):
+except ImportError:
     HAS_NPU = False

vllm_omni/diffusion/attention/backends/ring/ring_kernels.py

@@ -6,13 +6,7 @@ import math
 
 import torch
 
-from .ring_globals import (
-    HAS_AITER,
-    HAS_FA3,
-    HAS_FLASH_ATTN,
-    HAS_FLASHINFER,
-    fa3_fwd_func,
-)
+from .ring_globals import HAS_AITER, HAS_FLASH_ATTN, HAS_FLASH_ATTN_HOPPER, HAS_FLASHINFER, HAS_NPU
 
 _scaled_dot_product_flash_attention = torch.ops.aten._scaled_dot_product_flash_attention
 _scaled_dot_product_efficient_attention = torch.ops.aten._scaled_dot_product_efficient_attention
@@ -32,11 +26,23 @@ if HAS_FLASH_ATTN:
     import flash_attn
     from flash_attn.flash_attn_interface import _flash_attn_forward
 
+if HAS_FLASH_ATTN_HOPPER:
+    from flash_attn_interface import _flash_attn_backward as flash_attn_func_hopper_backward
+    from flash_attn_interface import _flash_attn_forward as flash_attn_forward_hopper
+    from flash_attn_interface import flash_attn_func as flash3_attn_func
+else:
+    flash_attn_forward_hopper = None
+    flash_attn_func_hopper_backward = None
+    flash3_attn_func = None
+
 if HAS_FLASHINFER:
     from flashinfer.prefill import single_prefill_with_kv_cache
 
     _LOG2_E = math.log2(math.e)
 
+if HAS_NPU:
+    import torch_npu
+
 
 def pytorch_attn_forward(
     q: torch.Tensor,
@@ -140,35 +146,52 @@ def flash_attn_forward(
     return block_out, block_lse
 
 
-def fa3_forward(q, k, v, dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes, return_softmax):
-    """FA3 forward pass for inference.
+def flash_attn3_func_forward(
+    q, k, v, dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes, return_softmax
+):
+    assert HAS_FLASH_ATTN_HOPPER
+    # current signature of flash_attn_forward_hopper:
+    # (q, k, v, softmax_scale, causal, window_size, descale_q=None, descale_k=None, descale_v=None, gqa_parallel=False)
 
-    FA3 supports Ampere, Ada, and Hopper GPUs. Dropout is ignored since FA3 is inference-only.
-    Uses low-level API (_flash_attn_forward) which always returns softmax_lse,
-    required for Ring Attention's correct accumulation.
-    """
-    assert HAS_FA3, "FA3 is not available"
-    assert fa3_fwd_func is not None, "FA3 low-level API (fa3_fwd_func) not available"
-
-    # Low-level API always returns (out, softmax_lse, S_dmask, rng_state)
-    out, softmax_lse, *_ = fa3_fwd_func(
-        q,
-        k,
-        v,
+    out, softmax_lse, *unused = flash_attn_forward_hopper(
+        q=q,
+        k=k,
+        v=v,
+        k_new=None,
+        v_new=None,
+        qv=None,
+        out=None,
+        cu_seqlens_q=None,
+        cu_seqlens_k=None,
+        cu_seqlens_k_new=None,
+        seqused_q=None,
+        seqused_k=None,
+        max_seqlen_q=None,
+        max_seqlen_k=None,
+        page_table=None,
+        kv_batch_idx=None,
+        leftpad_k=None,
+        rotary_cos=None,
+        rotary_sin=None,
+        seqlens_rotary=None,
+        q_descale=None,
+        k_descale=None,
+        v_descale=None,
         softmax_scale=softmax_scale,
-        causal=causal,
-        window_size_left=window_size[0] if window_size else -1,
-        window_size_right=window_size[1] if window_size else -1,
-        softcap=softcap if softcap else 0.0,
+        causal=False,
+        window_size=(-1, -1),
+        attention_chunk=0,
+        softcap=0.0,
+        rotary_interleaved=True,
+        scheduler_metadata=None,
+        num_splits=0,
+        pack_gqa=None,
+        sm_margin=0,
     )
 
     return out, softmax_lse
 
 
-# Legacy alias for backward compatibility
-flash_attn3_func_forward = fa3_forward
-
-
 def flash_attn_forward_aiter(
     q,
     k,
@@ -241,3 +264,20 @@ def flashinfer_attn_forward(
         raise ValueError(f"Invalid input shape: {q.shape}")
     lse = lse / _LOG2_E
     return out, lse
+
+
+def npu_attn_forward(q, k, v, softmax_scale=None, layout="BSND"):
+    assert HAS_NPU, "torch_npu is not available"
+    softmax_scale = q.shape[-1] ** (-0.5)
+    block_out, block_lse = torch_npu.npu_fused_infer_attention_score(
+        q,
+        k,
+        v,
+        num_heads=q.shape[-2],
+        input_layout=layout,
+        scale=softmax_scale,
+        softmax_lse_flag=True,
+        pre_tokens=65535,
+        next_tokens=65535,
+    )
+    return block_out, block_lse.squeeze(dim=-1)

vllm_omni/diffusion/attention/backends/ring/ring_selector.py

@@ -10,6 +10,7 @@ from functools import partial
 import torch
 
 from .ring_globals import (
+    HAS_NPU,
     HAS_SAGE_ATTENTION,
     HAS_SPARSE_SAGE_ATTENTION,
 )
@@ -27,6 +28,9 @@ if HAS_SAGE_ATTENTION:
 if HAS_SPARSE_SAGE_ATTENTION:
     from spas_sage_attn.autotune import SparseAttentionMeansim
 
+if HAS_NPU:
+    from torch_npu import npu_fused_infer_attention_score
+
 
 class AttnType(Enum):
     AITER = "aiter"
@@ -40,6 +44,7 @@ class AttnType(Enum):
     SAGE_FP8 = "sage_fp8"
     SAGE_FP8_SM90 = "sage_fp8_sm90"
     SPARSE_SAGE = "sparse_sage"
+    NPU = "npu"
 
     @classmethod
     def from_string(cls, s: str):
@@ -152,6 +157,11 @@ def select_flash_attn_impl(
 
         return fn
 
+    elif impl_type == AttnType.NPU:
+        if not HAS_NPU:
+            raise ImportError("torch_npu is not available!")
+        return npu_fused_infer_attention_score
+
     elif attn_processor is not None:
         return attn_processor