[Debug] Correct Unreasonable Long Timeout (#1175 )

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
[Debug] Clear Dockerfile.ci to accelerate build image (#1172 )
2026-02-03 08:49:51 +00:00 · 2026-02-03 13:43:45 +08:00 · 2026-02-03 13:19:36 +08:00 · 2026-02-03 12:24:45 +08:00 · 2026-02-03 12:23:15 +08:00 · 2026-02-02 20:36:56 +08:00
335 changed files with 24324 additions and 9182 deletions
@@ -18,11 +18,25 @@ steps:
  #     queue: "ascend"

  - label: "Simple Unit Test"
-    depends_on: ~
+    depends_on: image-build
    commands:
-      - ".buildkite/scripts/simple_test.sh"
+      - 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
    agents:
-      queue: "cpu_queue_premerge"
+      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"

  - label: "Diffusion Model Test"
    timeout_in_minutes: 20
@@ -41,11 +55,29 @@ steps:
          volumes:
            - "/fsx/hf_cache:/fsx/hf_cache"

+  - label: "Diffusion Images API LoRA E2E"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 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: "Diffusion Model CPU offloading Test"
    timeout_in_minutes: 20
    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:
@@ -132,7 +164,7 @@ steps:
    timeout_in_minutes: 20
    depends_on: image-build
    commands:
-      - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
+      - pytest -s -v tests/diffusion/test_diffusion_worker.py
    agents:
      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
    plugins:
@@ -146,31 +178,13 @@ steps:
          volumes:
            - "/fsx/hf_cache:/fsx/hf_cache"

-  - label: "Omni Model Test"
+
+  - label: "Benchmark 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
-      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_async_omni.py
+      - pytest -s -v tests/benchmarks/test_serve_cli.py
    agents:
      queue: "mithril-h100-pool"
    plugins:
@@ -200,12 +214,69 @@ 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_i2i_multi_image_input.py
+      - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
    agents:
      queue: "mithril-h100-pool"
    plugins:
@@ -234,3 +305,38 @@ 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
@@ -116,6 +116,9 @@ 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 \
@@ -148,6 +151,9 @@ 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 \
@@ -1,57 +0,0 @@
-#!/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/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
@@ -8,14 +8,32 @@ 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
@@ -26,25 +44,30 @@ 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 Parallelism Test"
-  timeout_in_minutes: 15
+- label: "Diffusion Sequence Parallelism Test"
+  timeout_in_minutes: 20
  agent_pool: mi325_2
  depends_on: amd-build
  mirror_hardwares: [amdproduction]
  grade: Blocking
  commands:
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - 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
@@ -52,9 +75,7 @@ steps:
  mirror_hardwares: [amdproduction]
  grade: Blocking
  commands:
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
+    - pytest -s -v tests/diffusion/test_diffusion_worker.py

 - label: "Omni Model Test Qwen2-5-Omni"
  timeout_in_minutes: 15
@@ -66,12 +87,6 @@ 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"
@@ -83,9 +98,10 @@ steps:
  commands:
    - 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_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py
+    - 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

 - label: "Diffusion Image Edit Test"
  timeout_in_minutes: 15
@@ -97,10 +113,4 @@ 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/online_serving/test_i2i_multi_image_input.py
+    - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
@@ -29,7 +29,10 @@ repos:
        # only for staged files

  - repo: https://github.com/rhysd/actionlint
-    rev: v1.7.9
+    # v1.7.8+ sets `go 1.24.0` in go.mod, which older Go toolchains (and most
+    # current CI images) cannot parse. Pin to v1.7.7 until actionlint fixes the
+    # go.mod directive.
+    rev: v1.7.7
    hooks:
      - id: actionlint
        files: ^\.github/workflows/.*\.ya?ml$
@@ -15,6 +15,8 @@ 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. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
 - [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm), 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.

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

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

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.14.0
+ARG VLLM_BASE_TAG=v0.15.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.14.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.15.0
 FROM ${BASE_IMAGE} AS final

 ARG COMMON_WORKDIR=/app
@@ -17,6 +17,10 @@ RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
 COPY . ${COMMON_WORKDIR}/vllm-omni
 RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"

+# When we are installing onnxruntime-rocm, we need to uninstall the system-installed onnxruntime first.
+# These are the dependencies of Qwen3-TTS.
+RUN uv pip uninstall onnxruntime --system && uv pip install --no-cache-dir onnxruntime-rocm sox --system
+
 RUN ln -sf /usr/bin/python3 /usr/bin/python

 CMD ["/bin/bash"]
@@ -25,4 +29,4 @@ ENTRYPOINT []

 #Set entrypoint for vllm-openai official images
 FROM final AS vllm-openai
-ENTRYPOINT ["vllm", "serve"]
+ENTRYPOINT ["vllm", "serve", "--omni"]
@@ -7,6 +7,7 @@ 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:
@@ -53,8 +53,10 @@ Engine classes for offline and online inference.

 Core scheduling and caching components.

+- [vllm_omni.core.sched.omni_ar_scheduler.KVCacheTransferData][]
 - [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][]
 - [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
+- [vllm_omni.core.sched.output.OmniCachedRequestData][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
@@ -22,3 +22,21 @@ 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.
@@ -0,0 +1,359 @@
+# 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>
@@ -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_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker # The specific worker used
+      worker_type: ar # 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_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      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_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      gpu_memory_utilization: 0.15
      enforce_eager: true
@@ -8,7 +8,7 @@ stage_args:
    engine_args:
      model_stage: thinker
      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      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_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      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_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      gpu_memory_utilization: 0.15
      enforce_eager: true
@@ -44,7 +44,7 @@ vllm_omni/

 # Step-by-step-implementation
 ## Step 1: Model Implementation
-The diffusion pipeline’s implementation follows **HuggingFace Diffusers**, and components that do not need modification can be imported directly.
+The diffusion pipeline’s implementation follows **HuggingFace Diffusers**.
 ### 1.1 Define the Pipeline Class
 Define the pipeline class, e.g., `QwenImageEditPipeline`, and initialize all required submodules, either from HuggingFace `diffusers` or custom implementations. In `QwenImageEditPipeline`, only `QwenImageTransformer2DModel` is re-implemented to support optimizations such as Ulysses-SP. When adding new models in the future, you can either reuse this re-implemented `QwenImageTransformer2DModel` or extend it as needed.

@@ -73,10 +73,75 @@ def get_qwen_image_edit_post_process_func(
 The forward function of `QwenImageEditPipeline` follows the HuggingFace `diffusers` design for the most part. The key differences are:
 + As described in the overview, arguments are passed through `OnniDiffusionRequest`, so we need to get user parameters from it accordingly.
 ```python
-prompt = req.prompt if req.prompt is not None else prompt
+prompt = req.prompt
 ```
 + pre/post-processing are handled by the framework elsewhere, so skip them.

+### 1.4 Replace some ops or layers in DiT component
+
+vLLM-Omni provides a set of optimized operators with better performance and built-in support for parallelism, including attention, rotary embeddings (RoPE), and linear layers.
+
+Below is an example showing how to replace standard Transformer attention and FFN layers with vLLM-Omni implementations:
+
+```python
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+
+class MyAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = Attention()
+        self.to_qkv = QKVParallelLinear()
+        self.to_out = RowParallelLinear()
+        self.rope = RotaryEmbedding(is_neox_style=False)
+
+    def forward(self, hidden_states):
+        qkv, _ = self.to_qkv(hidden_states)
+        q, k, v = qkv.split(...)
+        q, k = self.rope(...)
+        attn_output = self.attn(q, k, v)
+        output = self.to_out(attn_output)
+
+class MyFFN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = ColumnParallelLinear()
+        self.fc2 = RowParallelLinear()
+        self.act = F.gelu
+
+    def forward(self, hidden_states):
+        hidden, _ = self.fc1(hidden_states)
+        hidden = self.act(hidden)
+        output = self.fc2(hidden)
+        return output
+```
+
+In this example:
+
+ Attention uses vLLM-Omni’s optimized attention kernel together with parallel QKV projection and RoPE.
+
+ Linear layers are replaced with column- and row-parallel variants to enable tensor parallelism.
+
+ The FFN follows a standard two-layer structure and can be further optimized (e.g., using fused or merged projections) if needed.
+
+
+### 1.5 Provide a `_repeated_blocks` in DiT model
+`_repeated_blocks` is the small and frequently-repeated block(s) of a model -- typically a transformer layer.
+
+It's used for torch compile optimizations.
+```python
+_repeated_blocks = ["QwenImageTransformerBlock"]
+```
+
+
+### 1.6 (Optional) implement sequence parallelism
+vLLM-Omni has a non-intrusive `_sp_plan` that enable sequence parallel without modifying `forward()` logic.
+You can refer to [How to parallelize a new model](../../user_guide/diffusion/parallelism_acceleration.md)
+
+
+### 1.7 (Optional) integrate with Cache-Dit
+vLLM-Omni supports acceleration via [Cache-Dit](../../user_guide/diffusion/cache_dit_acceleration.md). Most models compatible with Diffusers can use Cache-Dit seamlessly. For new models, you can extend support by modifying`cache_dit_backend.py`
+
 ## Step 2: Extend OmniDiffusionRequest Fields
 User-provided inputs are ultimately passed to the model’s forward method through OmniDiffusionRequest, so we add the required fields here to support the new model.
 ```python
@@ -139,6 +204,165 @@ 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
+
+When submitting a pull request to add support for a new model, please include the following information in the PR description:
+
+ Output verification: provide generation outputs to verify correctness and model behavior.
+
+ Inference speed: provide a comparison with the corresponding implementation in Diffusers.
+
+ Parallelism support: specify the supported parallel sizes and any relevant limitations.
+
+ Cache acceleration: check whether the model can be accelerated using Cache-Dit or not.
+
+
+Providing these details helps reviewers evaluate correctness, performance improvements, and parallel scalability of the new model integration.
+
 # Testing
 For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.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 LLMs** and **Diffusion Models**.
+vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** 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. Start Profiling
+### 2. Profiling Omni-Modality Models

 It is best to limit profiling to one iteration to keep trace files manageable.

@@ -77,12 +77,12 @@ omni_llm.close()

 **Examples**:

-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)
+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)

-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)
+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)


-**For Diffusion Models as a single stage**
+### 3. Profiling diffusion models

 Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.

@@ -131,27 +131,10 @@ 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.

-**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
+### 4. Analyzing Omni Traces

 Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.

@@ -163,4 +146,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/latest/dev/profiling.html)
+**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/)
@@ -1,27 +1,35 @@
 # Disaggregated Inference for Omni-Modality Models

-This guide explains how to configure and use distributed connectors (vllm_omni/distributed/connectors) in vllm-omni for multi-stage pipelines.
+This guide explains how to configure and use distributed connectors
+(`vllm_omni/distributed/omni_connectors`) in vllm-omni for multi-stage pipelines.

-## 1. Overview
+Backend-specific setup lives in separate docs:
+
+- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
+- [MooncakeConnector](omni_connectors/mooncake_connector.md)
+- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
+
+## Overview

 Connectors enable data transfer between pipeline stages (e.g., Thinker -> Talker).
-Currently supported connectors operate in **D2H2D (Device-to-Host-to-Device)** mode:
-1. **SharedMemoryConnector**: Uses system shared memory.
-2. **MooncakeConnector**: Uses [Mooncake](https://github.com/kvcache-ai/Mooncake).
+Current connectors operate in D2H2D (device to host to device) mode.

-*   **SharedMemoryConnector (Default)**: Zero-copy (on host), lowest latency. Best for **single-node** deployments. Auto-configured if no connectors are specified.
-*   **MooncakeConnector**: TCP/RDMA based. Best for **multi-node** distributed deployments. Requires a Mooncake Master service.
+## Connector Choices

-## 2. API Design
+| Use Case | Recommended Connector | Notes |
+| :--- | :--- | :--- |
+| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
+| Multi node (Mooncake) | MooncakeConnector | Requires Mooncake Master + metadata server. |
+| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |

-The connector system is built around the `OmniConnectorBase` abstraction, which decouples data transport from stage logic.
+## Core API

-### Core Interface
+The connector system is built around `OmniConnectorBase`.

 ```python
 class OmniConnectorBase(ABC):
    @abstractmethod
-    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int, Optional[dict]]:
+    def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, Optional[dict]]:
        """
        Store data.
        Returns: (success, serialized_size, metadata)
@@ -29,145 +37,23 @@ class OmniConnectorBase(ABC):
        pass

    @abstractmethod
-    def get(self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
+    def get(self, from_stage: str, to_stage: str, get_key: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
        """
        Retrieve data.
-        Args: metadata - Transport-specific handles returned by put() (e.g., SHM name).
+        Args: metadata - transport-specific handles returned by put() (e.g., SHM name).
        Returns: (object, serialized_size)
        """
        pass
 ```

-### Key Concept: Metadata Passing
-Unlike a pure key-value store, some connectors (like `SharedMemoryConnector`) generate transient resources (e.g., a shared memory block name) during `put()`. This `metadata` **must be passed via the control plane** (e.g., HTTP headers, queue messages) from the producer stage to the consumer stage so `get()` can locate the data.
+### Metadata Passing

-## 3. Backends & Use Cases
+Some connectors (e.g., SharedMemoryConnector) generate transient resources during `put()`.
+This `metadata` must be passed through the control plane so `get()` can locate the data.

-### 3.1 SharedMemoryConnector
-**Best for:** Single-node, high-performance IPC.
+## Configuration Model

-*   **Mechanism:**
-    *   **Small Payloads (< Threshold)**: Data is serialized and passed directly "inline" within the `metadata` dictionary. This avoids the overhead of creating SHM blocks for tiny messages.
-    *   **Large Payloads (>= Threshold)**: Data is written to a named System Shared Memory block. The block name is returned in `metadata`.
-*   **Configuration:**
-    *   `shm_threshold_bytes`: Size in bytes to switch from inline to SHM (default: 64KB).
-
-### 3.2 MooncakeConnector
-**Best for:** Multi-node distributed inference.
-
-*   **Mechanism:** Uses Mooncake's distributed KVCache store.
-    *   **Data Plane**: TCP or RDMA for high-bandwidth transfer.
-    *   **Control Plane**: Uses a centralized Mooncake Master and Metadata Server.
-    *   **Keying**: Deterministic keys based on `request_id/from_stage_to_stage`.
-*   **Requirements**: Requires a running Mooncake Master service.
-
-## 4. Relationship with vLLM
-
-vLLM provides specialized distributed mechanisms for specific artifacts:
-*   **KV Transfer** (`vllm.distributed.kv_transfer`): Optimized for transferring KV caches between prefill and decode instances (using NCCL, Mooncake, etc.).
-*   **EC Transfer** (`vllm.distributed.ec_transfer`): Optimized for sharing encoder embeddings.
-*   **Device Communicators** (`vllm.distributed.device_communicators`): Low-level primitives (NCCL, SHM) for Tensor/Pipeline Parallelism.
-
-`vllm-omni` complements this by introducing a **Generalized Connector Abstraction** (`OmniConnector`) for multimodal pipelines. While vLLM's connectors are artifact-specific, `vllm-omni`:
-
-1.  **Unifies Transport**: Provides a single API (`put`/`get`) to transport *any* stage artifact (Input Embeddings, Hidden States, Audio/Image Tensors, KV Cache, Final Output) between arbitrary pipeline stages (e.g., AudioEncoder -> LLM -> AudioGenerator).
-2.  **Extends Connectivity**: Enables flexible construction of complex DAGs (Directed Acyclic Graphs) where stages can run in the same process, same node, or across nodes, using the most appropriate backend (SHM, Mooncake, etc.) for each edge.
-3.  **Wraps & Adapts**: Can internally utilize vLLM's specialized `kv_transfer` for KV paths while using generic transports (SHM/Mooncake) for other data types, presenting a consistent interface to the application layer.
-
-## 5. Installation (Mooncake)
-
-If using `MooncakeConnector`, install the library first:
-
-```bash
-# For CUDA-enabled systems (Recommended)
-pip install mooncake-transfer-engine
-
-# For non-CUDA systems
-pip install mooncake-transfer-engine-non-cuda
-```
-
-## 6. Using MooncakeConnector
-
-### 6.1 Start Mooncake Master
-
-Start the master service on your primary node:
-
-```bash
-# if you use mooncake SSD storage
-mkdir -p ./mc_storage
-
-mooncake_master \
-  --rpc_port=50051 \
-  --enable_http_metadata_server=true \
-  --http_metadata_server_host=0.0.0.0 \
-  --http_metadata_server_port=8080 \
-  --metrics_port=9003 \
-  --root_fs_dir=./mc_storage/ \
-  --cluster_id=mc-local-1 &
-```
-
-### 6.2 Configuration (YAML)
-
-Edit your stage config (e.g., `qwen2_5_omni.yaml`).
-
-**Step 1: Define Connector in Global Runtime**
-
-```yaml
-runtime:
-  connectors:
-    connector_of_mooncake:
-      name: MooncakeConnector
-      extra:
-        host: "127.0.0.1"           # Local Worker IP
-        metadata_server: "http://<MASTER_IP>:8080/metadata"
-        master: "<MASTER_IP>:50051"
-        segment: 512000000          # 512MB segment
-        localbuf: 64000000          # 64MB buffer
-        proto: "tcp"                # "tcp" or "rdma"
-    ```
-
-**Mooncake Configuration Parameters:**
-
-*   **host**: The hostname or IP address of the local machine (worker). Mooncake uses this to register itself in the metadata server so other nodes can find it.
-*   **metadata_server**: The URL of the metadata server. This is used for service discovery and connection establishment (e.g., exchanging QP information for RDMA).
-*   **master**: The address of the Mooncake Master Server (e.g., `<MASTER_IP>:50051`). This is used for global state management and control plane operations.
-*   **segment**: The size of the global memory segment in bytes (default: ~512MB). This defines the shared memory region accessible by Mooncake for data transfer.
-*   **localbuf**: The size of the local buffer in bytes (default: ~64MB). Used for local data buffering during transfer operations.
-*   **proto**: The transport protocol to use. Options:
-    *   `tcp`: Standard TCP/IP (easier setup, universal compatibility).
-    *   `rdma`: Remote Direct Memory Access (higher performance, requires RDMA-capable hardware).
-
-For more details, refer to the [Mooncake Repository](https://github.com/kvcache-ai/Mooncake).
-
-    **Step 2: Reference in Stages**
-
-Explicitly link stages using `input_connectors` and `output_connectors`:
-
-```yaml
-stage_args:
-  - stage_id: 0
-    # ...
-    output_connectors:
-      to_stage_1: connector_of_mooncake
-
-  - stage_id: 1
-    # ...
-    input_connectors:
-      from_stage_0: connector_of_mooncake
-```
-
-## 7. Using SharedMemoryConnector (Auto-Mode)
-
-**Best for single-node.**
-
-The system will automatically create `SharedMemoryConnector`s for any pipeline edge that does not have an explicit connector defined. This is inferred from:
-1.  `runtime.edges` list in the config.
-2.  `engine_input_source` dependencies defined in `stage_args`.
-
-### Threshold Configuration
-By default, payloads larger than **64KB** (default threshold) are transferred via shared memory, while smaller ones use the control queue (inline).
-
-To adjust this threshold (e.g., to 1GB), add the following to your `runtime.connectors`:
+Define connectors in runtime:

 ```yaml
 runtime:
@@ -175,26 +61,46 @@ runtime:
    connector_of_shared_memory:
      name: SharedMemoryConnector
      extra:
-        shm_threshold_bytes: 1024 # 1KB threshold
+        shm_threshold_bytes: 65536
 ```

-## 8. Summary
+Wire stages to connectors:

-| Use Case | Recommended Connector | Configuration |
-| :--- | :--- | :--- |
-| **Single Node** | `SharedMemoryConnector` | **None** (Automatic) or Custom Threshold |
-| **Multi Node** | `MooncakeConnector` | Explicit YAML + Mooncake Master |
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_shared_memory

-## 9. Operational Notes (important)
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+```

- **Fail-fast config validation**: the loader raises if any expected edge is missing a connector. Define `input_connectors`/`output_connectors` or rely on auto-SHM filling; otherwise startup aborts.
- **Missing payloads halt the stage**: workers expect connector payloads; if metadata or connector config is missing, the stage raises and stops. Verify connector wiring and metadata propagation before production.
+If a pipeline edge has no explicit connector, the system auto-creates a
+SharedMemoryConnector for that edge.

-## 10. Future Roadmap: Device-to-Device (D2D) Transport
+## Relationship with vLLM

-The current implementations (`SharedMemoryConnector`, `MooncakeConnector`) utilize a **D2H2D (Device-to-Host-to-Device)** data path. Tensors are moved to CPU memory (Host) for transport, which incurs PCIe overhead.
+vLLM provides specialized distributed mechanisms for specific artifacts:

-As outlined in the design RFC, future versions will introduce **D2D (Device-to-Device)** connectors:
+- KV Transfer (`vllm.distributed.kv_transfer`): optimized for KV caches.
+- EC Transfer (`vllm.distributed.ec_transfer`): optimized for encoder embeddings.
+- Device Communicators (`vllm.distributed.device_communicators`): low-level primitives (NCCL, SHM).

-*   **Goal**: Direct GPU-to-GPU transfer (via NCCL, UCX, or IPC) to minimize latency for large tensor payloads.
-*   **Mechanism**: The `OmniConnector` API allows `put()` to initiate a transfer and return a lightweight handle (metadata) via the control plane, while the heavy payload flows directly between devices.
+vllm-omni complements this with a generalized connector abstraction:
+
+1. Unifies transport via a single `put`/`get` API for any stage artifact.
+2. Enables DAG-style pipelines across processes or nodes with per-edge transports.
+3. Can wrap vLLM-specific transfers for KV paths while keeping a consistent interface.
+
+## Operational Notes
+
+- Fail-fast config validation: missing expected edges cause startup failures.
+- Missing payloads halt stages: verify connector wiring and metadata propagation.
+
+## Future Roadmap: D2D Transport
+
+Current connectors use D2H2D paths. Future versions will introduce direct
+device-to-device connectors (NCCL, UCX, IPC) to reduce latency for large
+tensor payloads.
@@ -0,0 +1,74 @@
+# MooncakeConnector
+
+## When to Use
+
+Best for multi-node distributed inference using Mooncake.
+
+## Installation
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+## Start Mooncake Master
+
+```bash
+# If you use Mooncake SSD storage
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+## Configuration
+
+Define the connector in runtime:
+
+```yaml
+runtime:
+  connectors:
+    connector_of_mooncake:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://<MASTER_IP>:8080/metadata"
+        master: "<MASTER_IP>:50051"
+        segment: 512000000
+        localbuf: 64000000
+        proto: "tcp"
+```
+
+Wire stages to the connector:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_mooncake
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_mooncake
+```
+
+Parameters:
+
+- host: local worker IP registered in the metadata server.
+- metadata_server: metadata server URL for discovery and setup.
+- master: Mooncake Master address.
+- segment: global memory segment size in bytes.
+- localbuf: local buffer size in bytes.
+- proto: transport protocol ("tcp" or "rdma").
+
+For more details, refer to the
+[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
@@ -0,0 +1,27 @@
+# SharedMemoryConnector
+
+## When to Use
+
+Best for single-node deployments where stages run on the same host. It is
+auto-configured when no explicit connector is specified for an edge.
+
+## How It Works
+
+- Small payloads (< threshold): serialized and passed inline in metadata.
+- Large payloads (>= threshold): stored in shared memory; the SHM name is
+  returned in metadata.
+
+## Configuration
+
+```yaml
+runtime:
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+```
+
+## Notes
+
+- Auto-mode uses SharedMemoryConnector if no connector is declared for an edge.
@@ -0,0 +1,101 @@
+# YuanrongConnector
+
+## When to Use
+
+Best for multi-node distributed inference using Yuanrong Datasystem.
+
+## Mechanism
+
+Uses Yuanrong Datasystem's distributed KV store (`datasystem.kv_client`).
+
+- Data Plane: TCP or RDMA for high-bandwidth transfer.
+- Control Plane: Yuanrong Datasystem workers and etcd.
+- Keying: deterministic keys based on `put_key` (often composed as `request_id:fromStage_toStage`).
+
+## Installation
+
+```bash
+pip install openyuanrong-datasystem
+```
+
+## Start etcd
+
+```bash
+# Download and install etcd (v3.5.12 or higher)
+ETCD_VERSION="v3.5.12"
+ETCD_ARCH="linux-arm64"
+wget https://github.com/etcd-io/etcd/releases/download/${ETCD_VERSION}/etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz
+tar -xvf etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz
+cd etcd-${ETCD_VERSION}-${ETCD_ARCH}
+sudo cp etcd etcdctl /usr/local/bin/
+
+# Start etcd
+etcd \
+  --name etcd-single \
+  --data-dir /tmp/etcd-data \
+  --listen-client-urls http://0.0.0.0:2379 \
+  --advertise-client-urls http://0.0.0.0:2379 \
+  --listen-peer-urls http://0.0.0.0:2380 \
+  --initial-advertise-peer-urls http://0.0.0.0:2380 \
+  --initial-cluster etcd-single=http://0.0.0.0:2380 &
+
+# Verify etcd is running
+etcdctl --endpoints "127.0.0.1:2379" put key "value"
+etcdctl --endpoints "127.0.0.1:2379" get key
+```
+
+For production environments, refer to the
+[official etcd clustering documentation](https://etcd.io/docs/current/op-guide/clustering/).
+
+## Start Datasystem Worker
+
+```bash
+# Replace ${ETCD_IP} with etcd node IP, ${WORKER_IP} with local node IP
+dscli start -w \
+  --worker_address "${WORKER_IP}:31501" \
+  --etcd_address "${ETCD_IP}:2379" \
+  --shared_memory_size_mb 20480
+```
+
+To stop the worker:
+
+```bash
+dscli stop --worker_address "${WORKER_IP}:31501"
+```
+
+## Configuration
+
+Define the connector in runtime:
+
+```yaml
+runtime:
+  connectors:
+    connector_of_yuanrong:
+      name: YuanrongConnector
+      extra:
+        host: "127.0.0.1"
+        port: 31501
+        get_sub_timeout_ms: 1000
+```
+
+Wire stages to the connector:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_yuanrong
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_yuanrong
+```
+
+Parameters:
+
+- host: datasystem worker host.
+- port: datasystem worker port.
+- get_sub_timeout_ms: get timeout in milliseconds (0 for no timeout).
+
+For more details, refer to the
+[Yuanrong Datasystem repository](https://atomgit.com/openeuler/yuanrong-datasystem).
@@ -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 and 0.12.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).
+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).


 vLLM-Omni is built based on vLLM. Please install it with command below.
@@ -23,15 +23,6 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.rocm -t vllm-omni-rocm .
 ```

-If you want to specify which GPU Arch to build for to cutdown build time:
-
-```bash
-DOCKER_BUILDKIT=1 docker build \
-  -f docker/Dockerfile.rocm \
-  --build-arg PYTORCH_ROCM_ARCH="gfx942;gfx950" \
-  -t vllm-omni-rocm .
-```
-
 #### Launch the docker image

 ##### Launch with OpenAI API Server
@@ -46,6 +46,46 @@ if __name__ == "__main__":
    images[0].save("coffee.png")
 ```

+You can pass a list of prompts and wait for them to process altogether, shown below.
+
+!!! info
+
+    However, it is not currently recommended to do so
+    because not all models support batch inference,
+    and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
+    This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(
+        model="Tongyi-MAI/Z-Image-Turbo",
+        # stage_configs_path="./stage-config.yaml",  # See below
+    )
+    prompts = [
+        "a cup of coffee on a table",
+        "a toy dinosaur on a sandy beach",
+        "a fox waking up in bed and yawning",
+    ]
+    omni_outputs = omni.generate(prompts)
+    for i_prompt, prompt_output in enumerate(omni_outputs):
+        this_request_output = prompt_output.request_output[0]
+        this_images = this_request_output.images
+        for i_image, image in enumerate(this_images):
+            image.save(f"p{i_prompt}-img{i_image}.jpg")
+            print("saved to", f"p{i_prompt}-img{i_image}.jpg")
+            # saved to p0-img0.jpg
+            # saved to p1-img0.jpg
+            # saved to p2-img0.jpg
+```
+
+!!! info
+
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
+    batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+
 For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)

 ## Online Serving with OpenAI-Completions API
@@ -33,6 +33,7 @@ th {
 |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
+|`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
 |`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
@@ -0,0 +1,205 @@
+# 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
+```
@@ -18,6 +18,7 @@ Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`. Cac

 ```python
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams

 # Simplest way: just enable cache-dit with default parameters
 omni = Omni(
@@ -27,7 +28,7 @@ omni = Omni(

 images = omni.generate(
    "a beautiful landscape",
-    num_inference_steps=50,
+    OmniDiffusionSamplingParams(num_inference_steps=50),
 )
 ```

@@ -1,19 +1,32 @@
 # 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 Model
+**Execution Flow**:
 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:

- **Python API**: set `enable_cpu_offload=True`.
+1. **Python API**: set `enable_cpu_offload=True`.

 ```python
 from vllm_omni import Omni
@@ -23,7 +36,66 @@ if __name__ == "__main__":
    m = Omni(model="Qwen/Qwen-Image",enable_cpu_offload=True)
 ```

- **CLI**: pass `--enable-cpu-offload` to the diffusion service entrypoint.
+2. **CLI**: pass `--enable-cpu-offload` to the diffusion service entrypoint.

-## Known Limitations
+### 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
@@ -18,18 +18,19 @@ 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` | ❌ | ❌ | ❌ | ❌ |
-
+| 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.
@@ -48,7 +49,7 @@ The following table shows which models are currently supported by parallelism me

 | Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel |
 |-------|------------------|------------|---------|--------------------------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ | ❌ |
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ❌ |

 ### Tensor Parallelism

@@ -66,10 +67,12 @@ omni = Omni(
 )

 outputs = omni.generate(
-    prompt="a cat reading a book",
-    num_inference_steps=9,
-    width=512,
-    height=512,
+    "a cat reading a book",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=9,
+        width=512,
+        height=512,
+    ),
 )
 ```

@@ -82,6 +85,7 @@ outputs = omni.generate(
 An example of offline inference script using [Ulysses-SP](https://arxiv.org/pdf/2309.14509) is shown below:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ulysses_degree = 2

@@ -90,7 +94,10 @@ omni = Omni(
    parallel_config=DiffusionParallelConfig(ulysses_degree=2)
 )

-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```

 See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
@@ -132,6 +139,7 @@ Ring-Attention ([arxiv paper](https://arxiv.org/abs/2310.01889)) splits the inpu
 An example of offline inference script using Ring-Attention is shown below:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ring_degree = 2

@@ -140,7 +148,10 @@ omni = Omni(
    parallel_config=DiffusionParallelConfig(ring_degree=2)
 )

-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```

 See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
@@ -182,6 +193,7 @@ You can combine both Ulysses-SP and Ring-Attention for larger scale parallelism.

 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig

 # Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
@@ -190,7 +202,10 @@ omni = Omni(
    parallel_config=DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)
 )

-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```

 ##### Online Serving
@@ -357,9 +372,9 @@ def forward(self, hidden_states, ...):

 ### CFG-Parallel

-##### Offline Inference
+#### Offline Inference

-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).
+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.

 An example of offline inference using CFG-Parallel (image-to-image) is shown below:

@@ -367,23 +382,50 @@ 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",
-    true_cfg_scale=4.0,
-    pil_image=input_image,
-    num_inference_steps=50,
+    {
+        "prompt": "turn this cat to a dog",
+        "negative_prompt": "low quality, blurry",
+        "multi_modal_data": {"image": input_image},
+    },
+    OmniDiffusionSamplingParams(
+        true_cfg_scale=4.0,
+        num_inference_steps=50,
+    ),
 )
 ```

 Notes:

- CFG-Parallel is only effective when **true CFG** is enabled (i.e., `true_cfg_scale > 1` and a `negative_prompt` is provided).
+- 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
+```

 #### How to parallelize a pipeline

@@ -396,58 +438,130 @@ 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.

-Below is an example of CFG-Parallel implementation:
+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:**

 ```python
-def diffuse(
+class QwenImageCFGParallelMixin(CFGParallelMixin):
+    """
+    Base Mixin class for Qwen Image pipelines providing shared CFG methods.
+    """
+
+    def diffuse(
        self,
-        ...
-        ):
-    # Enable CFG-parallel: rank0 computes positive, rank1 computes negative.
-    cfg_parallel_ready = do_true_cfg and get_classifier_free_guidance_world_size() > 1
+        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

-    self.transformer.do_true_cfg = do_true_cfg
+        for i, t in enumerate(timesteps):
+            timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)

-    if cfg_parallel_ready:
-        cfg_group = get_cfg_group()
-        cfg_rank = get_classifier_free_guidance_rank()
+            # 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_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]
+            # 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

-        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
-        ...
+            # 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
 ```
@@ -8,6 +8,7 @@ Enable TeaCache by setting `cache_backend` to `"tea_cache"`:

 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams

 # Simple configuration - model_type is automatically extracted from pipeline.__class__.__name__
 omni = Omni(
@@ -17,7 +18,12 @@ omni = Omni(
        "rel_l1_thresh": 0.2  # Optional, defaults to 0.2
    }
 )
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```

 ### Using Environment Variable
@@ -68,13 +74,19 @@ Controls the balance between speed and quality. Lower values prioritize quality,

 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams

 omni = Omni(
    model="Qwen/Qwen-Image",
    cache_backend="tea_cache",
    cache_config={"rel_l1_thresh": 0.2}
 )
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```

 ## Performance Tuning
@@ -112,6 +124,7 @@ th {
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
+| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |

 ### VideoGen

@@ -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` | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **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 |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ❌ | ❌ |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention |CFG-Parallel |
+|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:----------------:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ |


 ## Performance Benchmarks
@@ -96,6 +96,7 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m

 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams

 omni = Omni(
    model="Qwen/Qwen-Image",
@@ -103,13 +104,19 @@ omni = Omni(
    cache_config={"rel_l1_thresh": 0.2}  # Optional, defaults to 0.2
 )

-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```

 ### Using Cache-DiT

 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams

 omni = Omni(
    model="Qwen/Qwen-Image",
@@ -123,7 +130,12 @@ omni = Omni(
    }
 )

-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```

 ### Using Ulysses-SP
@@ -131,6 +143,7 @@ outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_st
 Run text-to-image:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ulysses_degree = 2

@@ -139,13 +152,17 @@ omni = Omni(
    parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
 )

-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```


 Run image-to-image:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ulysses_degree = 2

@@ -154,8 +171,13 @@ omni = Omni(
    parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
 )

-outputs = omni.generate(prompt="turn this cat to a dog",
-        pil_image=input_image, num_inference_steps=50)
+outputs = omni.generate(
+    {
+        "prompt": "turn this cat to a dog",
+        "multi_modal_data": {"image": input_image}
+    },
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
 ```

 ### Using Ring-Attention
@@ -163,6 +185,7 @@ outputs = omni.generate(prompt="turn this cat to a dog",
 Run text-to-image:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ring_degree = 2

@@ -171,7 +194,10 @@ omni = Omni(
    parallel_config=DiffusionParallelConfig(ring_degree=2)
 )

-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```

 ### Using CFG-Parallel
@@ -182,6 +208,7 @@ CFG-Parallel splits the CFG positive/negative branches across GPUs. Use it when

 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 cfg_parallel_size = 2

@@ -190,8 +217,13 @@ omni = Omni(
    parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size)
 )

-outputs = omni.generate(prompt="turn this cat to a dog",
-        pil_image=input_image, num_inference_steps=50, true_cfg_scale=4.0)
+outputs = omni.generate(
+    {
+        "prompt": "turn this cat to a dog",
+        "multi_modal_data": {"image": input_image}
+    },
+    OmniDiffusionSamplingParams(num_inference_steps=50, true_cfg_scale=4.0),
+)
 ```

 ## Documentation
@@ -0,0 +1,179 @@
+# 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**       |
@@ -47,6 +47,7 @@ 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.
@@ -52,6 +52,7 @@ 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).
@@ -0,0 +1,107 @@
+# LoRA-Inference
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/lora_inference>.
+
+This contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
+The example uses the  `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.
+
+## Overview
+
+Similar to vLLM, vLLM-omni uses a unified LoRA handling mechanism:
+
+- **Pre-loaded LoRA**: Loaded at initialization via `--lora-path` (pre-loaded into cache)
+- **Per-request LoRA**: Loaded on-demand. In the example, the LoRA is loaded via `--lora-request-path` in each request
+
+Both approaches use the same underlying mechanism - all LoRA adapters are handled uniformly through `set_active_adapter()`. If no LoRA request is provided in a request, all adapters are deactivated.
+
+## Usage
+
+### Pre-loaded LoRA (via --lora-path)
+
+Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_preloaded.png
+```
+
+**Note**: When using `--lora-path`, the adapter is loaded at init time with a stable ID derived from the adapter path. This example activates it automatically for the request.
+
+### Per-request LoRA (via --lora-request-path)
+
+Load a LoRA adapter on-demand for each request:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-request-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_per_request.png
+```
+
+### No LoRA
+
+If no LoRA request is provided, we will use the base model without any LoRA adapters:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_no_lora.png
+```
+
+## Parameters
+
+### LoRA Parameters
+
+- `--lora-path`: Path to LoRA adapter folder to pre-load at initialization (loads into cache with a stable ID derived from the path)
+- `--lora-request-path`: Path to LoRA adapter folder for per-request loading
+- `--lora-request-id`: Integer ID for the LoRA adapter (optional). If not provided and `--lora-request-path` is set, will derive a stable ID from the path.
+- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
+
+### Standard Parameters
+
+- `--prompt`: Text prompt for image generation (required)
+- `--seed`: Random seed for reproducibility (default: 42)
+- `--height`: Image height in pixels (default: 1024)
+- `--width`: Image width in pixels (default: 1024)
+- `--num_inference_steps`: Number of denoising steps (default: 50)
+- `--output`: Output file path (default: `lora_output.png`)
+
+## How LoRA Works
+
+All LoRA adapters are handled uniformly:
+
+1. **Initialization**: If `--lora-path` is provided, the adapter is loaded into cache with a stable ID derived from the adapter path
+2. **Per-request**: If `--lora-request-path` is provided, the adapter is loaded/activated for that request
+3. **No LoRA**: If no LoRA request is provided (`req.lora_request` is None), all adapters are deactivated
+
+The system uses LRU cache management - adapters are cached and evicted when the cache is full (unless pinned).
+
+## LoRA Adapter Format
+
+LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
+
+## Example materials
+
+??? abstract "lora_inference.py"
+    ``````py
+    --8<-- "examples/offline_inference/lora_inference/lora_inference.py"
+    ``````
@@ -23,7 +23,7 @@ if __name__ == "__main__":
    images[0].save("coffee.png")
 ```

-Or put more than one prompt in a request, processing them sequentially.
+Or put more than one prompt in a request.

 ```python
 from vllm_omni.entrypoints.omni import Omni
@@ -40,6 +40,41 @@ if __name__ == "__main__":
      image = output.request_output[0].images[0].save(f"{i}.jpg")
 ```

+!!! info
+
+    However, it is not currently recommended to do so
+    because not all models support batch inference,
+    and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
+    This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
+
+!!! info
+
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
+    batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+
+Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
+This is useful for models that support negative prompts.
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(model="Qwen/Qwen-Image")
+    outputs = omni.generate([
+      {
+        "prompt": "a cup of coffee on a table"，
+        "negative_prompt": "low resolution"
+      },
+      {
+        "prompt": "a toy dinosaur on a sandy beach"，
+        "negative_prompt": "cinematic, realistic"
+      }
+    ])
+    for i, output in enumerate(outputs):
+      image = output.request_output[0].images[0].save(f"{i}.jpg")
+```
+
 ## Local CLI Usage

 ```bash
@@ -60,6 +95,7 @@ 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).
@@ -28,6 +28,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).
+- `--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.
@@ -0,0 +1,232 @@
+# 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**       |
@@ -0,0 +1,69 @@
+# LoRA-Inference
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/lora_inference>.
+
+This example shows how to use **per-request LoRA** with vLLM-Omni diffusion models via the OpenAI-compatible Chat Completions API.
+
+> Note: The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
+> Note: This example uses `/v1/chat/completions`. LoRA payloads for other OpenAI endpoints are not implemented here.
+
+## Start Server
+
+```bash
+# Pick a diffusion model (examples)
+# export MODEL=stabilityai/stable-diffusion-3.5-medium
+# export MODEL=Qwen/Qwen-Image
+
+bash run_server.sh
+```
+
+## Call API (curl)
+
+```bash
+# Required: local LoRA folder on the server
+export LORA_PATH=/path/to/lora_adapter
+
+# Optional
+export SERVER=http://localhost:8091
+export PROMPT="A piece of cheesecake"
+export LORA_NAME=my_lora
+export LORA_SCALE=1.0
+# Optional: if omitted, the server derives a stable id from LORA_PATH.
+# export LORA_INT_ID=123
+
+bash run_curl_lora_inference.sh
+```
+
+## Call API (Python)
+
+```bash
+python openai_chat_client.py \
+  --prompt "A piece of cheesecake" \
+  --lora-path /path/to/lora_adapter \
+  --lora-name my_lora \
+  --lora-scale 1.0 \
+  --output output.png
+```
+
+## LoRA Format
+
+LoRA adapters should be in PEFT format, for example:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
+
+??? abstract "openai_chat_client.py"
+    ``````py
+    --8<-- "examples/online_serving/lora_inference/openai_chat_client.py"
+    ``````
+??? abstract "run_curl_lora_inference.sh"
+    ``````py
+    --8<-- "examples/online_serving/lora_inference/run_curl_lora_inference.sh"
+    ``````
+??? abstract "run_server.sh"
+    ``````py
+    --8<-- "examples/online_serving/lora_inference/run_server.sh"
+    ``````
@@ -15,6 +15,12 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```

+If you want to open async chunking for qwen3-omni, launch the server with command below
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+```
+
 If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
@@ -0,0 +1,177 @@
+# 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**       |
@@ -1,5 +1,8 @@
 import argparse
 import os
+from typing import cast
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType


 def parse_args():
@@ -50,21 +53,24 @@ def parse_args():
 def main():
    args = parse_args()
    model_name = args.model
+    prompts: list[OmniPromptType] = []
    try:
        # Preferred: load from txt file (one prompt per line)
        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
            with open(args.txt_prompts, encoding="utf-8") as f:
                lines = [ln.strip() for ln in f.readlines()]
-            args.prompts = [ln for ln in lines if ln != ""]
-            print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
+            prompts = [ln for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
+        else:
+            prompts = args.prompts
    except Exception as e:
        print(f"[Error] Failed to load prompts: {e}")
        raise

-    if args.prompts is None:
+    if not prompts:
        # Default prompt for text2img test if none provided
-        args.prompts = ["<|im_start|>A cute cat<|im_end|>"]
-        print(f"[Info] No prompts provided, using default: {args.prompts}")
+        prompts = ["<|im_start|>A cute cat<|im_end|>"]
+        print(f"[Info] No prompts provided, using default: {prompts}")
    omni_outputs = []

    from PIL import Image
@@ -77,21 +83,27 @@ def main():
        print("[Info] Running in img2img mode (Stage 1 only)")
        client = OmniDiffusion(model=model_name)

-        generate_kwargs = {
-            "prompt": args.prompts,
-            "seed": 52,
-            "need_kv_receive": False,
-            "num_inference_steps": args.steps,
-        }
-
        if args.image_path:
            if os.path.exists(args.image_path):
                loaded_image = Image.open(args.image_path).convert("RGB")
-                generate_kwargs["pil_image"] = loaded_image
+                prompts = [
+                    {
+                        "prompt": cast(str, p),
+                        "multi_modal_data": {"image": loaded_image},
+                    }
+                    for p in prompts
+                ]
            else:
                print(f"[Warning] Image path {args.image_path} does not exist.")

-        result = client.generate(**generate_kwargs)
+        result = client.generate(
+            prompts,
+            OmniDiffusionSamplingParams(
+                seed=52,
+                need_kv_receive=False,
+                num_inference_steps=args.steps,
+            ),
+        )

        # Ensure result is a list for iteration
        if not isinstance(result, list):
@@ -100,8 +112,6 @@ def main():
            omni_outputs = result

    else:
-        import copy
-
        from vllm_omni.entrypoints.omni import Omni

        omni_kwargs = {}
@@ -144,11 +154,11 @@ def main():
                prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
                formatted_prompts.append(prompt_dict)

-        params_list = copy.deepcopy(omni.default_sampling_params_list)
+        params_list = omni.default_sampling_params_list
        if args.modality == "text2img":
-            params_list[0]["max_tokens"] = 1
+            params_list[0].max_tokens = 1  # type: ignore # The first stage is a SamplingParam (vllm)
            if len(params_list) > 1:
-                params_list[1]["num_inference_steps"] = args.steps
+                params_list[1].num_inference_steps = args.steps  # type: ignore # The second stage is an OmniDiffusionSamplingParam

        omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))

@@ -1,6 +0,0 @@
-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}
@@ -77,10 +77,11 @@ from pathlib import Path
 import torch
 from PIL import Image

-from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
+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.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform


 def parse_args() -> argparse.Namespace:
@@ -279,6 +280,32 @@ 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()


@@ -300,12 +327,8 @@ def main():
    else:
        input_image = input_images

-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).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,
@@ -338,12 +361,15 @@ def main():
    # Initialize Omni with appropriate pipeline
    omni = Omni(
        model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        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,
        cache_backend=args.cache_backend,
        cache_config=cache_config,
        parallel_config=parallel_config,
        enforce_eager=args.enforce_eager,
+        enable_cpu_offload=args.enable_cpu_offload,
    )
    print("Pipeline loaded")

@@ -368,25 +394,28 @@ def main():
    print(f"{'=' * 60}\n")

    generation_start = time.perf_counter()
-    # Generate edited image
-    generate_kwargs = {
-        "prompt": args.prompt,
-        "pil_image": input_image,
-        "negative_prompt": args.negative_prompt,
-        "generator": generator,
-        "true_cfg_scale": args.cfg_scale,
-        "guidance_scale": args.guidance_scale,
-        "num_inference_steps": args.num_inference_steps,
-        "num_outputs_per_prompt": args.num_outputs_per_prompt,
-        "layers": args.layers,
-        "resolution": args.resolution,
-    }

    if profiler_enabled:
        print("[Profiler] Starting profiling...")
        omni.start_profile()

-    outputs = omni.generate(**generate_kwargs)
+    # Generate edited image
+    outputs = omni.generate(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+            "multi_modal_data": {"image": input_image},
+        },
+        OmniDiffusionSamplingParams(
+            generator=generator,
+            true_cfg_scale=args.cfg_scale,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_outputs_per_prompt,
+            layers=args.layers,
+            resolution=args.resolution,
+        ),
+    )
    generation_end = time.perf_counter()
    generation_time = generation_end - generation_start

@@ -412,7 +441,6 @@ 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
@@ -47,3 +47,9 @@ 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.
@@ -54,3 +54,9 @@ 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.
@@ -26,9 +26,11 @@ 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.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform


 def parse_args() -> argparse.Namespace:
@@ -58,6 +60,44 @@ 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()


@@ -74,8 +114,7 @@ def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> t

 def main():
    args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)

    # Load input image
    image = PIL.Image.open(args.image).convert("RGB")
@@ -92,37 +131,54 @@ 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,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        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,
        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(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        pil_image=image,
-        height=height,
-        width=width,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        guidance_scale_2=args.guidance_scale_high,
-        num_inference_steps=args.num_inference_steps,
-        num_frames=args.num_frames,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+            "multi_modal_data": {"image": image},
+        },
+        OmniDiffusionSamplingParams(
+            height=height,
+            width=width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            guidance_scale_2=args.guidance_scale_high,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+        ),
    )

    # Extract video frames from OmniRequestOutput
@@ -0,0 +1,98 @@
+# LoRA Inference Examples
+
+This directory contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
+The example uses the  `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.
+
+## Overview
+
+Similar to vLLM, vLLM-omni uses a unified LoRA handling mechanism:
+
+- **Pre-loaded LoRA**: Loaded at initialization via `--lora-path` (pre-loaded into cache)
+- **Per-request LoRA**: Loaded on-demand. In the example, the LoRA is loaded via `--lora-request-path` in each request
+
+Both approaches use the same underlying mechanism - all LoRA adapters are handled uniformly through `set_active_adapter()`. If no LoRA request is provided in a request, all adapters are deactivated.
+
+## Usage
+
+### Pre-loaded LoRA (via --lora-path)
+
+Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_preloaded.png
+```
+
+**Note**: When using `--lora-path`, the adapter is loaded at init time with a stable ID derived from the adapter path. This example activates it automatically for the request.
+
+### Per-request LoRA (via --lora-request-path)
+
+Load a LoRA adapter on-demand for each request:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-request-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_per_request.png
+```
+
+### No LoRA
+
+If no LoRA request is provided, we will use the base model without any LoRA adapters:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_no_lora.png
+```
+
+## Parameters
+
+### LoRA Parameters
+
+- `--lora-path`: Path to LoRA adapter folder to pre-load at initialization (loads into cache with a stable ID derived from the path)
+- `--lora-request-path`: Path to LoRA adapter folder for per-request loading
+- `--lora-request-id`: Integer ID for the LoRA adapter (optional). If not provided and `--lora-request-path` is set, will derive a stable ID from the path.
+- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
+
+### Standard Parameters
+
+- `--prompt`: Text prompt for image generation (required)
+- `--seed`: Random seed for reproducibility (default: 42)
+- `--height`: Image height in pixels (default: 1024)
+- `--width`: Image width in pixels (default: 1024)
+- `--num_inference_steps`: Number of denoising steps (default: 50)
+- `--output`: Output file path (default: `lora_output.png`)
+
+## How LoRA Works
+
+All LoRA adapters are handled uniformly:
+
+1. **Initialization**: If `--lora-path` is provided, the adapter is loaded into cache with a stable ID derived from the adapter path
+2. **Per-request**: If `--lora-request-path` is provided, the adapter is loaded/activated for that request
+3. **No LoRA**: If no LoRA request is provided (`req.lora_request` is None), all adapters are deactivated
+
+The system uses LRU cache management - adapters are cached and evicted when the cache is full (unless pinned).
+
+## LoRA Adapter Format
+
+LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+from pathlib import Path
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.lora.utils import stable_lora_int_id
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate images with LoRA adapters.")
+    parser.add_argument("--model", default="stabilityai/stable-diffusion-3.5-medium", help="Model name or path.")
+    parser.add_argument("--prompt", required=True, help="Text prompt for image generation.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for deterministic results.")
+    parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
+    parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")
+    parser.add_argument(
+        "--num_inference_steps",
+        type=int,
+        default=50,
+        help="Number of denoising steps for the diffusion sampler.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="lora_output.png",
+        help="Path to save the generated image (PNG).",
+    )
+    parser.add_argument(
+        "--lora-path",
+        type=str,
+        default=None,
+        help="Path to LoRA adapter folder to pre-load at initialization (PEFT format). "
+        "Note: pre-loading populates the cache; you still need to pass a lora_request to activate it.",
+    )
+    parser.add_argument(
+        "--lora-request-path",
+        type=str,
+        default=None,
+        help="Path to LoRA adapter folder for per-request activation (dynamic LoRA). "
+        "If --lora-request-id is not provided, a stable ID will be derived from this path.",
+    )
+    parser.add_argument(
+        "--lora-request-id",
+        type=int,
+        default=None,
+        help="Integer ID for the LoRA adapter (for dynamic LoRA). "
+        "If not provided and --lora-request-path is set, will derive a stable ID from the path.",
+    )
+    parser.add_argument(
+        "--lora-scale",
+        type=float,
+        default=1.0,
+        help="Scale factor for LoRA weights (default: 1.0).",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    model = args.model
+
+    omni_kwargs = {}
+
+    if args.lora_path:
+        omni_kwargs["lora_path"] = args.lora_path
+        print(f"Using static LoRA from: {args.lora_path}")
+
+    omni = Omni(model=model, **omni_kwargs)
+
+    lora_request = None
+    if args.lora_request_path:
+        if args.lora_request_id is None:
+            lora_request_id = stable_lora_int_id(args.lora_request_path)
+        else:
+            lora_request_id = args.lora_request_id
+
+        lora_name = Path(args.lora_request_path).stem
+        lora_request = LoRARequest(
+            lora_name=lora_name,
+            lora_int_id=lora_request_id,
+            lora_path=args.lora_request_path,
+        )
+        print(f"Using per-request LoRA: name={lora_name}, id={lora_request_id}, scale={args.lora_scale}")
+    elif args.lora_path:
+        # pre-loaded LoRA
+        lora_request_id = stable_lora_int_id(args.lora_path)
+        lora_request = LoRARequest(
+            lora_name="preloaded",
+            lora_int_id=lora_request_id,
+            lora_path=args.lora_path,
+        )
+        print(f"Activating pre-loaded LoRA: id={lora_request_id}, scale={args.lora_scale}")
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+    )
+
+    if lora_request:
+        sampling_params.lora_request = lora_request
+        sampling_params.lora_scale = args.lora_scale
+
+    outputs = omni.generate(args.prompt, sampling_params)
+
+    if not outputs or len(outputs) == 0:
+        raise ValueError("No output generated from omni.generate()")
+
+    if isinstance(outputs, list):
+        first_output = outputs[0]
+    else:
+        first_output = outputs
+
+    images = None
+    if hasattr(first_output, "images") and first_output.images:
+        images = first_output.images
+    elif hasattr(first_output, "request_output") and first_output.request_output:
+        req_out = first_output.request_output
+        if isinstance(req_out, list) and len(req_out) > 0:
+            req_out = req_out[0]
+        if hasattr(req_out, "images") and req_out.images:
+            images = req_out.images
+
+    if not images:
+        raise ValueError("No images found in request_output")
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    suffix = output_path.suffix or ".png"
+    stem = output_path.stem or "lora_output"
+    if len(images) <= 1:
+        images[0].save(output_path)
+        print(f"Saved generated image to {output_path}")
+    else:
+        for idx, img in enumerate(images):
+            save_path = output_path.parent / f"{stem}_{idx}{suffix}"
+            img.save(save_path)
+            print(f"Saved generated image to {save_path}")
+
+
+if __name__ == "__main__":
+    main()
@@ -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.multimodal_output["audio"]
+                audio_tensor = output.outputs[0].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}")
@@ -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.multimodal_output["audio"]
+                audio_tensor = output.outputs[0].multimodal_output["audio"]
                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")

                # Convert to numpy array and ensure correct format
@@ -13,6 +13,15 @@ Qwen3 TTS provides multiple task variants for speech generation:
 ## Setup
 Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.

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

 Run a single sample for a task:
@@ -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.multimodal_output["audio"]
+            audio_tensor = output.outputs[0].multimodal_output["audio"]
            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-            audio_samplerate = output.multimodal_output["sr"].item()
+            audio_samplerate = output.outputs[0].multimodal_output["sr"].item()
            # Convert to numpy array and ensure correct format
            audio_numpy = audio_tensor.float().detach().cpu().numpy()

@@ -0,0 +1,37 @@
+# Text-To-Audio
+
+The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts.
+
+## Prerequisites
+
+If you use a gated model (e.g., `stabilityai/stable-audio-open-1.0`), ensure you have access:
+
+1. **Accept Model License**: Visit the model page on Hugging Face (e.g., [stabilityai/stable-audio-open-1.0]) and accept the user agreement.
+2. **Authenticate**: Log in to Hugging Face locally to access the gated model.
+   ```bash
+   huggingface-cli login
+   ```
+
+## Local CLI Usage
+
+```bash
+python text_to_audio.py \
+  --model stabilityai/stable-audio-open-1.0 \
+  --prompt "The sound of a hammer hitting a wooden surface" \
+  --negative_prompt "Low quality" \
+  --seed 42 \
+  --guidance_scale 7.0 \
+  --audio_length 10.0 \
+  --num_inference_steps 100 \
+  --output stable_audio_output.wav
+```
+
+Key arguments:
+
+- `--prompt`: text description (string).
+- `--negative_prompt`: negative prompt for classifier-free guidance.
+- `--seed`: integer seed for deterministic generation.
+- `--guidance_scale`: classifier-free guidance scale.
+- `--audio_length`: audio duration in seconds.
+- `--num_inference_steps`: diffusion sampling steps.(more steps = higher quality, slower).
+- `--output`: path to save the generated WAV file.
@@ -21,7 +21,8 @@ import numpy as np
 import torch

 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.platform_utils import detect_device_type
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform


 def parse_args() -> argparse.Namespace:
@@ -117,8 +118,7 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410

 def main():
    args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)

    print(f"\n{'=' * 60}")
    print("Stable Audio Open - Text-to-Audio Generation")
@@ -142,17 +142,21 @@ def main():
    generation_start = time.perf_counter()

    # Generate audio
-    audio = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_outputs_per_prompt=args.num_waveforms,
-        extra={
-            "audio_start_in_s": args.audio_start,
-            "audio_end_in_s": audio_end_in_s,
+    outputs = omni.generate(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
        },
+        OmniDiffusionSamplingParams(
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_waveforms,
+            extra_args={
+                "audio_start_in_s": args.audio_start,
+                "audio_end_in_s": audio_end_in_s,
+            },
+        ),
    )

    generation_end = time.perf_counter()
@@ -166,6 +170,21 @@ def main():
    suffix = output_path.suffix or ".wav"
    stem = output_path.stem or "stable_audio_output"

+    # Extract audio from omni.generate() outputs
+    if not outputs:
+        raise ValueError("No output generated from omni.generate()")
+
+    output = outputs[0]
+    if not hasattr(output, "request_output") or not output.request_output:
+        raise ValueError("No request_output found in OmniRequestOutput")
+    request_output = output.request_output[0]
+    if not hasattr(request_output, "multimodal_output"):
+        raise ValueError("No multimodal_output found in request_output")
+
+    audio = request_output.multimodal_output.get("audio")
+    if audio is None:
+        raise ValueError("No audio output found in request_output")
+
    # Handle different output formats
    if isinstance(audio, torch.Tensor):
        audio = audio.cpu().float().numpy()
@@ -20,7 +20,7 @@ if __name__ == "__main__":
    images[0].save("coffee.png")
 ```

-Or put more than one prompt in a request, processing them sequentially.
+Or put more than one prompt in a request.

 ```python
 from vllm_omni.entrypoints.omni import Omni
@@ -37,6 +37,41 @@ if __name__ == "__main__":
      image = output.request_output[0].images[0].save(f"{i}.jpg")
 ```

+!!! info
+
+    However, it is not currently recommended to do so
+    because not all models support batch inference,
+    and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
+    This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
+
+!!! info
+
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
+    batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+
+Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
+This is useful for models that support negative prompts.
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(model="Qwen/Qwen-Image")
+    outputs = omni.generate([
+      {
+        "prompt": "a cup of coffee on a table"，
+        "negative_prompt": "low resolution"
+      },
+      {
+        "prompt": "a toy dinosaur on a sandy beach"，
+        "negative_prompt": "cinematic, realistic"
+      }
+    ])
+    for i, output in enumerate(outputs):
+      image = output.request_output[0].images[0].save(f"{i}.jpg")
+```
+
 ## Local CLI Usage

 ```bash
@@ -61,6 +96,12 @@ 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.

@@ -5,8 +5,9 @@ import gradio as gr
 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.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform

 ASPECT_RATIOS: dict[str, tuple[int, int]] = {
    "1:1": (1328, 1328),
@@ -61,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 = is_npu()
-    vae_use_tiling = is_npu()
+    vae_use_slicing = current_omni_platform.is_npu()
+    vae_use_tiling = current_omni_platform.is_npu()
    return Omni(
        model=model_name,
        vae_use_slicing=vae_use_slicing,
@@ -71,7 +72,6 @@ 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,15 +98,17 @@ 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=device).manual_seed(seed)
+        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
        outputs = omni.generate(
            prompt.strip(),
-            height=height,
-            width=width,
-            generator=generator,
-            true_cfg_scale=float(cfg_scale_value),
-            num_inference_steps=num_steps,
-            num_outputs_per_prompt=num_images,
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                generator=generator,
+                true_cfg_scale=float(cfg_scale_value),
+                num_inference_steps=num_steps,
+                num_outputs_per_prompt=num_images,
+            ),
        )
        images_outputs = []
        for output in outputs:
@@ -10,8 +10,9 @@ import torch

 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.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform


 def parse_args() -> argparse.Namespace:
@@ -72,6 +73,11 @@ 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,
@@ -96,23 +102,44 @@ def parse_args() -> argparse.Namespace:
        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(
        "--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()
-    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()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)

    # Configure cache based on backend type
    cache_config = None
@@ -156,12 +183,16 @@ def main():

    omni = Omni(
        model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        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,
        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,
    )

    if profiler_enabled:
@@ -183,15 +214,19 @@ def main():

    generation_start = time.perf_counter()
    outputs = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        height=args.height,
-        width=args.width,
-        generator=generator,
-        true_cfg_scale=args.cfg_scale,
-        guidance_scale=args.guidance_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_outputs_per_prompt=args.num_images_per_prompt,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+        },
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            true_cfg_scale=args.cfg_scale,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_images_per_prompt,
+        ),
    )
    generation_end = time.perf_counter()
    generation_time = generation_end - generation_start
@@ -9,10 +9,11 @@ python text_to_video.py \
  --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
  --negative_prompt "<optional quality filter>" \
  --height 480 \
-  --width 640 \
+  --width 832 \
  --num_frames 33 \
  --guidance_scale 4.0 \
  --guidance_scale_high 3.0 \
+  --flow_shift 12.0 \
  --num_inference_steps 40 \
  --fps 16 \
  --output t2v_out.mp4
@@ -21,10 +22,16 @@ python text_to_video.py \
 Key arguments:

 - `--prompt`: text description (string).
- `--height/--width`: output resolution (defaults 720x1280). Dimensions should align with Wan VAE downsampling (multiples of 8).
+- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
 - `--num_frames`: Number of frames (Wan default is 81).
- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high)..
+- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high).
 - `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
- `--boundary_ratio`: Boundary split ratio for low/high DiT.
+- `--boundary_ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../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.
@@ -3,14 +3,17 @@

 import argparse
 import os
+import time
 from pathlib import Path

 import numpy as np
 import torch

+from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform


 def parse_args() -> argparse.Namespace:
@@ -29,50 +32,173 @@ 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.")
+    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(
        "--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,
+        default=1,
+        help="Number of GPUs used for ulysses sequence parallelism.",
+    )
+    parser.add_argument(
+        "--ring_degree",
+        type=int,
+        default=1,
+        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.",
+    )
+
    return parser.parse_args()


 def main():
    args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).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
    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))

    omni = Omni(
        model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        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,
        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,
    )

    if profiler_enabled:
        print("[Profiler] Starting profiling...")
        omni.start_profile()

-    frames = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        height=args.height,
-        width=args.width,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        guidance_scale_2=args.guidance_scale_high,
-        num_inference_steps=args.num_inference_steps,
-        num_frames=args.num_frames,
+    # 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: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
    )
+    print(f"  Video size: {args.width}x{args.height}")
+    print(f"{'=' * 60}\n")
+
+    generation_start = time.perf_counter()
+    frames = omni.generate(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+        },
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            guidance_scale_2=args.guidance_scale_high,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+        ),
+    )
+    generation_end = time.perf_counter()
+    generation_time = generation_end - generation_start
+
+    # Print profiling results
+    print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")

    # Extract video frames from OmniRequestOutput
    if isinstance(frames, list) and len(frames) > 0:
@@ -0,0 +1,230 @@
+# 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**       |
@@ -0,0 +1,54 @@
+# Online LoRA Inference (Diffusion)
+
+This example shows how to use **per-request LoRA** with vLLM-Omni diffusion models via the OpenAI-compatible Chat Completions API.
+
+> Note: The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
+> Note: This example uses `/v1/chat/completions`. LoRA payloads for other OpenAI endpoints are not implemented here.
+
+## Start Server
+
+```bash
+# Pick a diffusion model (examples)
+# export MODEL=stabilityai/stable-diffusion-3.5-medium
+# export MODEL=Qwen/Qwen-Image
+
+bash run_server.sh
+```
+
+## Call API (curl)
+
+```bash
+# Required: local LoRA folder on the server
+export LORA_PATH=/path/to/lora_adapter
+
+# Optional
+export SERVER=http://localhost:8091
+export PROMPT="A piece of cheesecake"
+export LORA_NAME=my_lora
+export LORA_SCALE=1.0
+# Optional: if omitted, the server derives a stable id from LORA_PATH.
+# export LORA_INT_ID=123
+
+bash run_curl_lora_inference.sh
+```
+
+## Call API (Python)
+
+```bash
+python openai_chat_client.py \
+  --prompt "A piece of cheesecake" \
+  --lora-path /path/to/lora_adapter \
+  --lora-name my_lora \
+  --lora-scale 1.0 \
+  --output output.png
+```
+
+## LoRA Format
+
+LoRA adapters should be in PEFT format, for example:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""
+OpenAI-compatible chat client for diffusion LoRA inference.
+
+Example:
+  python openai_chat_client.py \
+    --server http://localhost:8091 \
+    --prompt "A piece of cheesecake" \
+    --lora-path /path/to/lora_adapter \
+    --lora-name my_lora \
+    --lora-scale 1.0 \
+    --output output.png
+"""
+
+import argparse
+import base64
+from pathlib import Path
+
+import requests
+
+
+def generate_image(
+    *,
+    prompt: str,
+    server_url: str,
+    height: int | None,
+    width: int | None,
+    num_inference_steps: int | None,
+    seed: int | None,
+    lora_name: str | None,
+    lora_path: str | None,
+    lora_scale: float | None,
+    lora_int_id: int | None,
+) -> bytes | None:
+    messages = [{"role": "user", "content": prompt}]
+
+    extra_body: dict = {}
+    if height is not None:
+        extra_body["height"] = height
+    if width is not None:
+        extra_body["width"] = width
+    if num_inference_steps is not None:
+        extra_body["num_inference_steps"] = num_inference_steps
+    if seed is not None:
+        extra_body["seed"] = seed
+
+    if lora_path:
+        lora_body: dict = {
+            "local_path": lora_path,
+            "name": lora_name or Path(lora_path).stem,
+        }
+        if lora_scale is not None:
+            lora_body["scale"] = float(lora_scale)
+        if lora_int_id is not None:
+            lora_body["int_id"] = int(lora_int_id)
+        extra_body["lora"] = lora_body
+
+    payload = {"messages": messages}
+    if extra_body:
+        payload["extra_body"] = extra_body
+
+    response = requests.post(
+        f"{server_url}/v1/chat/completions",
+        headers={"Content-Type": "application/json"},
+        json=payload,
+        timeout=300,
+    )
+    response.raise_for_status()
+    data = response.json()
+
+    content = data["choices"][0]["message"]["content"]
+    if isinstance(content, list) and content:
+        image_url = content[0].get("image_url", {}).get("url", "")
+        if image_url.startswith("data:image"):
+            _, b64_data = image_url.split(",", 1)
+            return base64.b64decode(b64_data)
+
+    raise RuntimeError(f"Unexpected response format: {content!r}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Diffusion LoRA OpenAI chat client")
+    parser.add_argument("--server", default="http://localhost:8091", help="Server URL")
+    parser.add_argument("--prompt", default="A piece of cheesecake", help="Text prompt")
+    parser.add_argument("--output", default="lora_online_output.png", help="Output image path")
+
+    parser.add_argument("--height", type=int, default=1024, help="Image height")
+    parser.add_argument("--width", type=int, default=1024, help="Image width")
+    parser.add_argument("--steps", type=int, default=50, help="num_inference_steps")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+
+    parser.add_argument("--lora-path", default=None, help="Server-local LoRA adapter folder (PEFT format)")
+    parser.add_argument("--lora-name", default=None, help="LoRA name (optional)")
+    parser.add_argument("--lora-scale", type=float, default=1.0, help="LoRA scale")
+    parser.add_argument(
+        "--lora-int-id",
+        type=int,
+        default=None,
+        help="LoRA integer id (cache key). If omitted, the server derives a stable id from lora_path.",
+    )
+
+    args = parser.parse_args()
+
+    image_bytes = generate_image(
+        prompt=args.prompt,
+        server_url=args.server,
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.steps,
+        seed=args.seed,
+        lora_name=args.lora_name,
+        lora_path=args.lora_path,
+        lora_scale=args.lora_scale if args.lora_path else None,
+        lora_int_id=args.lora_int_id if args.lora_path else None,
+    )
+
+    out_path = Path(args.output)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_bytes(image_bytes)
+    print(f"Saved: {out_path} ({len(image_bytes) / 1024:.1f} KiB)")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,61 @@
+#!/bin/bash
+# Online diffusion LoRA inference via OpenAI-compatible chat API.
+
+SERVER="${SERVER:-http://localhost:8091}"
+PROMPT="${PROMPT:-A piece of cheesecake}"
+
+LORA_PATH="${LORA_PATH:-}"
+LORA_NAME="${LORA_NAME:-lora}"
+LORA_SCALE="${LORA_SCALE:-1.0}"
+LORA_INT_ID="${LORA_INT_ID:-}"
+
+HEIGHT="${HEIGHT:-1024}"
+WIDTH="${WIDTH:-1024}"
+NUM_INFERENCE_STEPS="${NUM_INFERENCE_STEPS:-50}"
+SEED="${SEED:-42}"
+
+CURRENT_TIME=$(date +%Y%m%d%H%M%S)
+OUTPUT="${OUTPUT:-lora_online_output_${CURRENT_TIME}.png}"
+
+if [ -z "$LORA_PATH" ]; then
+  echo "ERROR: LORA_PATH is required (must be a server-local path)."
+  exit 1
+fi
+
+echo "Generating image with LoRA..."
+echo "Server: $SERVER"
+echo "Prompt: $PROMPT"
+echo "LoRA: name=$LORA_NAME id=${LORA_INT_ID:-auto} scale=$LORA_SCALE path=$LORA_PATH"
+echo "Output: $OUTPUT"
+
+LORA_INT_ID_FIELD=""
+if [ -n "$LORA_INT_ID" ]; then
+  LORA_INT_ID_FIELD=", \"int_id\": $LORA_INT_ID"
+fi
+
+curl -s "$SERVER/v1/chat/completions" \
+  -H "Content-Type: application/json" \
+  -d "{
+    \"messages\": [
+      {\"role\": \"user\", \"content\": \"$PROMPT\"}
+    ],
+    \"extra_body\": {
+      \"height\": $HEIGHT,
+      \"width\": $WIDTH,
+      \"num_inference_steps\": $NUM_INFERENCE_STEPS,
+      \"seed\": $SEED,
+      \"lora\": {
+        \"name\": \"$LORA_NAME\",
+        \"local_path\": \"$LORA_PATH\",
+        \"scale\": $LORA_SCALE$LORA_INT_ID_FIELD
+      }
+    }
+  }" | jq -r '.choices[0].message.content[0].image_url.url' | sed 's/^data:image[^,]*,\s*//' | base64 -d > "$OUTPUT"
+
+if [ -f "$OUTPUT" ]; then
+  echo "Image saved to: $OUTPUT"
+  echo "Size: $(du -h "$OUTPUT" | cut -f1)"
+else
+  echo "Failed to generate image"
+  exit 1
+fi
@@ -0,0 +1,20 @@
+#!/bin/bash
+# Online diffusion serving with vLLM-Omni (OpenAI-compatible API).
+
+MODEL="${MODEL:-stabilityai/stable-diffusion-3.5-medium}"
+PORT="${PORT:-8091}"
+
+echo "Starting vLLM-Omni diffusion server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+
+if [ -z "${VLLM_BIN:-}" ]; then
+  if command -v vllm-omni >/dev/null 2>&1; then
+    VLLM_BIN="vllm-omni"
+  else
+    VLLM_BIN="vllm"
+  fi
+fi
+
+"$VLLM_BIN" serve "$MODEL" --omni \
+  --port "$PORT"
@@ -12,6 +12,12 @@ Please refer to [README.md](../../../README.md)
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```

+If you want to open async chunking for qwen3-omni, launch the server with command below
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+```
+
 If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
@@ -1,4 +1,5 @@
 import base64
+import concurrent.futures
 import os
 from typing import NamedTuple

@@ -414,51 +415,67 @@ def run_multimodal_generation(args) -> None:
    else:
        output_modalities = None

-    chat_completion = client.chat.completions.create(
-        messages=[
-            get_system_prompt(),
-            prompt,
-        ],
-        model=model_name,
-        modalities=output_modalities,
-        extra_body=extra_body,
-        stream=args.stream,
-    )
+    # Test multiple concurrent completions
+    num_concurrent_requests = args.num_concurrent_requests

+    with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
+        # Submit multiple completion requests concurrently
+        futures = [
+            executor.submit(
+                client.chat.completions.create,
+                messages=[
+                    get_system_prompt(),
+                    prompt,
+                ],
+                model=model_name,
+                modalities=output_modalities,
+                extra_body=extra_body,
+                stream=args.stream,
+            )
+            for _ in range(num_concurrent_requests)
+        ]
+
+        # Wait for all requests to complete and collect results
+        chat_completions = [future.result() for future in concurrent.futures.as_completed(futures)]
+
+    assert len(chat_completions) == num_concurrent_requests
    count = 0
    if not args.stream:
-        for choice in chat_completion.choices:
-            if choice.message.audio:
-                audio_data = base64.b64decode(choice.message.audio.data)
-                audio_file_path = f"audio_{count}.wav"
-                with open(audio_file_path, "wb") as f:
-                    f.write(audio_data)
-                print(f"Audio saved to {audio_file_path}")
-                count += 1
-            elif choice.message.content:
-                print("Chat completion output from text:", choice.message.content)
-    else:
-        printed_content = False
-        for chunk in chat_completion:
-            for choice in chunk.choices:
-                if hasattr(choice, "delta"):
-                    content = getattr(choice.delta, "content", None)
-                else:
-                    content = None
-
-                if getattr(chunk, "modality", None) == "audio" and content:
-                    audio_data = base64.b64decode(content)
+        # Verify all completions succeeded
+        for chat_completion in chat_completions:
+            for choice in chat_completion.choices:
+                if choice.message.audio:
+                    audio_data = base64.b64decode(choice.message.audio.data)
                    audio_file_path = f"audio_{count}.wav"
                    with open(audio_file_path, "wb") as f:
                        f.write(audio_data)
-                    print(f"\nAudio saved to {audio_file_path}")
+                    print(f"Audio saved to {audio_file_path}")
                    count += 1
+                elif choice.message.content:
+                    print("Chat completion output from text:", choice.message.content)
+    else:
+        printed_content = False
+        for chat_completion in chat_completions:
+            for chunk in chat_completion:
+                for choice in chunk.choices:
+                    if hasattr(choice, "delta"):
+                        content = getattr(choice.delta, "content", None)
+                    else:
+                        content = None

-                elif getattr(chunk, "modality", None) == "text":
-                    if not printed_content:
-                        printed_content = True
-                        print("\ncontent:", end="", flush=True)
-                    print(content, end="", flush=True)
+                    if getattr(chunk, "modality", None) == "audio" and content:
+                        audio_data = base64.b64decode(content)
+                        audio_file_path = f"audio_{count}.wav"
+                        with open(audio_file_path, "wb") as f:
+                            f.write(audio_data)
+                        print(f"\nAudio saved to {audio_file_path}")
+                        count += 1
+
+                    elif getattr(chunk, "modality", None) == "text":
+                        if not printed_content:
+                            printed_content = True
+                            print("\ncontent:", end="", flush=True)
+                        print(content, end="", flush=True)


 def parse_args():
@@ -517,6 +534,12 @@ def parse_args():
        action="store_true",
        help="Stream the response.",
    )
+    parser.add_argument(
+        "--num-concurrent-requests",
+        type=int,
+        default=1,
+        help="Number of concurrent requests to send. Default is 1.",
+    )

    return parser.parse_args()

@@ -13,7 +13,7 @@ stage_args:
    engine_args:
      model_stage: thinker
      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
      gpu_memory_utilization: 0.9
      enforce_eager: true
@@ -0,0 +1,175 @@
+# Qwen3-TTS Online Serving
+
+This directory contains examples for running Qwen3-TTS models with vLLM-Omni's online serving API.
+
+## Supported Models
+
+| Model | Task Type | Description |
+|-------|-----------|-------------|
+| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
+
+## Quick Start
+
+### 1. Start the Server
+
+```bash
+# CustomVoice model (default)
+./run_server.sh
+
+# Or specify task type
+./run_server.sh CustomVoice
+./run_server.sh VoiceDesign
+./run_server.sh Base
+```
+
+### 2. Run the Client
+
+```bash
+# CustomVoice: Use predefined speaker
+python openai_speech_client.py \
+    --text "你好，我是通义千问" \
+    --voice Vivian \
+    --language Chinese
+
+# CustomVoice with style instruction
+python openai_speech_client.py \
+    --text "今天天气真好" \
+    --voice Ryan \
+    --instructions "用开心的语气说"
+
+# VoiceDesign: Describe the voice style
+python openai_speech_client.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
+    --task-type VoiceDesign \
+    --text "哥哥，你回来啦" \
+    --instructions "体现撒娇稚嫩的萝莉女声，音调偏高"
+
+# Base: Voice cloning
+python openai_speech_client.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --task-type Base \
+    --text "Hello, this is a cloned voice" \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Original transcript of the reference audio"
+```
+
+### 3. Using curl
+
+```bash
+# Simple TTS request
+curl -X POST http://localhost:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "Vivian",
+        "language": "English"
+    }' --output output.wav
+
+# With style instruction
+curl -X POST http://localhost:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "I am so excited!",
+        "voice": "Vivian",
+        "instructions": "Speak with great enthusiasm"
+    }' --output excited.wav
+
+# List available voices in CustomVoice models
+curl http://localhost:8000/v1/audio/voices
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/speech
+```
+
+This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
+
+### Request Body
+
+```json
+{
+    "input": "Text to synthesize",
+    "voice": "Vivian",
+    "response_format": "wav",
+    "task_type": "CustomVoice",
+    "language": "Auto",
+    "instructions": "Optional style instructions",
+    "ref_audio": "URL or base64 for voice cloning",
+    "ref_text": "Reference audio transcript",
+    "x_vector_only_mode": false,
+    "max_new_tokens": 2048
+}
+```
+
+> **Note:** The `model` field is optional when serving a single model, as the server already knows which model is loaded.
+
+### Response
+
+Returns audio data in the requested format (default: WAV).
+
+## Parameters
+
+### Standard OpenAI Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | required | Text to synthesize |
+| `voice` | string | "Vivian" | Speaker/voice name |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
+| `model` | string | optional | Model name (optional when serving single model) |
+
+### Qwen3-TTS Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base |
+| `language` | string | "Auto" | Language: Auto, Chinese, English, Japanese, Korean |
+| `instructions` | string | "" | Voice style/emotion instructions |
+| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
+
+### Voice Clone Parameters (Base task)
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `ref_audio` | string | Yes* | Reference audio (file path, URL, or base64) |
+| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
+| `x_vector_only_mode` | bool | false | Use speaker embedding only (no ICL) |
+
+## Python Usage
+
+```python
+import httpx
+
+# Simple request
+response = httpx.post(
+    "http://localhost:8000/v1/audio/speech",
+    json={
+        "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+        "input": "Hello world",
+        "voice": "Vivian",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Limitations
+
+- **No streaming**: Audio is generated completely before being returned. Streaming will be supported after the pipeline is disaggregated (see RFC #938).
+- **Single request**: Batch processing is not yet optimized for online serving.
+
+## Troubleshooting
+
+1. **Connection refused**: Make sure the server is running on the correct port
+2. **Out of memory**: Reduce `--gpu-memory-utilization` in run_server.sh
+3. **Unsupported speaker**: Check supported speakers via model documentation
+4. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning
@@ -0,0 +1,240 @@
+"""OpenAI-compatible client for Qwen3-TTS via /v1/audio/speech endpoint.
+
+This script demonstrates how to use the OpenAI-compatible speech API
+to generate audio from text using Qwen3-TTS models.
+
+Examples:
+    # CustomVoice task (predefined speaker)
+    python openai_speech_client.py --text "Hello, how are you?" --voice Vivian
+
+    # CustomVoice with emotion instruction
+    python openai_speech_client.py --text "I'm so happy!" --voice Vivian \
+        --instructions "Speak with excitement"
+
+    # VoiceDesign task (voice from description)
+    python openai_speech_client.py --text "Hello world" \
+        --task-type VoiceDesign \
+        --instructions "A warm, friendly female voice"
+
+    # Base task (voice cloning)
+    python openai_speech_client.py --text "Hello world" \
+        --task-type Base \
+        --ref-audio "https://example.com/reference.wav" \
+        --ref-text "This is the reference transcript"
+"""
+
+import argparse
+import base64
+import os
+
+import httpx
+
+# Default server configuration
+DEFAULT_API_BASE = "http://localhost:8000"
+DEFAULT_API_KEY = "EMPTY"
+
+
+def encode_audio_to_base64(audio_path: str) -> str:
+    """Encode a local audio file to base64 data URL."""
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+    # Detect MIME type from extension
+    audio_path_lower = audio_path.lower()
+    if audio_path_lower.endswith(".wav"):
+        mime_type = "audio/wav"
+    elif audio_path_lower.endswith((".mp3", ".mpeg")):
+        mime_type = "audio/mpeg"
+    elif audio_path_lower.endswith(".flac"):
+        mime_type = "audio/flac"
+    elif audio_path_lower.endswith(".ogg"):
+        mime_type = "audio/ogg"
+    else:
+        mime_type = "audio/wav"  # Default
+
+    with open(audio_path, "rb") as f:
+        audio_bytes = f.read()
+    audio_b64 = base64.b64encode(audio_bytes).decode("utf-8")
+    return f"data:{mime_type};base64,{audio_b64}"
+
+
+def run_tts_generation(args) -> None:
+    """Run TTS generation via OpenAI-compatible /v1/audio/speech API."""
+
+    # Build request payload
+    payload = {
+        "model": args.model,
+        "input": args.text,
+        "voice": args.voice,
+        "response_format": args.response_format,
+    }
+
+    # Add optional parameters
+    if args.instructions:
+        payload["instructions"] = args.instructions
+    if args.task_type:
+        payload["task_type"] = args.task_type
+    if args.language:
+        payload["language"] = args.language
+    if args.max_new_tokens:
+        payload["max_new_tokens"] = args.max_new_tokens
+
+    # Voice clone parameters (Base task)
+    if args.ref_audio:
+        if args.ref_audio.startswith(("http://", "https://")):
+            payload["ref_audio"] = args.ref_audio
+        else:
+            payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
+    if args.ref_text:
+        payload["ref_text"] = args.ref_text
+    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}")
+    print("Generating audio...")
+
+    # Make the API call
+    api_url = f"{args.api_base}/v1/audio/speech"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {args.api_key}",
+    }
+
+    with httpx.Client(timeout=300.0) as client:
+        response = client.post(api_url, json=payload, headers=headers)
+
+    if response.status_code != 200:
+        print(f"Error: {response.status_code}")
+        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:
+        f.write(response.content)
+    print(f"Audio saved to: {output_path}")
+
+
+def parse_args():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="OpenAI-compatible client for Qwen3-TTS via /v1/audio/speech",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+
+    # Server configuration
+    parser.add_argument(
+        "--api-base",
+        type=str,
+        default=DEFAULT_API_BASE,
+        help=f"API base URL (default: {DEFAULT_API_BASE})",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=DEFAULT_API_KEY,
+        help="API key (default: EMPTY)",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+        help="Model name/path",
+    )
+
+    # Task configuration
+    parser.add_argument(
+        "--task-type",
+        "-t",
+        type=str,
+        default=None,
+        choices=["CustomVoice", "VoiceDesign", "Base"],
+        help="TTS task type (default: CustomVoice)",
+    )
+
+    # Input text
+    parser.add_argument(
+        "--text",
+        type=str,
+        required=True,
+        help="Text to synthesize",
+    )
+
+    # Voice/speaker
+    parser.add_argument(
+        "--voice",
+        type=str,
+        default="Vivian",
+        help="Speaker/voice name (default: Vivian). Options: Vivian, Ryan, etc.",
+    )
+    parser.add_argument(
+        "--language",
+        type=str,
+        default=None,
+        help="Language: Auto, Chinese, English, etc.",
+    )
+    parser.add_argument(
+        "--instructions",
+        type=str,
+        default=None,
+        help="Voice style/emotion instructions",
+    )
+
+    # Base (voice clone) parameters
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Reference audio file path or URL for voice cloning (Base task)",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Reference audio transcript for voice cloning (Base task)",
+    )
+    parser.add_argument(
+        "--x-vector-only",
+        action="store_true",
+        help="Use x-vector only mode for voice cloning (no ICL)",
+    )
+
+    # Generation parameters
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=None,
+        help="Maximum new tokens to generate",
+    )
+
+    # Output
+    parser.add_argument(
+        "--response-format",
+        type=str,
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
+        help="Audio output format (default: wav)",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        default=None,
+        help="Output audio file path (default: tts_output.wav)",
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_tts_generation(args)
@@ -0,0 +1,40 @@
+#!/bin/bash
+# Launch vLLM-Omni server for Qwen3-TTS models
+#
+# Usage:
+#   ./run_server.sh                           # Default: CustomVoice model
+#   ./run_server.sh CustomVoice               # CustomVoice model
+#   ./run_server.sh VoiceDesign               # VoiceDesign model
+#   ./run_server.sh Base                      # Base (voice clone) model
+
+set -e
+
+TASK_TYPE="${1:-CustomVoice}"
+
+case "$TASK_TYPE" in
+    CustomVoice)
+        MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
+        ;;
+    VoiceDesign)
+        MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
+        ;;
+    Base)
+        MODEL="Qwen/Qwen3-TTS-12Hz-1.7B-Base"
+        ;;
+    *)
+        echo "Unknown task type: $TASK_TYPE"
+        echo "Supported: CustomVoice, VoiceDesign, Base"
+        exit 1
+        ;;
+esac
+
+echo "Starting Qwen3-TTS server with model: $MODEL"
+
+vllm-omni serve "$MODEL" \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --host 0.0.0.0 \
+    --port 8000 \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Qwen-Image OpenAI-compatible chat client for image generation.
+Qwen-Image OpenAI-compatible image generation client.

 Usage:
    python openai_chat_client.py --prompt "A beautiful landscape" --output output.png
@@ -25,14 +25,14 @@ def generate_image(
    negative_prompt: str | None = None,
    num_outputs_per_prompt: int = 1,
 ) -> bytes | None:
-    """Generate an image using the chat completions API.
+    """Generate an image using the images generation API.

    Args:
        prompt: Text description of the image
        server_url: Server URL
        height: Image height in pixels
        width: Image width in pixels
-        steps: Number of inference steps
+        steps: Number of diffusion steps
        true_cfg_scale: Qwen-Image CFG scale
        seed: Random seed
        negative_prompt: Negative prompt
@@ -41,33 +41,31 @@ def generate_image(
    Returns:
        Image bytes or None if failed
    """
-    messages = [{"role": "user", "content": prompt}]
+    payload: dict[str, object] = {
+        "prompt": prompt,
+        "response_format": "b64_json",
+        "n": num_outputs_per_prompt,
+    }
+
+    if width is not None and height is not None:
+        payload["size"] = f"{width}x{height}"
+    elif width is not None:
+        payload["size"] = f"{width}x{width}"
+    elif height is not None:
+        payload["size"] = f"{height}x{height}"

-    # Build extra_body with generation parameters
-    extra_body = {}
-    if height is not None:
-        extra_body["height"] = height
-    if width is not None:
-        extra_body["width"] = width
    if steps is not None:
-        extra_body["num_inference_steps"] = steps
+        payload["num_inference_steps"] = steps
    if true_cfg_scale is not None:
-        extra_body["true_cfg_scale"] = true_cfg_scale
-    if seed is not None:
-        extra_body["seed"] = seed
+        payload["true_cfg_scale"] = true_cfg_scale
    if negative_prompt:
-        extra_body["negative_prompt"] = negative_prompt
-    extra_body["num_outputs_per_prompt"] = num_outputs_per_prompt
+        payload["negative_prompt"] = negative_prompt
+    if seed is not None:
+        payload["seed"] = seed

-    # Build request payload
-    payload = {"messages": messages}
-    if extra_body:
-        payload["extra_body"] = extra_body
-
-    # Send request
    try:
        response = requests.post(
-            f"{server_url}/v1/chat/completions",
+            f"{server_url}/v1/images/generations",
            headers={"Content-Type": "application/json"},
            json=payload,
            timeout=300,
@@ -75,15 +73,13 @@ def generate_image(
        response.raise_for_status()
        data = response.json()

-        # Extract image from response
-        content = data["choices"][0]["message"]["content"]
-        if isinstance(content, list) and len(content) > 0:
-            image_url = content[0].get("image_url", {}).get("url", "")
-            if image_url.startswith("data:image"):
-                _, b64_data = image_url.split(",", 1)
-                return base64.b64decode(b64_data)
+        items = data.get("data")
+        if isinstance(items, list) and items:
+            first = items[0].get("b64_json") if isinstance(items[0], dict) else None
+            if isinstance(first, str):
+                return base64.b64decode(first)

-        print(f"Unexpected response format: {content}")
+        print(f"Unexpected response format: {data}")
        return None

    except Exception as e:
@@ -1,35 +1,10 @@
 #!/bin/bash
 # Qwen-Image text-to-image curl example

-SERVER="${SERVER:-http://localhost:8091}"
-PROMPT="${PROMPT:-a good boy in the ocean}"
-CURRENT_TIME=$(date +%Y%m%d%H%M%S)
-OUTPUT="${OUTPUT:-qwen_image_output_${CURRENT_TIME}.png}"
-
-echo "Generating image..."
-echo "Prompt: $PROMPT"
-echo "Output: $OUTPUT"
-
-curl -s "$SERVER/v1/chat/completions" \
+curl -X POST http://localhost:8091/v1/images/generations \
  -H "Content-Type: application/json" \
-  -d "{
-    \"messages\": [
-      {\"role\": \"user\", \"content\": \"$PROMPT\"}
-    ],
-    \"extra_body\": {
-      \"height\": 1024,
-      \"width\": 1024,
-      \"num_inference_steps\": 100,
-      \"true_cfg_scale\": 4.0,
-      \"seed\": 42,
-      \"num_outputs_per_prompt\": 1
-    }
-  }" | jq -r '.choices[0].message.content[0].image_url.url' | sed 's/^data:image[^,]*,\s*//' | base64 -d > "$OUTPUT"
-
-if [ -f "$OUTPUT" ]; then
-    echo "Image saved to: $OUTPUT"
-    echo "Size: $(du -h "$OUTPUT" | cut -f1)"
-else
-    echo "Failed to generate image"
-    exit 1
-fi
+  -d '{
+    "prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
+    "size": "1024x1024",
+    "seed": 42
+  }' | jq -r '.data[0].b64_json' | base64 -d > dragon.png
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"

 [project]
 name = "vllm-omni"
-version = "0.12.0rc1"
+version = "0.14.0"
 description = "A framework for efficient model inference with omni-modality models"
 readme = "README.md"
 requires-python = ">=3.10,<3.14"
@@ -39,9 +39,11 @@ dependencies = [
    "cache-dit==1.2.0",
    "tqdm>=4.66.0",
    "torchsde>=0.2.6",  # Required for Stable Audio scheduler
-    "fa3-fwd", # flash attention 3, maintained by @ZJY0516
+    "fa3-fwd==0.0.1", # flash attention 3, maintained by @ZJY0516
    "openai-whisper>=20250625",
    "imageio[ffmpeg]>=2.37.2",
+    "onnxruntime>=1.19.0",
+    "sox>=1.5.0",
    # "vllm==0.14.0",  # TODO: fix the entrypoints overwrite problem
 ]

@@ -57,7 +59,8 @@ dev = [
    "psutil>=7.2.0",
    "soundfile>=0.13.1",
    "imageio[ffmpeg]>=0.6.0",
-    "opencv-python>=4.12.0.88"
+    "opencv-python>=4.12.0.88",
+    "mooncake-transfer-engine==0.3.8.post1"
 ]

 docs = [
@@ -200,4 +203,5 @@ extend-ignore-identifiers-re = [
    ".*nothink.*",
    ".*NOTHINK.*",
    ".*nin.*",
+    "Ono_Anna",
 ]
@@ -0,0 +1,60 @@
+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}"
@@ -10,6 +10,7 @@ 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
@@ -22,6 +23,7 @@ import psutil
 import pytest
 import torch
 import yaml
+from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 from vllm.logger import init_logger
 from vllm.utils.network_utils import get_open_port

@@ -52,31 +54,98 @@ def default_vllm_config():

@pytest.fixture(autouse=True)
 def clean_gpu_memory_between_tests():
-    if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
-        yield
+    print("\n=== PRE-TEST GPU CLEANUP ===")
+    _run_pre_test_cleanup()
+    yield
+    _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:
+        print("GPU cleanup disabled")
        return

-    # Wait for GPU memory to be cleared before starting the test
-    import gc
-
-    from tests.utils import wait_for_gpu_memory_to_clear
+    print("Pre-test GPU status:")

    num_gpus = torch.cuda.device_count()
    if num_gpus > 0:
        try:
+            from tests.utils import wait_for_gpu_memory_to_clear
+
            wait_for_gpu_memory_to_clear(
                devices=list(range(num_gpus)),
-                threshold_ratio=0.1,
+                threshold_ratio=0.05,
            )
-        except ValueError as e:
-            logger.info("Failed to clean GPU memory: %s", e)
+        except Exception as e:
+            print(f"Pre-test cleanup note: {e}")

-    yield

-    # Clean up GPU memory after the test
+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")
+        return
+
    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
        gc.collect()
+        torch.cuda.empty_cache()
+
+        print("Post-test GPU status:")
+        _print_gpu_processes()
+
+
+def _print_gpu_processes():
+    """Print GPU information including nvidia-smi and system processes"""
+
+    print("\n" + "=" * 80)
+    print("NVIDIA GPU Information (nvidia-smi)")
+    print("=" * 80)
+
+    try:
+        nvidia_result = subprocess.run(
+            ["nvidia-smi"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+
+        if nvidia_result.returncode == 0:
+            lines = nvidia_result.stdout.strip().split("\n")
+            for line in lines[:20]:
+                print(line)
+
+            if len(lines) > 20:
+                print(f"... (showing first 20 of {len(lines)} lines)")
+        else:
+            print("nvidia-smi command failed")
+
+    except (subprocess.TimeoutExpired, FileNotFoundError):
+        print("nvidia-smi not available or timed out")
+    except Exception as e:
+        print(f"Error running nvidia-smi: {e}")
+
+    print("\n" + "=" * 80)
+    print("Detailed GPU Processes (nvidia-smi pmon)")
+    print("=" * 80)
+
+    try:
+        pmon_result = subprocess.run(
+            ["nvidia-smi", "pmon", "-c", "1"],
+            capture_output=True,
+            text=True,
+            timeout=3,
+        )
+
+        if pmon_result.returncode == 0 and pmon_result.stdout.strip():
+            print(pmon_result.stdout)
+        else:
+            print("No active GPU processes found via nvidia-smi pmon")
+
+    except Exception:
+        print("nvidia-smi pmon not available")
+
+    print("\n" + "=" * 80)
+    print("System Processes with GPU keywords")
+    print("=" * 80)


 def dummy_messages_from_mix_data(
@@ -528,7 +597,36 @@ def convert_audio_to_text(audio_data):
        temperature=0.0,
        word_timestamps=True,
        condition_on_previous_text=False,
-        initial_prompt="Please transcribe with proper word spacing.",
+    )["text"]
+    if text:
+        return text
+    else:
+        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
@@ -538,75 +636,238 @@ def convert_audio_to_text(audio_data):

 def modify_stage_config(
    yaml_path: str,
-    stage_updates: dict[int, dict[str, Any]],
+    updates: dict[str, Any],
+    deletes: dict[str, Any] = None,
 ) -> str:
    """
-    Batch modify configurations for multiple stages in a YAML file.
+    Modify configurations in a YAML file, supporting both top-level and stage-specific modifications,
+    including addition, modification, and deletion of configurations.

    Args:
        yaml_path: Path to the YAML configuration file.
-        stage_updates: Dictionary where keys are stage IDs and values are dictionaries of
-                      modifications for that stage. Each modification dictionary uses
-                      dot-separated paths as keys and new configuration values as values.
-                      Example: {
-                          0: {'engine_args.max_model_len': 5800},
-                          1: {'runtime.max_batch_size': 2}
-                      }
+        updates: Dictionary containing both top-level and stage-specific modifications to add or update.
+                Format: {
+                    'async_chunk': True,
+                    'stage_args': {
+                        0: {'engine_args.max_model_len': 5800},
+                        1: {'runtime.max_batch_size': 2}
+                    }
+                }
+        deletes: Dictionary containing configurations to delete.
+                Format: {
+                    'old_config': None,  # Delete entire key
+                    'stage_args': {
+                        0: ['engine_args.old_param'],
+                        1: ['runtime.unused_setting']
+                    }
+                }

    Returns:
        str: Path to the newly created modified YAML file with timestamp suffix.
-
-    Example:
-        >>> output_file = modify_stage_config(
-        ...     'config.yaml',
-        ...     {
-        ...         0: {'engine_args.max_model_len': 5800},
-        ...         1: {'runtime.max_batch_size': 2}
-        ...     }
-        ... )
-        >>> print(f"Modified configuration saved to: {output_file}")
-        Modified configuration saved to: config_1698765432.yaml
    """
    path = Path(yaml_path)
    if not path.exists():
        raise FileNotFoundError(f"yaml does not exist: {path}")
+
    try:
        with open(yaml_path, encoding="utf-8") as f:
            config = yaml.safe_load(f) or {}
    except Exception as e:
        raise ValueError(f"Cannot parse YAML file: {e}")

-    stage_args = config.get("stage_args", [])
-    if not stage_args:
-        raise ValueError("the stage_args does not exist")
+    # Helper function to apply update
+    def apply_update(config_dict: dict, key_path: str, value: Any) -> None:
+        """Apply update to dictionary using dot-separated path."""
+        # Handle direct list assignment (e.g., engine_input_source: [1, 2])
+        if "." not in key_path:
+            # Simple key, set directly
+            config_dict[key_path] = value
+            return

-    for stage_id, config_dict in stage_updates.items():
-        target_stage = None
-        for stage in stage_args:
-            if stage.get("stage_id") == stage_id:
-                target_stage = stage
-                break
+        current = config_dict
+        keys = key_path.split(".")

-        if target_stage is None:
-            available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
-            raise KeyError(f"Stage ID {stage_id} is not exist, available IDs: {available_ids}")
+        for i in range(len(keys) - 1):
+            key = keys[i]

-        for key_path, value in config_dict.items():
-            current = target_stage
-            keys = key_path.split(".")
-            for i in range(len(keys) - 1):
-                key = keys[i]
+            # Handle list indices
+            if key.isdigit() and isinstance(current, list):
+                index = int(key)
+                if index < 0:
+                    raise ValueError(f"Negative list index not allowed: {index}")
+                if index >= len(current):
+                    # Expand list if needed
+                    while len(current) <= index:
+                        # If we need to go deeper (more keys after this), create a dict
+                        # Otherwise, create None placeholder
+                        current.append({} if i < len(keys) - 2 else None)
+                current = current[index]
+            elif isinstance(current, dict):
+                # Handle dictionary keys
                if key not in current:
-                    raise KeyError(f"the {'.'.join(keys[: i + 1])} does not exist")
-
-                elif not isinstance(current[key], dict) and i < len(keys) - 2:
-                    raise ValueError(f"{'.'.join(keys[: i + 1])}' cannot continue deeper because it's not a dict")
+                    # If there are more keys after this, create appropriate structure
+                    if i < len(keys) - 1:
+                        # Check if next key is a digit (list index) or string (dict key)
+                        if keys[i + 1].isdigit():
+                            current[key] = []
+                        else:
+                            current[key] = {}
+                    else:
+                        # This is the last key, create based on value type
+                        current[key] = [] if isinstance(value, list) else {}
+                elif not isinstance(current[key], (dict, list)) and i < len(keys) - 1:
+                    # If current value is not dict/list but we need to go deeper, replace it
+                    if keys[i + 1].isdigit():
+                        current[key] = []
+                    else:
+                        current[key] = {}
                current = current[key]
-            current[keys[-1]] = value
+            else:
+                # Current is not a dict or list, cannot traverse further
+                raise TypeError(
+                    f"Cannot access {'.'.join(keys[: i + 1])} as a dict/list. It's a {type(current).__name__}"
+                )
+
+        # Set the final value
+        last_key = keys[-1]
+        if isinstance(current, list) and last_key.isdigit():
+            # Setting a value in a list by index
+            index = int(last_key)
+            if index < 0:
+                raise ValueError(f"Negative list index not allowed: {index}")
+            if index >= len(current):
+                # Expand list if needed
+                while len(current) <= index:
+                    current.append(None)
+            current[index] = value
+        elif isinstance(current, dict):
+            # Special case: if the value is a list and we're setting a top-level key
+            # Example: updating engine_input_source with [1, 2]
+            current[last_key] = value
+        else:
+            # Current is not a dict, cannot set key
+            raise TypeError(f"Cannot set value at {key_path}. Current type is {type(current).__name__}, expected dict.")
+
+    # Helper function to delete by path
+    def delete_by_path(config_dict: dict, path: str) -> None:
+        """Delete configuration by dot-separated path."""
+        if not path:
+            return
+
+        current = config_dict
+        keys = path.split(".")
+
+        # Traverse to the parent
+        for i in range(len(keys) - 1):
+            key = keys[i]
+
+            # Handle list indices
+            if key.isdigit() and isinstance(current, list):
+                index = int(key)
+                if index < 0 or index >= len(current):
+                    raise KeyError(f"List index {index} out of bounds")
+                current = current[index]
+            elif isinstance(current, dict):
+                if key not in current:
+                    raise KeyError(f"Path {'.'.join(keys[: i + 1])} does not exist")
+                current = current[key]
+            else:
+                raise TypeError(
+                    f"Cannot access {'.'.join(keys[: i + 1])} as a dict/list. It's a {type(current).__name__}"
+                )
+
+        # Delete the item
+        last_key = keys[-1]
+
+        if isinstance(current, list) and last_key.isdigit():
+            index = int(last_key)
+            if index < 0 or index >= len(current):
+                raise KeyError(f"List index {index} out of bounds")
+            del current[index]
+        elif isinstance(current, dict) and last_key in current:
+            del current[last_key]
+        else:
+            raise KeyError(f"Path {path} does not exist")
+
+    # Apply deletions first
+    if deletes:
+        for key, value in deletes.items():
+            if key == "stage_args":
+                if value and isinstance(value, dict):
+                    stage_args = config.get("stage_args", [])
+                    if not stage_args:
+                        raise ValueError("stage_args does not exist in config")
+
+                    for stage_id, delete_paths in value.items():
+                        if not delete_paths:
+                            continue
+
+                        # Find stage by ID
+                        target_stage = None
+                        for stage in stage_args:
+                            if stage.get("stage_id") == stage_id:
+                                target_stage = stage
+                                break
+
+                        if target_stage is None:
+                            available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
+                            raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}")
+
+                        # Delete specified paths in this stage
+                        for path in delete_paths:
+                            if path:  # Skip empty paths
+                                delete_by_path(target_stage, path)
+            elif "." in key:
+                # Delete using dot-separated path
+                delete_by_path(config, key)
+            elif value is None and key in config:
+                # Delete entire key
+                del config[key]
+
+    # Apply updates
+    for key, value in updates.items():
+        if key == "stage_args":
+            if value and isinstance(value, dict):
+                stage_args = config.get("stage_args", [])
+                if not stage_args:
+                    raise ValueError("stage_args does not exist in config")
+
+                for stage_id, stage_updates in value.items():
+                    # Find stage by ID
+                    target_stage = None
+                    for stage in stage_args:
+                        if stage.get("stage_id") == stage_id:
+                            target_stage = stage
+                            break
+
+                    if target_stage is None:
+                        available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
+                        raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}")
+
+                    # Apply updates to this stage
+                    for path, val in stage_updates.items():
+                        # Check if this is a simple key (not dot-separated)
+                        # Example: 'engine_input_source' vs 'engine_args.max_model_len'
+                        if "." not in path:
+                            # Direct key assignment (e.g., updating a list value)
+                            target_stage[path] = val
+                        else:
+                            # Dot-separated path (e.g., nested dict access)
+                            apply_update(target_stage, path, val)
+        elif "." in key:
+            # Apply using dot-separated path
+            apply_update(config, key, value)
+        else:
+            # Direct top-level key
+            config[key] = value
+
+    # Save to new file with timestamp
+    timestamp = int(time.time())
+    base_name = yaml_path.rsplit(".", 1)[0] if "." in yaml_path else yaml_path
+    output_path = f"{base_name}_{timestamp}.yaml"

-    output_path = f"{yaml_path.split('.')[0]}_{int(time.time())}.yaml"
    with open(output_path, "w", encoding="utf-8") as f:
-        yaml.dump(config, f, default_flow_style=False, sort_keys=False, allow_unicode=True, indent=2)
+        yaml.dump(config, f, default_flow_style=None, sort_keys=False, allow_unicode=True, indent=2)

    return output_path

@@ -621,6 +882,9 @@ 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
@@ -656,7 +920,7 @@ class OmniServer:
        )

        # Wait for server to be ready
-        max_wait = 600  # 10 minutes
+        max_wait = 1200  # 20 minutes
        start_time = time.time()
        while time.time() - start_time < max_wait:
            try:
@@ -673,24 +937,32 @@ class OmniServer:
        raise RuntimeError(f"Server failed to start within {max_wait} seconds")

    def _kill_process_tree(self, pid):
-        """kill process and its children"""
+        """kill process and its children with verification"""
        try:
            parent = psutil.Process(pid)
            children = parent.children(recursive=True)
+
+            # Get all PIDs first
+            all_pids = [pid] + [child.pid for child in children]
+
+            # Terminate children
            for child in children:
                try:
                    child.terminate()
                except psutil.NoSuchProcess:
                    pass

+            # Wait for children
            gone, still_alive = psutil.wait_procs(children, timeout=10)

+            # Kill remaining children
            for child in still_alive:
                try:
                    child.kill()
                except psutil.NoSuchProcess:
                    pass

+            # Terminate parent
            try:
                parent.terminate()
                parent.wait(timeout=10)
@@ -700,6 +972,24 @@ class OmniServer:
                except psutil.NoSuchProcess:
                    pass

+            # VERIFICATION: Check if all processes are gone
+            time.sleep(1)  # Give system time
+            alive_processes = []
+            for check_pid in all_pids:
+                if psutil.pid_exists(check_pid):
+                    alive_processes.append(check_pid)
+
+            if alive_processes:
+                print(f"Warning: Processes still alive: {alive_processes}")
+                # Optional: Try system kill
+                import subprocess
+
+                for alive_pid in alive_processes:
+                    try:
+                        subprocess.run(["kill", "-9", str(alive_pid)], timeout=2)
+                    except Exception as e:
+                        print(f"Cleanup failed: {e}")
+
        except psutil.NoSuchProcess:
            pass

@@ -709,31 +999,7 @@ class OmniServer:

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.proc:
-            try:
-                parent = psutil.Process(self.proc.pid)
-                children = parent.children(recursive=True)
-                for child in children:
-                    try:
-                        child.terminate()
-                    except psutil.NoSuchProcess:
-                        pass
-
-                gone, still_alive = psutil.wait_procs(children, timeout=10)
-
-                for child in still_alive:
-                    try:
-                        child.kill()
-                    except psutil.NoSuchProcess:
-                        pass
-
-                try:
-                    parent.terminate()
-                    parent.wait(timeout=10)
-                except (psutil.NoSuchProcess, psutil.TimeoutExpired):
-                    try:
-                        parent.kill()
-                    except psutil.NoSuchProcess:
-                        pass
-
-            except psutil.NoSuchProcess:
-                pass
+            self._kill_process_tree(self.proc.pid)
+        _run_pre_test_cleanup(enable_force=True)
+        _run_post_test_cleanup(enable_force=True)
+        cleanup_dist_env_and_memory()
@@ -6,7 +6,6 @@ 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 (
@@ -19,15 +18,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
    initialize_model_parallel,
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
-from vllm_omni.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.")
+from vllm_omni.platforms import current_omni_platform


 def update_environment_variables(envs_dict: dict[str, str]):
@@ -171,7 +162,7 @@ def test_sequence_parallel(
    sequence_parallel_size = ulysses_degree * ring_degree

    # Skip if not enough GPUs available
-    available_gpus = torch_device.device_count()
+    available_gpus = current_omni_platform.get_device_count()
    if available_gpus < sequence_parallel_size:
        pytest.skip(f"Test requires {sequence_parallel_size} GPUs but only {available_gpus} available")

@@ -338,13 +329,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_platform.seed_everything(RANDOM_SEED)
+    current_omni_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"{device_type}:{local_rank}")
-    torch_device.set_device(device)
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
    torch.set_default_device(device)
    torch.set_default_dtype(dtype)

@@ -423,7 +414,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)
-            torch_device.manual_seed_all(42)
+            current_omni_platform.seed_everything(42)
            full_hidden_states = torch.randn(
                (batch_size, seq_len, hidden_size),
                dtype=dtype,
@@ -0,0 +1,423 @@
+# 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})")
@@ -14,15 +14,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
    init_distributed_environment,
    initialize_model_parallel,
 )
-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.")
+from vllm_omni.platforms import current_omni_platform


 def update_environment_variables(envs_dict: dict[str, str]):
@@ -49,7 +41,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 = torch_device.device_count()
+    available_gpus = current_omni_platform.get_device_count()
    if available_gpus < world_size:
        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")

@@ -85,8 +77,8 @@ def _test_4d_identity_worker(
 ):
    """Worker function for test_4d_identity."""
    # Set device
-    device = torch.device(f"{device_type}:{local_rank}")
-    torch_device.set_device(device)
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)

    # Set environment variables for distributed training
    update_environment_variables(
@@ -183,7 +175,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 = torch_device.device_count()
+    available_gpus = current_omni_platform.get_device_count()
    if available_gpus < world_size:
        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")

@@ -219,8 +211,8 @@ def _test_5d_identity_worker(
 ):
    """Worker function for test_5d_identity."""
    # Set device
-    device = torch.device(f"{device_type}:{local_rank}")
-    torch_device.set_device(device)
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)

    # Set environment variables for distributed training
    update_environment_variables(
@@ -316,7 +308,7 @@ def test_ring_p2p(
 ):
    """Test Ring P2P communication (send_recv)."""
    # Skip if not enough GPUs available
-    available_gpus = torch_device.device_count()
+    available_gpus = current_omni_platform.get_device_count()
    if available_gpus < world_size:
        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")

@@ -339,8 +331,8 @@ def _test_ring_p2p_worker(
    import sys

    # Set device
-    device = torch.device(f"{device_type}:{local_rank}")
-    torch_device.set_device(device)
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)

    # Set env vars
    # Use a different port to avoid conflict with other tests if run in parallel
@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from vllm_omni.diffusion.lora.layers.base_linear import DiffusionBaseLinearLayerWithLoRA
+
+
+@dataclass
+class _DummyLoRAConfig:
+    fully_sharded_loras: bool = False
+
+
+class _DummyQuantMethod:
+    def __init__(self, weight: torch.Tensor):
+        self._weight = weight
+
+    def apply(self, _base_layer, x: torch.Tensor, bias: torch.Tensor | None):
+        y = x @ self._weight.t()
+        if bias is not None:
+            y = y + bias
+        return y
+
+
+def test_diffusion_base_linear_apply_multi_slice():
+    # Build a fake diffusion LoRA layer with 2 slices and rank=2.
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    layer.tp_size = 1
+    layer.lora_config = _DummyLoRAConfig()
+
+    in_dim = 3
+    out_slices = (2, 1)
+    rank = 2
+
+    # Base weight: identity-ish mapping to make base output easy to reason about.
+    base_weight = torch.tensor(
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+    layer.base_layer = type("Base", (), {})()
+    layer.base_layer.quant_method = _DummyQuantMethod(base_weight)
+
+    # Allocate stacked weights: (max_loras=1, 1, rank, in_dim) and (1, 1, out_slice, rank)
+    a0 = torch.zeros((1, 1, rank, in_dim))
+    b0 = torch.zeros((1, 1, out_slices[0], rank))
+    a1 = torch.zeros((1, 1, rank, in_dim))
+    b1 = torch.zeros((1, 1, out_slices[1], rank))
+
+    # Slice 0: delta0 = (x @ A0.T) @ B0.T
+    A0 = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])  # (2, 3)
+    B0 = torch.tensor([[1.0, 0.0], [0.0, 1.0]])  # (2, 2)
+    a0[0, 0, :, :] = A0
+    b0[0, 0, :, :] = B0
+
+    # Slice 1: delta1 = (x @ A1.T) @ B1.T
+    A1 = torch.tensor([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]])  # (2, 3)
+    B1 = torch.tensor([[2.0, 0.0]])  # (1, 2)
+    a1[0, 0, :, :] = A1
+    b1[0, 0, :, :] = B1
+
+    layer.lora_a_stacked = (a0, a1)
+    layer.lora_b_stacked = (b0, b1)
+    layer.output_slices = out_slices
+
+    x = torch.tensor([[1.0, 2.0, 3.0]])
+    out = layer.apply(x)
+
+    # Base output is identity: [1,2,3]
+    base_out = x @ base_weight.t()
+    # delta0:
+    # (x @ A0.T) = [1,2]
+    # [1,2] @ B0.T = [1,2]
+    delta0 = torch.tensor([[1.0, 2.0]])
+    # delta1:
+    # (x @ A1.T) = [3,1]
+    # [3,1] @ B1.T = [6]
+    delta1 = torch.tensor([[6.0]])
+    expected = torch.cat([base_out[:, :2] + delta0, base_out[:, 2:3] + delta1], dim=-1)
+    assert torch.allclose(out, expected)
+
+
+def test_diffusion_base_linear_reset_lora_disables_fast_path(monkeypatch):
+    # Verify that after reset_lora(), apply() skips LoRA matmuls even if the
+    # LoRA tensors are still allocated and non-empty.
+    from vllm.lora.layers.base_linear import BaseLinearLayerWithLoRA
+
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    layer.tp_size = 1
+    layer.lora_config = _DummyLoRAConfig()
+
+    in_dim = 2
+    out_dim = 2
+    rank = 1
+
+    base_weight = torch.eye(in_dim)
+    layer.base_layer = type("Base", (), {})()
+    layer.base_layer.quant_method = _DummyQuantMethod(base_weight)
+
+    a = torch.ones((1, 1, rank, in_dim))
+    b = torch.tensor([[[[1.0], [2.0]]]])  # (1,1,out_dim,rank)
+
+    layer.lora_a_stacked = (a,)
+    layer.lora_b_stacked = (b,)
+    layer.output_slices = (out_dim,)
+    layer._diffusion_lora_active_slices = (True,)
+
+    x = torch.tensor([[1.0, 2.0]])
+    out_active = layer.apply(x)
+    assert torch.allclose(out_active, torch.tensor([[4.0, 8.0]]))
+
+    monkeypatch.setattr(BaseLinearLayerWithLoRA, "reset_lora", lambda self, index: None)
+    layer.reset_lora(0)
+
+    assert layer._diffusion_lora_active_slices == (False,)
+    out_inactive = layer.apply(x)
+    assert torch.allclose(out_inactive, x)
+
+
+def test_diffusion_base_linear_apply_respects_inactive_slices():
+    # Build a fake diffusion LoRA layer with 2 slices and rank=2.
+    layer = DiffusionBaseLinearLayerWithLoRA.__new__(DiffusionBaseLinearLayerWithLoRA)
+    layer.tp_size = 1
+    layer.lora_config = _DummyLoRAConfig()
+
+    in_dim = 3
+    out_slices = (2, 1)
+    rank = 2
+
+    base_weight = torch.tensor(
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+    )
+    layer.base_layer = type("Base", (), {})()
+    layer.base_layer.quant_method = _DummyQuantMethod(base_weight)
+
+    a0 = torch.zeros((1, 1, rank, in_dim))
+    b0 = torch.zeros((1, 1, out_slices[0], rank))
+    a1 = torch.zeros((1, 1, rank, in_dim))
+    b1 = torch.zeros((1, 1, out_slices[1], rank))
+
+    A0 = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])  # (2, 3)
+    B0 = torch.tensor([[1.0, 0.0], [0.0, 1.0]])  # (2, 2)
+    a0[0, 0, :, :] = A0
+    b0[0, 0, :, :] = B0
+
+    A1 = torch.tensor([[0.0, 0.0, 1.0], [1.0, 0.0, 0.0]])  # (2, 3)
+    B1 = torch.tensor([[2.0, 0.0]])  # (1, 2)
+    a1[0, 0, :, :] = A1
+    b1[0, 0, :, :] = B1
+
+    layer.lora_a_stacked = (a0, a1)
+    layer.lora_b_stacked = (b0, b1)
+    layer.output_slices = out_slices
+    layer._diffusion_lora_active_slices = (True, False)
+
+    x = torch.tensor([[1.0, 2.0, 3.0]])
+    out = layer.apply(x)
+
+    # Only the first slice should be adapted.
+    expected = torch.tensor([[2.0, 4.0, 3.0]])
+    assert torch.allclose(out, expected)
@@ -0,0 +1,343 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import torch
+from vllm.lora.lora_weights import LoRALayerWeights
+from vllm.lora.utils import get_supported_lora_modules
+from vllm.model_executor.layers.linear import LinearBase
+
+from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
+from vllm_omni.lora.request import LoRARequest
+
+
+class _DummyLoRALayer:
+    def __init__(self, n_slices: int, output_slices: tuple[int, ...]):
+        self.n_slices = n_slices
+        self.output_slices = output_slices
+        self.set_calls: list[
+            tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
+        ] = []
+        self.reset_calls: int = 0
+
+    def set_lora(self, index: int, lora_a, lora_b):
+        assert index == 0
+        self.set_calls.append((lora_a, lora_b))
+
+    def reset_lora(self, index: int):
+        assert index == 0
+        self.reset_calls += 1
+
+
+class _FakeLinearBase(LinearBase):
+    def __init__(self):
+        torch.nn.Module.__init__(self)
+
+
+def test_lora_manager_supported_modules_are_stable_with_wrapped_layers(monkeypatch):
+    # Simulate a pipeline that already contains LoRA wrappers where the original
+    # LinearBase is nested under ".base_layer".
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    class _DummyBaseLayerWithLoRA(torch.nn.Module):
+        def __init__(self, base_layer: torch.nn.Module):
+            super().__init__()
+            self.base_layer = base_layer
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    pipeline = torch.nn.Module()
+    pipeline.transformer = torch.nn.Module()
+    pipeline.transformer.foo = _DummyBaseLayerWithLoRA(_FakeLinearBase())
+
+    # vLLM helper would see only the nested LinearBase and yield "base_layer".
+    assert get_supported_lora_modules(pipeline) == ["base_layer"]
+
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    assert "foo" in manager._supported_lora_modules
+    assert "base_layer" not in manager._supported_lora_modules
+
+
+def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    class _DummyBaseLayerWithLoRA(torch.nn.Module):
+        def __init__(self, base_layer: torch.nn.Module):
+            super().__init__()
+            self.base_layer = base_layer
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
+        if isinstance(layer, _FakeLinearBase):
+            return _DummyBaseLayerWithLoRA(layer)
+        return layer
+
+    replace_calls: list[str] = []
+
+    def _fake_replace_submodule(root: torch.nn.Module, module_name: str, submodule: torch.nn.Module):
+        replace_calls.append(module_name)
+        setattr(root, module_name, submodule)
+
+    monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
+    monkeypatch.setattr(manager_mod, "replace_submodule", _fake_replace_submodule)
+
+    pipeline = torch.nn.Module()
+    pipeline.transformer = torch.nn.Module()
+    pipeline.transformer.foo = _FakeLinearBase()
+
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    peft_helper = type("_PH", (), {"r": 1})()
+
+    manager._replace_layers_with_lora(peft_helper)
+    manager._replace_layers_with_lora(peft_helper)
+
+    # Only the top-level layer should have been replaced; nested ".base_layer"
+    # must be skipped to avoid nesting LoRA wrappers.
+    assert replace_calls == ["foo"]
+
+
+def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    class _DummyBaseLayerWithLoRA(torch.nn.Module):
+        def __init__(self, base_layer: torch.nn.Module):
+            super().__init__()
+            self.base_layer = base_layer
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
+        return _DummyBaseLayerWithLoRA(layer)
+
+    replace_calls: list[str] = []
+
+    def _fake_replace_submodule(root: torch.nn.Module, module_name: str, submodule: torch.nn.Module):
+        replace_calls.append(module_name)
+        setattr(root, module_name, submodule)
+
+    monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion)
+    monkeypatch.setattr(manager_mod, "replace_submodule", _fake_replace_submodule)
+
+    pipeline = torch.nn.Module()
+    pipeline.packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+    pipeline.transformer = torch.nn.Module()
+    pipeline.transformer.to_qkv = _FakeLinearBase()
+
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    # Treat the dummy layer as a packed 3-slice projection so the manager uses
+    # `packed_modules_mapping` to decide replacement based on target_modules.
+    monkeypatch.setattr(manager, "_get_packed_modules_list", lambda _module: ["q", "k", "v"])
+
+    peft_helper = type("_PH", (), {"r": 1, "target_modules": ["to_q"]})()
+    manager._replace_layers_with_lora(peft_helper)
+
+    assert replace_calls == ["to_qkv"]
+
+
+def test_lora_manager_activates_fused_lora_on_packed_layer():
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    packed_layer = _DummyLoRALayer(n_slices=3, output_slices=(2, 1, 1))
+    manager._lora_modules = {"transformer.blocks.0.attn.to_qkv": packed_layer}
+
+    rank = 2
+    A = torch.ones((rank, 4))
+    B = torch.arange(0, sum(packed_layer.output_slices) * rank, dtype=torch.bfloat16).view(-1, rank)
+    lora = LoRALayerWeights(
+        module_name="transformer.blocks.0.attn.to_qkv",
+        rank=rank,
+        lora_alpha=rank,
+        lora_a=A,
+        lora_b=B,
+    )
+    manager._registered_adapters = {
+        7: type(
+            "LM",
+            (),
+            {
+                "id": 7,
+                "loras": {"transformer.blocks.0.attn.to_qkv": lora},
+                "get_lora": lambda self, k: self.loras.get(k),
+            },
+        )()
+    }
+    manager._adapter_scales = {7: 0.5}
+
+    manager._activate_adapter(7)
+
+    assert packed_layer.reset_calls == 0
+    assert len(packed_layer.set_calls) == 1
+    lora_a_list, lora_b_list = packed_layer.set_calls[0]
+    assert isinstance(lora_a_list, list)
+    assert isinstance(lora_b_list, list)
+    assert len(lora_a_list) == 3
+    assert len(lora_b_list) == 3
+    assert all(torch.allclose(a, A) for a in lora_a_list)
+    # B should be split into 3 slices and scaled.
+    b0, b1, b2 = lora_b_list
+    assert b0.shape[0] == 2 and b1.shape[0] == 1 and b2.shape[0] == 1
+    assert torch.allclose(torch.cat([b0, b1, b2], dim=0), B * 0.5)
+
+
+def test_lora_manager_activates_packed_lora_from_sublayers():
+    pipeline = torch.nn.Module()
+    pipeline.packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
+    manager = DiffusionLoRAManager(
+        pipeline=pipeline,
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    packed_layer = _DummyLoRALayer(n_slices=3, output_slices=(2, 1, 1))
+    manager._lora_modules = {"transformer.blocks.0.attn.to_qkv": packed_layer}
+
+    rank = 2
+    loras: dict[str, LoRALayerWeights] = {}
+    for name, out_dim in zip(["to_q", "to_k", "to_v"], [2, 1, 1]):
+        loras[f"transformer.blocks.0.attn.{name}"] = LoRALayerWeights(
+            module_name=f"transformer.blocks.0.attn.{name}",
+            rank=rank,
+            lora_alpha=rank,
+            lora_a=torch.ones((rank, 4)) * (1 if name == "to_q" else 2),
+            lora_b=torch.ones((out_dim, rank)) * (3 if name == "to_q" else 4),
+        )
+
+    manager._registered_adapters = {
+        1: type("LM", (), {"id": 1, "loras": loras, "get_lora": lambda self, k: self.loras.get(k)})()
+    }
+    manager._adapter_scales = {1: 2.0}
+
+    manager._activate_adapter(1)
+
+    assert packed_layer.reset_calls == 0
+    assert len(packed_layer.set_calls) == 1
+    lora_a_list, lora_b_list = packed_layer.set_calls[0]
+    assert isinstance(lora_a_list, list)
+    assert isinstance(lora_b_list, list)
+    assert len(lora_a_list) == 3
+    assert len(lora_b_list) == 3
+    # Scale should apply to B only.
+    assert torch.allclose(lora_b_list[0], torch.ones((2, rank)) * 3 * 2.0)
+    assert torch.allclose(lora_b_list[1], torch.ones((1, rank)) * 4 * 2.0)
+    assert torch.allclose(lora_b_list[2], torch.ones((1, rank)) * 4 * 2.0)
+
+
+def _dummy_lora_request(adapter_id: int) -> LoRARequest:
+    return LoRARequest(
+        lora_name=f"adapter_{adapter_id}",
+        lora_int_id=adapter_id,
+        lora_path=f"/tmp/adapter_{adapter_id}",
+    )
+
+
+def test_lora_manager_evicts_lru_adapter_when_cache_full(monkeypatch):
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=2,
+    )
+
+    def _fake_load(_req: LoRARequest):
+        lora_model = type("LM", (), {"id": _req.lora_int_id})()
+        peft_helper = type("PH", (), {})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+
+    req1 = _dummy_lora_request(1)
+    req2 = _dummy_lora_request(2)
+    req3 = _dummy_lora_request(3)
+
+    manager.set_active_adapter(req1, lora_scale=1.0)
+    manager.set_active_adapter(req2, lora_scale=1.0)
+
+    # Touch adapter 1 so adapter 2 becomes LRU.
+    manager.set_active_adapter(req1, lora_scale=1.0)
+
+    manager.set_active_adapter(req3, lora_scale=1.0)
+
+    assert set(manager.list_adapters()) == {1, 3}
+
+
+def test_lora_manager_does_not_evict_pinned_adapter(monkeypatch):
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=2,
+    )
+
+    def _fake_load(_req: LoRARequest):
+        lora_model = type("LM", (), {"id": _req.lora_int_id})()
+        peft_helper = type("PH", (), {})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+
+    manager.set_active_adapter(_dummy_lora_request(1), lora_scale=1.0)
+    assert manager.pin_adapter(1)
+
+    manager.set_active_adapter(_dummy_lora_request(2), lora_scale=1.0)
+    manager.set_active_adapter(_dummy_lora_request(3), lora_scale=1.0)
+
+    assert set(manager.list_adapters()) == {1, 3}
+
+
+def test_lora_manager_warns_when_all_adapters_pinned(monkeypatch):
+    manager = DiffusionLoRAManager(
+        pipeline=torch.nn.Module(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=2,
+    )
+
+    def _fake_load(_req: LoRARequest):
+        lora_model = type("LM", (), {"id": _req.lora_int_id})()
+        peft_helper = type("PH", (), {})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+
+    manager.set_active_adapter(_dummy_lora_request(1), lora_scale=1.0)
+    manager.set_active_adapter(_dummy_lora_request(2), lora_scale=1.0)
+
+    assert manager.pin_adapter(1)
+    assert manager.pin_adapter(2)
+
+    manager.max_cached_adapters = 1
+    manager._evict_if_needed()
+
+    assert set(manager.list_adapters()) == {1, 2}
@@ -2,9 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 """
-Unit tests for GPUDiffusionWorker class.
+Unit tests for DiffusionWorker class.

-This module tests the GPUDiffusionWorker implementation:
+This module tests the DiffusionWorker 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.gpu_diffusion_worker import GPUDiffusionWorker
+from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker


@pytest.fixture
@@ -33,18 +33,19 @@ def mock_od_config():

@pytest.fixture
 def mock_gpu_worker(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)
+    """Create a DiffusionWorker with mocked initialization."""
+    with patch.object(DiffusionWorker, "init_device"):
+        worker = DiffusionWorker(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 TestGPUDiffusionWorkerLoadWeights:
-    """Test GPUDiffusionWorker.load_weights method."""
+class TestDiffusionWorkerLoadWeights:
+    """Test DiffusionWorker.load_weights method."""

    def test_load_weights_calls_pipeline(self, mock_gpu_worker):
        """Test that load_weights delegates to model_runner.load_weights."""
@@ -75,20 +76,21 @@ class TestGPUDiffusionWorkerLoadWeights:
        assert result == set()


-class TestGPUDiffusionWorkerSleep:
-    """Test GPUDiffusionWorker.sleep method."""
+class TestDiffusionWorkerSleep:
+    """Test DiffusionWorker.sleep method."""

-    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
    @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_level_1(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+    def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_gpu_worker):
        """Test sleep mode level 1 (offload weights only)."""
        # Setup memory info mocks
-        # 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
+        # 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
        ]
+        mock_platform.get_device_total_memory.return_value = 8 * 1024**3

        # Setup allocator mock
        mock_allocator = Mock()
@@ -104,15 +106,16 @@ class TestGPUDiffusionWorkerSleep:
        # 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.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
    @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_level_2(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+    def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_gpu_worker):
        """Test sleep mode level 2 (offload all, save buffers)."""
        # Setup memory info mocks
-        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_free_memory.side_effect = [
+            1 * 1024**3,  # Before sleep
+            5 * 1024**3,  # After sleep (freed 4GB)
        ]
+        mock_platform.get_device_total_memory.return_value = 8 * 1024**3

        # Setup allocator mock
        mock_allocator = Mock()
@@ -141,15 +144,16 @@ class TestGPUDiffusionWorkerSleep:
        assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
        assert "buffer2" in mock_gpu_worker._sleep_saved_buffers

-    @patch("vllm_omni.diffusion.worker.gpu_diffusion_worker.torch.cuda.mem_get_info")
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
    @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_mem_info, mock_gpu_worker):
+    def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_platform, mock_gpu_worker):
        """Test that sleep validates memory was actually freed."""
        # Simulate memory increase (should trigger assertion error)
-        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_free_memory.side_effect = [
+            3 * 1024**3,  # Before sleep: 3GB free
+            1 * 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)
@@ -160,8 +164,8 @@ class TestGPUDiffusionWorkerSleep:
            mock_gpu_worker.sleep(level=1)


-class TestGPUDiffusionWorkerWakeUp:
-    """Test GPUDiffusionWorker.wake_up method."""
+class TestDiffusionWorkerWakeUp:
+    """Test DiffusionWorker.wake_up method."""

    @patch("vllm.device_allocator.cumem.CuMemAllocator")
    def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker):
@@ -1,189 +1,251 @@
-import unittest
-from unittest.mock import MagicMock, patch
-
+import pytest
 import torch

-from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline
+from tests.utils import hardware_test
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    OmniKVCacheConfig,
+    OmniKVTransferManager,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams


-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."""
-
+class MockConnector:
    def __init__(self):
-        self.device = "cpu"
-        self.od_config = MagicMock()
-        self.bagel = MagicMock()
-        self.tokenizer = MagicMock()
-        self.language_model = MagicMock()
-        self.new_token_ids = {}
+        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


-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
-        req = OmniDiffusionRequest(prompt="test")
-        from types import SimpleNamespace
-
-        req.past_key_values = SimpleNamespace(**layer_blocks)
-        req.kv_metadata = transfer_data.metadata
-
-        # 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.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)
-        req = OmniDiffusionRequest(prompt="integration_test")
-        req.past_key_values = data_dict["layer_blocks"]
-        req.kv_metadata = data_dict["metadata"]
-
-        # 4. Receiver (Injection Simulation)
-        # Use the logic verification again
-        pass
+@pytest.fixture
+def mock_connector():
+    return MockConnector()


-if __name__ == "__main__":
-    unittest.main()
+@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
@@ -7,10 +7,12 @@ Pytest configuration and fixtures for vllm-omni tests.
 from typing import Any

 import pytest
+from vllm import TextPrompt
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
-from vllm.sampling_params import SamplingParams

+from tests.conftest import _run_post_test_cleanup, _run_pre_test_cleanup
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

 PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
@@ -49,6 +51,9 @@ 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

@@ -63,7 +68,7 @@ class OmniRunner:
            **kwargs,
        )

-    def get_default_sampling_params_list(self) -> list[SamplingParams]:
+    def get_default_sampling_params_list(self) -> list[OmniSamplingParams]:
        """
        Get a list of default sampling parameters for all stages.

@@ -81,7 +86,7 @@ class OmniRunner:
        videos: PromptVideoInput = None,
        mm_processor_kwargs: dict[str, Any] | None = None,
        modalities: list[str] | None = None,
-    ) -> list[dict[str, Any]]:
+    ) -> list[TextPrompt]:
        """
        Construct Omni input format from prompts and multimodal data.

@@ -174,7 +179,7 @@ class OmniRunner:
                f"<|im_start|>assistant\n"
            )

-            input_dict: dict[str, Any] = {"prompt": full_prompt}
+            input_dict: TextPrompt = {"prompt": full_prompt}
            if multi_modal_data:
                input_dict["multi_modal_data"] = multi_modal_data
            if modalities:
@@ -188,8 +193,8 @@ class OmniRunner:

    def generate(
        self,
-        prompts: list[dict[str, Any]],
-        sampling_params_list: list[SamplingParams] | None = None,
+        prompts: list[TextPrompt],
+        sampling_params_list: list[OmniSamplingParams] | None = None,
    ) -> list[OmniRequestOutput]:
        """
        Generate outputs for the given prompts.
@@ -211,7 +216,7 @@ class OmniRunner:
    def generate_multimodal(
        self,
        prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
+        sampling_params_list: list[OmniSamplingParams] | None = None,
        system_prompt: str | None = None,
        audios: PromptAudioInput = None,
        images: PromptImageInput = None,
@@ -248,7 +253,7 @@ class OmniRunner:
    def generate_audio(
        self,
        prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
+        sampling_params_list: list[OmniSamplingParams] | None = None,
        system_prompt: str | None = None,
        audios: PromptAudioInput = None,
        mm_processor_kwargs: dict[str, Any] | None = None,
@@ -275,7 +280,7 @@ class OmniRunner:
    def generate_video(
        self,
        prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
+        sampling_params_list: list[OmniSamplingParams] | None = None,
        system_prompt: str | None = None,
        videos: PromptVideoInput = None,
        mm_processor_kwargs: dict[str, Any] | None = None,
@@ -302,7 +307,7 @@ class OmniRunner:
    def generate_image(
        self,
        prompts: list[str] | str,
-        sampling_params_list: list[SamplingParams] | None = None,
+        sampling_params_list: list[OmniSamplingParams] | None = None,
        system_prompt: str | None = None,
        images: PromptImageInput = None,
        mm_processor_kwargs: dict[str, Any] | None = None,
@@ -335,6 +340,7 @@ class OmniRunner:
        self.close()
        del self.omni
        cleanup_dist_env_and_memory()
+        _run_post_test_cleanup(enable_force=True)

    def close(self):
        """Close and cleanup the Omni instance."""
@@ -0,0 +1,85 @@
+# 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
@@ -0,0 +1,83 @@
+# 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
@@ -10,7 +10,7 @@ stage_args:
    engine_args:
      model_stage: thinker
      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      worker_type: ar
      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_cls: vllm_omni.worker.npu.npu_ar_worker.NPUARWorker
+      worker_type: ar
      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_cls: vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker
+      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      gpu_memory_utilization: 0.15
      enforce_eager: true
@@ -11,7 +11,7 @@ stage_args:
    engine_args:
      model_stage: thinker
      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      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_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      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_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      gpu_memory_utilization: 0.15
      enforce_eager: true
@@ -7,21 +7,21 @@
 stage_args:
  - stage_id: 0
    runtime:
-      devices: "0,1"
+      devices: "0"
      max_batch_size: 1
    engine_args:
      model_stage: thinker
      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
+      gpu_memory_utilization: 0.9
      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: 2
+      tensor_parallel_size: 1
      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_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.3
+       gpu_memory_utilization: 0.6
       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: "0"
+      devices: "1"
      max_batch_size: 1
    engine_args:
      model_stage: code2wav
      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      enforce_eager: true
      trust_remote_code: true
@@ -11,7 +11,7 @@ stage_args:
    engine_args:
      model_stage: thinker
      model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      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_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      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_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      gpu_memory_utilization: 0.15
      enforce_eager: true
@@ -7,21 +7,21 @@
 stage_args:
  - stage_id: 0
    runtime:
-      devices: "0,1"
+      devices: "0"
      max_batch_size: 1
    engine_args:
      model_stage: thinker
      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.6
+      gpu_memory_utilization: 0.9
      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: 2
+      tensor_parallel_size: 1
      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_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-       gpu_memory_utilization: 0.3
+       gpu_memory_utilization: 0.6
       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: "0"
+      devices: "1"
      max_batch_size: 1
    engine_args:
      model_stage: code2wav
      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
      enforce_eager: true
      trust_remote_code: true
@@ -0,0 +1,308 @@
+# 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
@@ -15,6 +15,8 @@ from pathlib import Path
 import pytest
 import torch

+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[2]
 if str(REPO_ROOT) not in sys.path:
@@ -40,42 +42,51 @@ def test_cache_dit(model_name: str):
        "residual_diff_threshold": 0.24,
        "max_continuous_cached_steps": 3,
    }
+    m = None
+    try:
+        m = Omni(
+            model=model_name,
+            cache_backend="cache_dit",
+            cache_config=cache_config,
+        )

-    m = Omni(
-        model=model_name,
-        cache_backend="cache_dit",
-        cache_config=cache_config,
-    )
+        # Use minimal settings for fast testing
+        height = 256
+        width = 256
+        num_inference_steps = 4  # Minimal steps for fast test

-    # Use minimal settings for fast testing
-    height = 256
-    width = 256
-    num_inference_steps = 4  # Minimal steps for fast test
+        outputs = m.generate(
+            "a photo of a cat sitting on a laptop keyboard",
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=0.0,
+                generator=torch.Generator("cuda").manual_seed(42),
+                num_outputs_per_prompt=1,  # Single output for speed
+            ),
+        )
+        # Extract images from request_output[0]['images']
+        first_output = outputs[0]
+        assert first_output.final_output_type == "image"
+        if not hasattr(first_output, "request_output") or not first_output.request_output:
+            raise ValueError("No request_output found in OmniRequestOutput")

-    outputs = m.generate(
-        "a photo of a cat sitting on a laptop keyboard",
-        height=height,
-        width=width,
-        num_inference_steps=num_inference_steps,
-        guidance_scale=0.0,
-        generator=torch.Generator("cuda").manual_seed(42),
-        num_outputs_per_prompt=1,  # Single output for speed
-    )
-    # Extract images from request_output[0]['images']
-    first_output = outputs[0]
-    assert first_output.final_output_type == "image"
-    if not hasattr(first_output, "request_output") or not first_output.request_output:
-        raise ValueError("No request_output found in OmniRequestOutput")
+        req_out = first_output.request_output[0]
+        if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+            raise ValueError("Invalid request_output structure or missing 'images' key")

-    req_out = first_output.request_output[0]
-    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
-        raise ValueError("Invalid request_output structure or missing 'images' key")
+        images = req_out.images

-    images = req_out.images
-
-    # Verify generation succeeded
-    assert images is not None
-    assert len(images) == 1
-    # Check image size
-    assert images[0].width == width
-    assert images[0].height == height
+        # Verify generation succeeded
+        assert images is not None
+        assert len(images) == 1
+        # Check image size
+        assert images[0].width == width
+        assert images[0].height == height
+    except Exception as e:
+        print(f"Test failed with error: {e}")
+        raise
+    finally:
+        if m is not None and hasattr(m, "close"):
+            m.close()
@@ -3,9 +3,11 @@ 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.utils.platform_utils import is_npu, is_rocm
+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]
@@ -17,34 +19,41 @@ from vllm_omni import Omni
 models = ["riverclouds/qwen_image_random"]


-@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):
-    def inference(offload: bool = True):
-        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_cpu_offload=offload)
-        torch.cuda.reset_peak_memory_stats(device=device_index)
-        height = 256
-        width = 256
+def inference(model_name: str, offload: bool = True):
+    current_omni_platform.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_cpu_offload=offload)
+    torch.cuda.reset_peak_memory_stats(device=device_index)
+    height = 256
+    width = 256

-        m.generate(
-            "a photo of a cat sitting on a laptop keyboard",
+    m.generate(
+        "a photo of a cat sitting on a laptop keyboard",
+        OmniDiffusionSamplingParams(
            height=height,
            width=width,
            num_inference_steps=9,
            guidance_scale=0.0,
            generator=torch.Generator("cuda").manual_seed(42),
-        )
-        peak = monitor.peak_used_mb
-        monitor.stop()
+        ),
+    )
+    peak = monitor.peak_used_mb
+    monitor.stop()

-        return peak
+    return peak

-    offload_peak_memory = inference(offload=True)
-    no_offload_peak_memory = inference(offload=False)
+
+@pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported")
+@pytest.mark.parametrize("model_name", models)
+def test_cpu_offload_diffusion_model(model_name: str):
+    try:
+        no_offload_peak_memory = inference(model_name, offload=False)
+        cleanup_dist_env_and_memory()
+        offload_peak_memory = inference(model_name, offload=True)
+    except Exception:
+        pytest.fail("Inference failed")
    print(f"Offload peak memory: {offload_peak_memory} MB")
    print(f"No offload peak memory: {no_offload_peak_memory} MB")
    assert offload_peak_memory + 2500 < no_offload_peak_memory, (
@@ -0,0 +1,110 @@
+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"
+    )
@@ -0,0 +1,138 @@
+import json
+import os
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
+# 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
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+
+# This test is specific to Z-Image LoRA behavior. Keep it focused on a single
+# model to reduce runtime and avoid extra downloads.
+models = ["Tongyi-MAI/Z-Image-Turbo"]
+
+
+@pytest.mark.parametrize("model_name", models)
+def test_diffusion_model(model_name: str, tmp_path: Path):
+    def _extract_images(outputs: list[OmniRequestOutput]):
+        if not outputs:
+            raise ValueError("Empty outputs from Omni.generate()")
+        first_output = outputs[0]
+        assert first_output.final_output_type == "image"
+        if not hasattr(first_output, "request_output") or not first_output.request_output:
+            raise ValueError("No request_output found in OmniRequestOutput")
+
+        req_out = first_output.request_output[0]
+        if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+            raise ValueError("Invalid request_output structure or missing 'images' key")
+        return req_out.images
+
+    def _write_zimage_lora(adapter_dir: Path) -> str:
+        adapter_dir.mkdir(parents=True, exist_ok=True)
+
+        # Z-Image transformer uses dim=3840 by default (see ZImageTransformer2DModel).
+        dim = 3840
+        module_name = "transformer.layers.0.attention.to_qkv"
+        rank = 1
+        lora_a = torch.zeros((rank, dim), dtype=torch.float32)
+        lora_a[0, 0] = 1.0
+
+        # QKVParallelLinear packs (Q, K, V). With tp=1 and n_kv_heads==n_heads in Z-Image,
+        # each slice is `dim`, so total out dim is `3 * dim`.
+        lora_b = torch.zeros((3 * dim, rank), dtype=torch.float32)
+        # Apply a visible delta to the Q slice only to keep the perturbation bounded.
+        lora_b[:dim, 0] = 0.1
+
+        save_file(
+            {
+                f"base_model.model.{module_name}.lora_A.weight": lora_a,
+                f"base_model.model.{module_name}.lora_B.weight": lora_b,
+            },
+            str(adapter_dir / "adapter_model.safetensors"),
+        )
+        (adapter_dir / "adapter_config.json").write_text(
+            json.dumps(
+                {
+                    "r": rank,
+                    "lora_alpha": rank,
+                    "target_modules": [module_name],
+                }
+            ),
+            encoding="utf-8",
+        )
+        return str(adapter_dir)
+
+    m = Omni(model=model_name)
+    try:
+        # high resolution may cause OOM on L4
+        height = 256
+        width = 256
+        prompt = "a photo of a cat sitting on a laptop keyboard"
+
+        outputs = m.generate(
+            prompt,
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                num_inference_steps=2,
+                guidance_scale=0.0,
+                generator=torch.Generator("cuda").manual_seed(42),
+                num_outputs_per_prompt=1,
+            ),
+        )
+        images = _extract_images(outputs)
+
+        assert len(images) == 1
+        # check image size
+        assert images[0].width == width
+        assert images[0].height == height
+
+        # Real LoRA E2E: generate again with a real on-disk PEFT adapter and
+        # verify that output changes.
+        if model_name == "Tongyi-MAI/Z-Image-Turbo":
+            from vllm_omni.lora.request import LoRARequest
+            from vllm_omni.lora.utils import stable_lora_int_id
+
+            lora_dir = _write_zimage_lora(tmp_path / "zimage_lora")
+            lora_request = LoRARequest(
+                lora_name="test",
+                lora_int_id=stable_lora_int_id(lora_dir),
+                lora_path=lora_dir,
+            )
+            outputs_lora = m.generate(
+                prompt,
+                OmniDiffusionSamplingParams(
+                    height=height,
+                    width=width,
+                    num_inference_steps=2,
+                    guidance_scale=0.0,
+                    generator=torch.Generator("cuda").manual_seed(42),
+                    num_outputs_per_prompt=1,
+                    lora_request=lora_request,
+                    lora_scale=2.0,
+                ),
+            )
+            images_lora = _extract_images(outputs_lora)
+            assert len(images_lora) == 1
+            assert images_lora[0].width == width
+            assert images_lora[0].height == height
+
+            import numpy as np
+
+            diff = np.abs(np.array(images[0], dtype=np.int16) - np.array(images_lora[0], dtype=np.int16)).mean()
+            assert diff > 0.0
+    finally:
+        m.close()
@@ -25,6 +25,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.models.ovis_image.pipeline_ovis_image import OvisImagePipeline
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams


@pytest.fixture
@@ -162,11 +163,13 @@ def test_basic_generation(ovis_pipeline):
    """Test the forward pass logic."""
    # Setup request
    req = OmniDiffusionRequest(
-        prompt="A photo of a cat",
-        height=256,
-        width=256,
-        num_inference_steps=2,
-        guidance_scale=1.0,
+        prompts=["A photo of a cat"],
+        sampling_params=OmniDiffusionSamplingParams(
+            height=256,
+            width=256,
+            num_inference_steps=2,
+            guidance_scale=1.0,
+        ),
    )

    output = ovis_pipeline(req)
@@ -184,12 +187,18 @@ def test_basic_generation(ovis_pipeline):
 def test_guidance_scale(ovis_pipeline):
    """Test that classifier-free guidance path is taken when scale > 1.0."""
    req = OmniDiffusionRequest(
-        prompt="A photo of a cat",
-        negative_prompt="bad quality",
-        height=256,
-        width=256,
-        num_inference_steps=1,
-        guidance_scale=2.0,  # Trigger CFG
+        prompts=[
+            {
+                "prompt": "A photo of a cat",
+                "negative_prompt": "bad quality",
+            }
+        ],
+        sampling_params=OmniDiffusionSamplingParams(
+            height=256,
+            width=256,
+            num_inference_steps=1,
+            guidance_scale=2.0,  # Trigger CFG
+        ),
    )

    ovis_pipeline(req)
@@ -200,9 +209,11 @@ def test_resolution_check(ovis_pipeline):
    """Test resolution divisible validation logic if present."""
    # Pass odd resolution
    req = OmniDiffusionRequest(
-        prompt="test",
-        height=250,  # Not divisible by 16 (8*2)
-        width=250,
+        prompts=["test"],
+        sampling_params=OmniDiffusionSamplingParams(
+            height=250,  # Not divisible by 16 (8*2)
+            width=250,
+        ),
    )

    # Should warn but proceed (as per code I read earlier) or resize?
@@ -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.utils import is_npu, is_rocm
+from vllm_omni.platforms import current_omni_platform

 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 is_npu():
+if current_omni_platform.is_npu():
    stage_config = str(Path(__file__).parent / "stage_configs" / "npu" / "qwen2_5_omni_ci.yaml")
-elif is_rocm():
+elif current_omni_platform.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].multimodal_output["audio"]
+        audio_tensor = audio_output.request_output[0].outputs[0].multimodal_output["audio"]
        assert audio_tensor is not None
        assert audio_tensor.numel() > 0

@@ -7,19 +7,21 @@ E2E offline tests for Omni model with video input and audio output.
 import os

 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
 from pathlib import Path

 import pytest
 from vllm.assets.video import VideoAsset

-from vllm_omni.utils import is_rocm
+from vllm_omni.platforms import current_omni_platform

 from .conftest import OmniRunner

 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]

 # CI stage config for 2xH100-80G GPUs or AMD GPU MI325
-if is_rocm():
+if current_omni_platform.is_rocm():
    # ROCm stage config optimized for MI325 GPU
    stage_configs = [str(Path(__file__).parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
 else:
@@ -73,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].multimodal_output["audio"]
+        audio_tensor = audio_output.request_output[0].outputs[0].multimodal_output["audio"]
        assert audio_tensor is not None
        assert audio_tensor.numel() > 0
@@ -2,11 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 """
-System test for Ulysses sequence parallel backend.
+System test for Sequence Parallel (SP) backends: Ulysses and Ring attention.

-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.
+Tests verify that SP inference produces correct outputs compared to baseline.
 """

 import gc
@@ -14,6 +12,7 @@ import os
 import sys
 import time
 from pathlib import Path
+from typing import NamedTuple

 import numpy as np
 import pytest
@@ -21,6 +20,8 @@ import torch
 import torch.distributed as dist
 from PIL import Image

+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
 # ruff: noqa: E402
 REPO_ROOT = Path(__file__).resolve().parents[3]
 if str(REPO_ROOT) not in sys.path:
@@ -28,27 +29,24 @@ if str(REPO_ROOT) not in sys.path:

 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.diffusion.distributed.parallel_state import device_count
-from vllm_omni.diffusion.envs import get_device_name
-
-models = ["riverclouds/qwen_image_random"]
+from vllm_omni.platforms import current_omni_platform

+# 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


-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)
+class InferenceResult(NamedTuple):
+    """Result of an inference run."""

-
-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()
+    images: list[Image.Image]
+    elapsed_ms: float


 def _cleanup_distributed():
@@ -60,201 +58,223 @@ def _cleanup_distributed():
        os.environ.pop(key, None)

    gc.collect()
-
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+    if current_omni_platform.is_available():
+        current_omni_platform.empty_cache()
+        current_omni_platform.synchronize()

    time.sleep(5)


-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(
+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(
        model=model_name,
-        parallel_config=baseline_parallel_config,
+        parallel_config=parallel_config,
        dtype=dtype,
        attention_backend=attn_backend,
    )

    try:
-        outputs = baseline.generate(
+        # 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(
            PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=4,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
+            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),
+                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:
-        baseline.close()
+        omni.close()
        _cleanup_distributed()


-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,
-    )
+# =============================================================================
+# Correctness & Performance Tests
+# =============================================================================

-    try:
-        outputs = sp.generate(
-            PROMPT,
+# 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",
+                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,
-            num_inference_steps=4,
-            guidance_scale=0.0,
-            generator=torch.Generator(get_device_name()).manual_seed(seed),
-            num_outputs_per_prompt=1,
+            warmup=sp_warmup,
        )
-        return outputs[0].request_output[0].images
-    finally:
-        sp.close()
-        _cleanup_distributed()
+        assert len(sp_result.images) == 1

+        # Compare outputs (correctness)
+        mean_diff, max_diff = _diff_metrics(baseline.images[0], sp_result.images[0])

-@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
+        # 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)

-    images = _run_baseline(model_name, dtype, attn_backend, height, width, seed)
+        # 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)")

-    assert images is not None
-    assert len(images) == 1
-    assert images[0].width == width
-    assert images[0].height == height
+        print(f"[{sp_mode}] diff: 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."
+        # 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}"
        )

-    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}"
-    )
+    # 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)