Update meetup reference in README

Signed-off-by: Roger Wang <hey@rogerw.io>
Update latest news section in README.md
2026-03-15 22:49:59 -07:00 · 2026-03-15 22:49:03 -07:00 · 2026-03-15 22:47:42 -07:00 · 2026-03-15 22:48:12 +08:00 · 2026-03-15 18:27:07 +08:00 · 2026-03-15 12:18:12 +08:00
696 changed files with 120244 additions and 14690 deletions
@@ -3,347 +3,38 @@ steps:
    key: image-build
    commands:
      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
    agents:
      queue: "cpu_queue_premerge"

-  # - label: "Test on NPU"
-  #   depends_on: ~
-  #   key: npu-test
-  #   commands:
-  #     - ".buildkite/scripts/hardware_ci/run_npu_test.sh"
-  #   agents:
-  #     queue: "ascend"
-
-  - label: "Simple Unit Test"
+  # L2 Test
+  - label: "Upload Ready Pipeline"
    depends_on: image-build
+    key: upload-ready-pipeline
+    if: build.branch != "main"
    commands:
-    - |
-      pytest -v -s \
-        tests/entrypoints/ \
-        tests/diffusion/cache/ \
-        tests/diffusion/lora/ \
-        tests/model_executor/models/qwen2_5_omni/test_audio_length.py \
-        tests/worker/ \
-        tests/distributed/omni_connectors/test_kv_flow.py \
-        --cov=vllm_omni \
-        --cov-branch \
-        --cov-report=term-missing \
-        --cov-report=html \
-        --cov-report=xml
+      - buildkite-agent pipeline upload .buildkite/test-ready.yml
    agents:
-      queue: "gpu_1_queue"
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
+      queue: "cpu_queue_premerge"

-  - label: "Diffusion Model Test"
-    timeout_in_minutes: 20
+  # L3 Test
+  - label: "Upload Merge Pipeline"
    depends_on: image-build
+    key: upload-merge-pipeline
+    if: build.branch == "main"
    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
+      - buildkite-agent pipeline upload .buildkite/test-merge.yml
    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"
+      queue: "cpu_queue_premerge"

-  - label: "Diffusion Images API LoRA E2E"
-    timeout_in_minutes: 20
+  # L4 Test
+  - label: "Upload Nightly Pipeline"
    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
    commands:
-      - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
    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:
-      - 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: "Audio Generation Model Test"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.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 Cache Backend Test"
-    timeout_in_minutes: 15
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.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 Sequence Parallelism Test"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.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
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Diffusion Tensor Parallelism Test"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.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
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Diffusion GPU Worker Test"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/diffusion/test_diffusion_worker.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
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-
-  - label: "Benchmark Test"
-    timeout_in_minutes: 15
-    depends_on: image-build
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/benchmarks/test_serve_cli.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: "Omni Model Test"
-    timeout_in_minutes: 15
-    depends_on: image-build
-    commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-    - docker#v5.2.0:
-        image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-        always-pull: true
-        propagate-environment: true
-        environment:
-        - "HF_HOME=/fsx/hf_cache"
-        volumes:
-        - "/fsx/hf_cache:/fsx/hf_cache"
-
-  # - label: "Omni Model Test with H100"
-  #   timeout_in_minutes: 30
-  #   depends_on: image-build
-  #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-  #     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-  #     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
-  #     - pytest -s -v tests/e2e/online_serving/test_async_omni.py
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 2
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
-
-  - label: "Diffusion Image Edit Test with H100 (1 GPU)"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
-    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
-
-  # - 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
+      queue: "cpu_queue_premerge"
@@ -0,0 +1,38 @@
+steps:
+  - label: ":docker: Build image"
+    key: image-build
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L2 Test
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-ready.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L3 Test
+  - label: "Upload Merge Pipeline"
+    depends_on: image-build
+    key: upload-merge-pipeline
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-merge.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L4 Test
+  - label: "Upload Nightly Pipeline"
+    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
+    agents:
+      queue: "cpu_queue_premerge"
@@ -0,0 +1,56 @@
+#!/bin/bash
+
+# This script build the XPU docker image and run the offline inference inside the container.
+set -ex
+
+omni_source_dir=$(git rev-parse --show-toplevel)
+
+base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}"
+image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}"
+container_name="xpu_${BUILDKITE_COMMIT}_$(
+    tr -dc A-Za-z0-9 </dev/urandom | head -c 10
+    echo
+)"
+
+cd "${omni_source_dir}"
+if [ -z "$(docker images -q "${base_image_name}")" ]; then
+    docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
+fi
+
+# Try building the docker image
+docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
+
+# Setup cleanup
+remove_docker_container() {
+    docker rm -f "${container_name}" || true
+    docker image rm -f "${image_name}" || true
+    docker system prune -f || true
+}
+trap remove_docker_container EXIT
+
+HF_CACHE="${HF_CACHE:-$(realpath ~)/.cache/huggingface}"
+mkdir -p "${HF_CACHE}"
+HF_MOUNT="/root/.cache/huggingface"
+
+time timeout -k 30 30m docker run \
+    --device /dev/dri:/dev/dri \
+    --net=host \
+    --ipc=host \
+    -v /dev/dri/by-path:/dev/dri/by-path \
+    -v "${HF_CACHE}:${HF_MOUNT}" \
+    --security-opt seccomp=unconfined \
+    --entrypoint="" \
+    -e VLLM_LOGGING_LEVEL \
+    -e VLLM_OMNI_LOGGING_LEVEL \
+    -e HF_TOKEN \
+    -e ZE_AFFINITY_MASK \
+    --name "${container_name}" \
+    "${image_name}" \
+    bash -c '
+    set -e
+    echo $ZE_AFFINITY_MASK
+    pip install tblib==3.1.0
+    cd /workspace/vllm-omni
+    pytest -v -s -m "core_model and xpu and B60"
+    pytest -v -s -m "advanced_model and xpu and B60"
+'
@@ -33,6 +33,7 @@ steps:
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+    - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py

 - label: "Diffusion Cache Backend Test"
  timeout_in_minutes: 15
@@ -44,7 +45,7 @@ steps:
    - export GPU_ARCHS=gfx942
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
+    - pytest -s -v -m 'core_model and cache and diffusion and not distributed_rocm and MI325'

 - label: "Diffusion Sequence Parallelism Test"
  timeout_in_minutes: 20
@@ -56,7 +57,7 @@ steps:
    - 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
+    - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model

 - label: "Diffusion Tensor Parallelism Test"
  timeout_in_minutes: 20
@@ -66,7 +67,7 @@ steps:
    - 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
+    - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py

 - label: "Diffusion GPU Worker Test"
  timeout_in_minutes: 20
@@ -88,6 +89,7 @@ steps:
    - export VLLM_LOGGING_LEVEL=DEBUG
    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    - pytest -s -v tests/engine/test_async_omni_engine_abort.py

 - label: "Omni Model Test Qwen3-Omni"
  timeout_in_minutes: 15
@@ -101,7 +103,7 @@ steps:
    - 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
@@ -0,0 +1,308 @@
+steps:
+  - label: "Simple Unit Test"
+    depends_on: upload-merge-pipeline
+    commands:
+      - "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
+    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 Images API LoRA E2E"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    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: upload-merge-pipeline
+    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:
+      - 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: "Audio Generation Model Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.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 Cache Backend Test"
+    timeout_in_minutes: 15
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
+    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 Sequence Parallelism Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.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
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Tensor Parallelism Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.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
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion GPU Worker Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v tests/diffusion/test_diffusion_worker.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
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Benchmark & Engine Test"
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 15m bash -c '
+                export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                set +e
+                pytest -s -v tests/benchmarks/test_serve_cli.py
+                EXIT1=$$?
+                pytest -s -v tests/engine/test_async_omni_engine_abort.py
+                EXIT2=$$?
+                exit $$((EXIT1 | EXIT2))
+        '
+    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
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Omni Model Test"
+    timeout_in_minutes: 15
+    depends_on: upload-merge-pipeline
+    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: upload-merge-pipeline
+    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 -m "advanced_model" --run-level "advanced_model"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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: upload-merge-pipeline
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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
+
+  # - label: "Bagel Text2Img Model Test with H100"
+  #   timeout_in_minutes: 30
+  #   depends_on: upload-merge-pipeline
+  #   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: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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
@@ -1,45 +0,0 @@
-steps:
-  - label: ":docker: Build image"
-    key: image-build
-    commands:
-      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
-      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
-      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
-      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
-    agents:
-      queue: "cpu_queue_premerge"
-
-  - label: "Omni Model Test with H100"
-    timeout_in_minutes: 180
-    depends_on: image-build
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.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
@@ -0,0 +1,178 @@
+steps:
+  - label: ":full_moon: Omni Model Test with H100"
+    timeout_in_minutes: 90
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - |
+        set +e
+        pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
+        EXIT1=$$?
+        pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+        EXIT2=$$?
+        exit $$((EXIT1 | EXIT2))
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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: ":full_moon: Omni Model Test"
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
+    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
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: ":full_moon: Diffusion Model Test with H100"
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    # if: build.env("NIGHTLY") == "1"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
+    timeout_in_minutes: 30
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - |
+        huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
+    agents:
+      queue: "gpu_4_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
+    key: nightly-performance
+    timeout_in_minutes: 180
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export BENCHMARK_DIR=tests
+      - pytest -s -v tests/perf/scripts/run_benchmark.py
+      - buildkite-agent artifact upload "tests/*.json"
+      - python tools/nightly/buildkite_testcase_statistics.py -o tests/buildkite_testcase_statistics.html
+      - buildkite-agent artifact upload "tests/*.html"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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: ":email: Nightly Collection & Email"
+    key: nightly-perf-distribution
+    depends_on:
+    - nightly-performance
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - pip install openpyxl
+      - export DEFAULT_INPUT_DIR=tests
+      - export DEFAULT_OUTPUT_DIR=tests
+      - buildkite-agent artifact download "tests/*.json" . --step nightly-performance
+      - buildkite-agent artifact download "tests/*.html" . --step nightly-performance
+      - python tools/nightly/generate_nightly_perf_excel.py
+      - python tools/nightly/send_nightly_email.py --report-file "tests/*.xlsx, tests/*.html"
+      - buildkite-agent artifact upload "tests/*.xlsx"
+    agents:
+      queue: "cpu_queue_premerge"
@@ -0,0 +1,377 @@
+steps:
+  - label: "Simple Unit Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model CPU offloading Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.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: "Audio Generation Model Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.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 Cache Backend Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 15m pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
+    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 Sequence Parallelism Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
+    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
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion GPU Worker Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.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
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+
+  - label: "Benchmark & Engine Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 15m bash -c '
+                export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                set +e
+                pytest -s -v tests/benchmarks/test_serve_cli.py
+                EXIT1=$$?
+                pytest -s -v tests/engine/test_async_omni_engine_abort.py
+                EXIT2=$$?
+                exit $$((EXIT1 | EXIT2))
+        '
+    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
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+
+  - label: "Omni Model Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 17m bash -c '
+          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"
+  #   depends_on: upload-ready-pipeline
+  #   commands:
+  #     - |
+  #       timeout 20m bash -c '
+  #         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 -m "core_model" --run-level "core_model"
+  #       '
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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: "Qwen3-TTS E2E Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
+        '
+    agents:
+      queue: "gpu_4_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  # - label: "Diffusion Image Edit Test with H100 (1 GPU)"
+  #   depends_on: upload-ready-pipeline
+  #   commands:
+  #     - |
+  #       timeout 20m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+  #       '
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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
+
+  - label: "Bagel Text2Img Model Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 30m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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
+
+  - label: "Bagel Img2Img Model Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 30m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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
+
+  - label: "Bagel Online Serving Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 40m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          export VLLM_IMAGE_FETCH_TIMEOUT=60
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/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
@@ -14,7 +14,7 @@
        soft_fail: false
        commands:
          - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
-          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target final --progress plain ."
+          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
          - "docker push {{ docker_image_amd }}"
        key: "amd-build"
        env:
@@ -12,10 +12,10 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 <summary> Essential Elements of an Effective PR Description Checklist </summary>

 - [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
- [ ] The test plan, such as providing test command.
- [ ] The test results, such as pasting the results comparison before and after, or e2e results
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft.
+- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
+- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
+- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
+- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
 </details>

 **BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
@@ -0,0 +1,629 @@
+#!/usr/bin/env python3
+"""
+PR Reviewer using GLM API for vllm-omni project.
+"""
+
+import json
+import logging
+import os
+import sys
+import time
+from dataclasses import dataclass
+from typing import Any, TypedDict
+
+import requests
+
+
+# Type definitions for API responses
+class PRDetails(TypedDict):
+    """Type definition for GitHub PR details response."""
+
+    title: str
+    body: str
+    number: int
+    state: str
+    user: dict[str, Any]
+
+
+class GLMMessage(TypedDict):
+    """Type definition for GLM API message."""
+
+    role: str
+    content: str
+
+
+class GLMChoice(TypedDict):
+    """Type definition for GLM API choice."""
+
+    message: GLMMessage
+    finish_reason: str
+
+
+class GLMResponse(TypedDict):
+    """Type definition for GLM API response."""
+
+    choices: list[GLMChoice]
+    usage: dict[str, int] | None
+
+
+class GitHubComment(TypedDict):
+    """Type definition for GitHub comment."""
+
+    id: int
+    body: str
+    created_at: str
+    user: dict[str, Any]
+
+
+# Configuration
+TRIGGER_PHRASE: str = "@vllm-omni-reviewer"
+DEFAULT_GLM_API_URL: str = "https://open.bigmodel.cn/api/paas/v4/chat/completions"  # noqa: E501
+DEFAULT_GLM_MODEL: str = "glm-5"
+DEFAULT_COOLDOWN_MINUTES: int = 5
+DEFAULT_MAX_RETRIES: int = 3
+DEFAULT_RETRY_DELAY: float = 1.0
+MAX_DIFF_SIZE: int = 100_000  # Maximum diff size in characters
+
+
+@dataclass
+class Config:
+    """Configuration for the PR reviewer."""
+
+    glm_api_url: str
+    glm_model: str
+    cooldown_minutes: int
+    max_retries: int
+    retry_delay: float
+    max_diff_size: int
+
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="[PR Reviewer] %(asctime)s - %(levelname)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def get_config() -> Config:
+    """Load configuration from environment variables with defaults."""
+    return Config(
+        glm_api_url=os.getenv("GLM_API_URL", DEFAULT_GLM_API_URL),
+        glm_model=os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL),
+        cooldown_minutes=int(
+            os.getenv(
+                "PR_REVIEWER_COOLDOWN_MINUTES",
+                str(DEFAULT_COOLDOWN_MINUTES),
+            )
+        ),
+        max_retries=int(
+            os.getenv(
+                "PR_REVIEWER_MAX_RETRIES",
+                str(DEFAULT_MAX_RETRIES),
+            )
+        ),
+        retry_delay=float(os.getenv("PR_REVIEWER_RETRY_DELAY", str(DEFAULT_RETRY_DELAY))),
+        max_diff_size=int(os.getenv("PR_REVIEWER_MAX_DIFF_SIZE", str(MAX_DIFF_SIZE))),  # noqa: E501
+    )
+
+
+def get_env_var(name: str) -> str:
+    """
+    Get an environment variable or raise an error.
+
+    Args:
+        name: Name of the environment variable.
+
+    Returns:
+        The value of the environment variable.
+
+    Raises:
+        SystemExit: If the environment variable is not set.
+    """
+    value = os.environ.get(name)
+    if not value:
+        logger.error(f"Environment variable {name} is not set")
+        sys.exit(1)
+    return value
+
+
+def check_trigger(comment_body: str) -> bool:
+    """
+    Check if the comment contains the trigger phrase.
+
+    Args:
+        comment_body: The body of the comment to check.
+
+    Returns:
+        True if the trigger phrase is found, False otherwise.
+    """
+    return TRIGGER_PHRASE in comment_body
+
+
+def fetch_pr_diff(
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    max_size: int = MAX_DIFF_SIZE,
+) -> str | None:
+    """
+    Fetch the diff for a pull request.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        max_size: Maximum diff size in characters.
+
+    Returns:
+        The diff content as a string, or None if fetching failed.
+        Returns empty string if diff is larger than max_size.
+    """
+    url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3.diff",
+    }
+
+    logger.info(f"Fetching PR diff from {url}")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code == 200:
+        diff: str = response.text
+        if len(diff) > max_size:
+            logger.warning(
+                f"Diff size ({len(diff)} bytes) exceeds maximum "
+                f"({max_size} bytes), truncating to first "
+                f"{max_size} bytes"
+            )
+            return diff[:max_size] + "\n\n... [Diff truncated due to size] ..."
+        logger.info(f"Successfully fetched diff ({len(diff)} bytes)")
+        return diff
+    else:
+        logger.error(f"Failed to fetch PR diff: {response.status_code}")
+        logger.error(f"Response: {response.text}")
+        return None
+
+
+def fetch_pr_details(
+    repo_name: str,
+    pr_number: int,
+    token: str,
+) -> PRDetails | None:
+    """
+    Fetch PR details including title and description.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+
+    Returns:
+        A dictionary containing PR details, or None if fetching failed.
+    """
+    url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    logger.info(f"Fetching PR details from {url}")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code == 200:
+        return response.json()
+    else:
+        logger.error(f"Failed to fetch PR details: {response.status_code}")
+        return None
+
+
+def build_review_prompt(pr_title: str, pr_description: str, diff: str) -> str:
+    """
+    Build the prompt for the GLM-4.7 API.
+
+    Args:
+        pr_title: The title of the pull request.
+        pr_description: The description/body of the pull request.
+        diff: The diff content of the pull request.
+
+    Returns:
+        The formatted prompt string for the API.
+    """
+    return f"""You are an expert code reviewer for the VLLM-Omni project. \
+Please review the following pull request:
+
+## Pull Request Details
+**Title:** {pr_title}
+
+**Description:**
+{pr_description if pr_description else "No description provided."}
+
+## Code Changes (Diff)
+{diff}
+
+## Review Guidelines
+
+Please provide a comprehensive code review with the following sections:
+
+### 1. Overview
+- Brief summary of the changes
+- Overall assessment (positive, neutral, or concerns)
+
+### 2. Code Quality
+- Code style and consistency
+- Potential bugs or edge cases
+- Performance considerations
+- Error handling
+
+### 3. Architecture & Design
+- Integration with existing codebase
+- Design patterns and best practices
+- Potential improvements
+
+### 4. Security & Safety
+- Security concerns (if any)
+- Resource management
+- Input validation
+
+### 5. Testing & Documentation
+- Test coverage considerations
+- Documentation completeness
+- Examples and usage clarity
+
+### 6. Specific Suggestions
+- Line-by-line specific feedback (use `file:line` format)
+- Concrete actionable suggestions
+- Code examples for improvements (if applicable)
+
+### 7. Approval Status
+- **LGTM** (Looks Good To Me) if the PR is ready to merge
+- **LGTM with suggestions** if the PR is good but has minor suggestions
+- **Changes requested** if significant changes are needed
+
+## Important Notes
+- Be constructive and helpful
+- Focus on objective technical feedback
+- Acknowledge good practices when you see them
+- Prioritize critical issues over nitpicks
+- If the diff is empty or minimal, acknowledge this and provide
+  any relevant context-specific guidance
+
+Please format your response in Markdown with clear section headers.
+"""
+
+
+def validate_glm_response(data: dict[str, Any]) -> str | None:
+    """
+    Validate and extract content from GLM API response.
+
+    Args:
+        data: The response data from GLM API.
+
+    Returns:
+        The review content string if valid, None otherwise.
+    """
+    # Check if choices exists and is a non-empty list
+    if "choices" not in data:
+        logger.error("GLM API response missing 'choices' field")
+        logger.error(f"Response structure: {json.dumps(data, indent=2)}")
+        return None
+
+    choices = data["choices"]
+    if not isinstance(choices, list):
+        logger.error(f"GLM API 'choices' is not a list: {type(choices)}")
+        return None
+
+    if len(choices) == 0:
+        logger.error("GLM API 'choices' is an empty list")
+        return None
+
+    # Check if first choice has message
+    try:
+        first_choice = choices[0]
+        if not isinstance(first_choice, dict):
+            logger.error(f"GLM API choice is not a dict: {type(first_choice)}")
+            return None
+
+        if "message" not in first_choice:
+            logger.error("GLM API choice missing 'message' field")
+            logger.error(f"Choice structure: {json.dumps(first_choice, indent=2)}")  # noqa: E501
+            return None
+
+        message = first_choice["message"]
+        if not isinstance(message, dict):
+            logger.error(f"GLM API message is not a dict: {type(message)}")
+            return None
+
+        if "content" not in message:
+            logger.error("GLM API message missing 'content' field")
+            logger.error(f"Message structure: {json.dumps(message, indent=2)}")
+            return None
+
+        content = message["content"]
+        if not isinstance(content, str):
+            logger.error(f"GLM API content is not a string: {type(content)}")
+            return None
+
+        return content
+
+    except (KeyError, IndexError, TypeError) as e:
+        logger.error(f"Failed to parse GLM API response: {e}")
+        logger.error(f"Response: {json.dumps(data, indent=2)}")
+        return None
+
+
+def call_glm_api(prompt: str, api_key: str, config: Config) -> str | None:
+    """
+    Call the GLM-4.7 API to get code review with retry logic.
+
+    Args:
+        prompt: The prompt to send to the API.
+        api_key: The GLM API key.
+        config: Configuration object.
+
+    Returns:
+        The review content as a string, or None if all retries failed.
+    """
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+    }
+
+    payload: dict[str, Any] = {
+        "model": config.glm_model,
+        "messages": [{"role": "user", "content": prompt}],
+        "temperature": 0.3,
+        "max_tokens": 32000,
+        "top_p": 0.9,
+    }
+
+    last_error: str | None = None
+
+    for attempt in range(config.max_retries):
+        try:
+            logger.info(f"Calling GLM API ({config.glm_model}) - Attempt {attempt + 1}/{config.max_retries}")
+            response = requests.post(
+                config.glm_api_url,
+                headers=headers,
+                json=payload,
+                timeout=120,
+            )
+
+            if response.status_code == 200:
+                data = response.json()
+                review = validate_glm_response(data)
+                if review:
+                    logger.info(f"Successfully received review ({len(review)} chars)")  # noqa: E501
+                    return review
+                else:
+                    last_error = "Failed to validate API response structure"
+                    logger.error(last_error)
+            else:
+                last_error = f"GLM API request failed: {response.status_code} - {response.text}"
+                logger.error(last_error)
+
+        except requests.exceptions.Timeout:
+            last_error = f"GLM API request timed out (attempt {attempt + 1})"
+            logger.error(last_error)
+        except requests.exceptions.RequestException as e:
+            last_error = f"GLM API request exception: {e}"
+            logger.error(last_error)
+        except json.JSONDecodeError as e:
+            last_error = f"Failed to decode GLM API response as JSON: {e}"
+            logger.error(last_error)
+
+        # Exponential backoff before retry
+        if attempt < config.max_retries - 1:
+            wait_time: float = config.retry_delay * (2**attempt)
+            logger.info(f"Waiting {wait_time}s before retry...")  # noqa: E501
+            time.sleep(wait_time)
+
+    logger.error(
+        f"All {config.max_retries} attempts failed. Last error: {last_error}"  # noqa: E501
+    )
+    return None
+
+
+def check_cooldown(  # noqa: E501
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    cooldown_minutes: int,
+) -> bool:
+    """
+    Check if the PR is within the cooldown period.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        cooldown_minutes: Cooldown period in minutes.
+
+    Returns:
+        True if within cooldown period (should skip), False otherwise.
+    """
+    from datetime import datetime, timedelta
+
+    url: str = (
+        f"https://api.github.com/repos/{repo_name}/issues/"
+        f"{pr_number}/comments"  # noqa: E501
+    )
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    logger.info(f"Checking cooldown period ({cooldown_minutes} minutes)")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code != 200:
+        logger.warning(f"Failed to check cooldown: {response.status_code}, proceeding with review")
+        return False
+
+    comments: list[dict[str, Any]] = response.json()
+    cutoff_time: datetime = datetime.utcnow() - timedelta(minutes=cooldown_minutes)  # noqa: E501
+
+    for comment in reversed(comments):
+        # Check if this is a bot comment
+        body: str = comment.get("body", "")
+        if "VLLM-Omni PR Review" in body or "PR Reviewer Bot" in body:
+            created_at_str: str = comment.get("created_at", "")
+            try:
+                # Parse GitHub timestamp format
+                created_at = datetime.fromisoformat(created_at_str.replace("Z", "+00:00"))
+                created_at = created_at.replace(tzinfo=None)
+                if created_at > cutoff_time:
+                    logger.info(f"PR is within cooldown period (last review: {created_at_str})")
+                    return True
+            except ValueError:
+                logger.warning(f"Failed to parse comment timestamp: {created_at_str}")  # noqa: E501
+                continue
+
+    logger.info("PR is outside cooldown period, proceeding with review")
+    return False
+
+
+def post_review_comment(  # noqa: E501
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    review: str,
+) -> bool:
+    """
+    Post the review as a comment on the PR.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        review: The review content to post.
+
+    Returns:
+        True if posting succeeded, False otherwise.
+    """
+    url: str = (
+        f"https://api.github.com/repos/{repo_name}/issues/"
+        f"{pr_number}/comments"  # noqa: E501
+    )
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    # Format the review comment
+    comment_body: str = f"""## 🤖 VLLM-Omni PR Review
+
+{review}
+
+---
+*This review was generated automatically by the VLLM-Omni PR Reviewer Bot
+using {os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL)}.*
+"""
+
+    payload: dict[str, str] = {"body": comment_body}
+
+    logger.info(f"Posting review comment to PR #{pr_number}")
+    response = requests.post(url, headers=headers, json=payload, timeout=30)
+
+    if response.status_code == 201:
+        logger.info("Successfully posted review comment")
+        return True
+    else:
+        logger.error(f"Failed to post comment: {response.status_code}")
+        logger.error(f"Response: {response.text}")
+        return False
+
+
+def main() -> int:
+    """
+    Main entry point for the PR reviewer bot.
+
+    Returns:
+        0 on success, 1 on error.
+    """
+    logger.info("VLLM-Omni PR Reviewer Bot starting...")
+
+    # Load configuration
+    config: Config = get_config()
+    logger.info(
+        f"Configuration: model={config.glm_model}, "
+        f"cooldown={config.cooldown_minutes}min, "
+        f"max_retries={config.max_retries}"
+    )
+
+    # Get environment variables
+    token: str = get_env_var("GITHUB_TOKEN")
+    api_key: str = get_env_var("GLM_API_KEY")
+    repo_name: str = get_env_var("REPO_NAME")
+    pr_number_str: str = get_env_var("PR_NUMBER")
+    comment_body: str = get_env_var("COMMENT_BODY")
+
+    try:
+        pr_number: int = int(pr_number_str)
+    except ValueError:
+        logger.error(f"Invalid PR number: {pr_number_str}")
+        return 1
+
+    logger.info(f"Repository: {repo_name}")
+    logger.info(f"PR Number: {pr_number}")
+
+    # Check if the comment contains the trigger phrase
+    if not check_trigger(comment_body):
+        logger.info(
+            f"Comment does not contain trigger phrase '{TRIGGER_PHRASE}', exiting"  # noqa: E501
+        )
+        return 0
+
+    logger.info("Trigger phrase detected! Starting review process...")
+
+    # Check cooldown period
+    if check_cooldown(repo_name, pr_number, token, config.cooldown_minutes):
+        logger.info("Skipping review due to cooldown period")
+        return 0
+
+    # Fetch PR details
+    logger.info("Step 1/4: Fetching PR details...")
+    pr_details: PRDetails | None = fetch_pr_details(repo_name, pr_number, token)  # noqa: E501
+    if not pr_details:
+        logger.error("Failed to fetch PR details")
+        return 1
+
+    pr_title: str = pr_details.get("title", "Unknown")
+    pr_description: str = pr_details.get("body", "")
+
+    logger.info(f"PR Title: {pr_title}")
+
+    # Fetch PR diff
+    logger.info("Step 2/4: Fetching PR diff...")
+    diff: str | None = fetch_pr_diff(repo_name, pr_number, token, config.max_diff_size)
+    if diff is None:
+        logger.error("Failed to fetch PR diff")
+        return 1
+
+    if not diff:
+        logger.warning("Warning: Empty diff - this might be a draft PR or no code changes")
+
+    # Build prompt
+    logger.info("Step 3/4: Building review prompt...")
+    prompt: str = build_review_prompt(pr_title, pr_description, diff)
+
+    # Call GLM API
+    logger.info("Step 4/4: Calling GLM API...")
+    review: str | None = call_glm_api(prompt, api_key, config)
+    if not review:
+        logger.error("Failed to get review from GLM API")
+        return 1
+
+    # Post review comment
+    logger.info("Posting review comment...")
+    if not post_review_comment(repo_name, pr_number, token, review):
+        logger.error("Failed to post review comment")
+        return 1
+
+    logger.info("PR review completed successfully!")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
@@ -1,27 +0,0 @@
-name: ComfyUI Integration -- Publish to Comfy registry
-on:
-  workflow_dispatch:
-  push:
-    tags:
-      - "*"
-    paths:
-      - "apps/ComfyUI-vLLM-Omni/**"
-
-permissions:
-  issues: write
-
-defaults:
-  run:
-    working-directory: apps/ComfyUI-vLLM-Omni
-
-jobs:
-  publish-node:
-    name: Publish Custom Node to registry
-    runs-on: ubuntu-latest
-    steps:
-      - name: ♻️ Check out code
-        uses: actions/checkout@v4
-      - name: 📦 Publish Custom Node
-        uses: Comfy-Org/publish-node-action@main
-        with:
-          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
@@ -1,40 +0,0 @@
-name: ComfyUI Integration -- Validate backwards compatibility
-
-on:
-  pull_request:
-    branches:
-      - master
-      - main
-    paths:
-      - "apps/ComfyUI-vLLM-Omni/**"
-
-defaults:
-  run:
-    working-directory: apps/ComfyUI-vLLM-Omni
-
-jobs:
-  check-base-path:
-    runs-on: ubuntu-latest
-    outputs:
-      exists: ${{ steps.check.outputs.exists }}
-    defaults:
-      run:
-        working-directory: .
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          ref: ${{ github.event.pull_request.base.sha }}
-          fetch-depth: 1
-      - id: check
-        run: |
-          if [ -d "apps/ComfyUI-vLLM-Omni" ]; then
-            echo "exists=true" >> "$GITHUB_OUTPUT"
-          else
-            echo "exists=false" >> "$GITHUB_OUTPUT"
-          fi
-  validate:
-    needs: check-base-path
-    if: needs.check-base-path.outputs.exists == 'true'
-    runs-on: ubuntu-latest
-    steps:
-      - uses: comfy-org/node-diff@main
@@ -0,0 +1,62 @@
+name: VLLM-Omni PR Reviewer
+
+on:
+  issue_comment:
+    types: [created]
+
+permissions:
+  contents: read
+  pull-requests: write
+  issues: write
+
+jobs:
+  pr-reviewer:
+    name: Review Pull Request
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    # Only run when the comment is from a collaborator/owner/member
+    if: |
+      github.event_name == 'issue_comment' &&
+      github.event.issue.pull_request != null &&
+      contains(github.event.comment.body, '@vllm-omni-reviewer') &&
+      (github.event.comment.author_association == 'MEMBER' ||
+       github.event.comment.author_association == 'COLLABORATOR' ||
+       github.event.comment.author_association == 'OWNER')
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4.2.2
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v5.3.0
+        with:
+          python-version: '3.11'
+          cache: 'pip'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip==24.0
+          pip install requests==2.31.0 pyyaml==6.0.1
+
+      - name: Run PR Reviewer
+        id: reviewer
+        env:
+          GLM_API_KEY: ${{ secrets.GLM_API_KEY }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}
+          COMMENT_BODY: ${{ github.event.comment.body }}
+          REPO_NAME: ${{ github.repository }}
+          PR_REVIEWER_COOLDOWN_MINUTES: 5
+          PR_REVIEWER_MAX_RETRIES: 3
+        run: |
+          python .github/scripts/pr_reviewer.py 2>&1 | tee "/tmp/pr_review_${PR_NUMBER}.log"
+
+      - name: Upload review logs
+        if: always()
+        uses: actions/upload-artifact@v4.4.3
+        with:
+          name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
+          path: /tmp/pr_review_*.log
+          retention-days: 7
+          if-no-files-found: ignore
@@ -242,3 +242,8 @@ configs/development.yaml
 Dockerfile.dev
 discussion
 tmp_test
+
+# Auto-generated version file (created by setuptools_scm during build)
+vllm_omni/_version.py
+# output files
+*.wav
@@ -8,6 +8,7 @@ repos:
      # list of supported hooks: https://pre-commit.com/hooks.html
      - id: check-yaml
        args: ["--unsafe"]
+        exclude: "examples/online_serving/chart-helm/templates/"
      - id: debug-statements
      - id: end-of-file-fixer
      - id: mixed-line-ending
@@ -1,6 +1,6 @@
 <p align="center">
  <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
+    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
    <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
  </picture>
 </p>
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>

 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
 </p>


@@ -17,6 +17,9 @@ Easy, fast, and cheap omni-modality model serving for everyone

 *Latest News* 🔥

+- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
+- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
+- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
 - [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
@@ -47,46 +47,35 @@ If no, check your shell running the ComfyUI process. There may be some error mes
 This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):

 - **Generate Image** for text-to-image and image-to-image tasks
- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
+- **Generate Video** for text-to-video and image-to-video tasks
+- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
 - **TTS** and **TTS Voice Clone** for TTS tasks

 This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)

-> [!INFO]
+> [!NOTE]
 > The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.

 To build a simple workflow yourself,

 - Drag a generation node onto the canvas.
 - Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
 - If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
    - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
    - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
    - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.

-**The following features are tested**:
-
- Single-node workflows for
-    - Multimodal Comprehension (e.g., Qwen Omni, BAGEL)
-    - Text-to-Image Generation (e.g., Qwen-Image)
-    - Image-to-Image Generation (e.g., Qwen-Image-Edit)
-    - TTS (e.g., Qwen TTS, including VoiceDesign, VoiceClone, CustomVoice)
-
-**The following features are not currently tested**. They may work or break. You are welcomed to test it out and offer comments.
-
- Multi-node workflow that connects multiple model services together.
-
 ## Screenshots and Examples

-### Multimodal comprehension (e.g., Qwen Omni series, BAGEL)
+### Multimodal understanding (e.g., Qwen Omni series, BAGEL)

-(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Annotated Example**)
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Multimodal Understanding**)

 <p align="center">
  <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg" width=55%>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.jpg">
+    <img alt="vLLM-Omni multimodal understanding" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.jpg" width=55%>
  </picture>
 </p>

@@ -98,7 +87,7 @@ You can configure per-stage sampling parameters for multi-stage models.
 <p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
+    <img alt="vLLM-Omni multiple stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
  </picture>
 </p>

@@ -109,13 +98,27 @@ You can configure per-stage sampling parameters for multi-stage models.
 <p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
+    <img alt="vLLM-Omni image generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
  </picture>
 </p>

 > [!TIP]
 > The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.

+### Text-to-video and image-to-video generation (e.g., Wan)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg">
+    <img alt="vLLM-Omni video generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
+
 ### TTS (e.g., Qwen TTS series)

 (Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)
@@ -123,13 +126,24 @@ You can configure per-stage sampling parameters for multi-stage models.
 <p align="center">
  <picture>
    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg">
-    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
+    <img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
  </picture>
 </p>

 > [!TIP]
 > There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node.

+### Chaining multiple model services
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Chaining Services**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg">
+    <img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg" width=55%>
+  </picture>
+</p>
+
 ## Develop

 Follow the [development convention and rules of vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/).
@@ -6,47 +6,56 @@ __all__ = [
    "WEB_DIRECTORY",
 ]

-__author__ = """Zeyu Huang"""
-__email__ = "11222265+fhfuih@users.noreply.github.com"
+__author__ = """vLLM-Omni Team"""
+__email__ = "vllm-omni@vllm.ai"
 __version__ = "0.0.1"

 from .comfyui_vllm_omni.nodes import (
    VLLMOmniARSampling,
-    VLLMOmniComprehension,
    VLLMOmniDiffusionSampling,
    VLLMOmniGenerateImage,
+    VLLMOmniGenerateVideo,
    VLLMOmniQwenTTSParams,
+    VLLMOmniRemoteLoRA,
    VLLMOmniSamplingParamsList,
    VLLMOmniTTS,
+    VLLMOmniUnderstanding,
    VLLMOmniVoiceClone,
+    VLLMOmniWanParams,
 )

 # A dictionary that contains all nodes you want to export with their names
 NODE_CLASS_MAPPINGS = {
    # === Generation ===
    "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
-    "VLLMOmniComprehension": VLLMOmniComprehension,
+    "VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
+    "VLLMOmniUnderstanding": VLLMOmniUnderstanding,
    "VLLMOmniTTS": VLLMOmniTTS,
    "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
    # === Params ===
    "VLLMOmniARSampling": VLLMOmniARSampling,
    "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
    "VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
+    "VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
    "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
+    "VLLMOmniWanParams": VLLMOmniWanParams,
 }

 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
    # === Generation ===
    "VLLMOmniGenerateImage": "Generate Image",
-    "VLLMOmniComprehension": "Multimodality Comprehension",
+    "VLLMOmniGenerateVideo": "Generate Video",
+    "VLLMOmniUnderstanding": "Multimodality Understanding",
    "VLLMOmniTTS": "TTS (Text to Speech)",
    "VLLMOmniVoiceClone": "TTS Voice Cloning",
    # === Params ===
    "VLLMOmniARSampling": "AR Sampling Params",
    "VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
    "VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
+    "VLLMOmniRemoteLoRA": "LoRA",
    "VLLMOmniQwenTTSParams": "Qwen TTS Params",
+    "VLLMOmniWanParams": "Wan Video Params",
 }

 WEB_DIRECTORY = "./web"
@@ -1,4 +1,4 @@
-from typing import Literal, cast
+from typing import Literal

 import torch
 from comfy_api.input import AudioInput, VideoInput
@@ -6,7 +6,13 @@ from comfy_api.input import AudioInput, VideoInput
 from .utils.api_client import VLLMOmniClient
 from .utils.logger import get_logger
 from .utils.models import lookup_model_spec
-from .utils.types import AudioFormat
+from .utils.types import (
+    AudioFormat,
+    AutoregressionSamplingParams,
+    DiffusionSamplingParams,
+    QwenTTSModelSpecificParams,
+    WanModelSpecificParams,
+)
 from .utils.validators import (
    add_sampling_parameters_to_stage,
    validate_model_and_sampling_params_types,
@@ -51,6 +57,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                # "video": ("VIDEO",),
                # "audio": ("AUDIO",),
                "sampling_params": ("SAMPLING_PARAMS",),
+                "lora": ("REMOTE_LORA",),
            },
        }

@@ -71,9 +78,11 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
        audio: AudioInput | None = None,  # Hidden & unused
        video: VideoInput | None = None,  # Hidden & unused
        sampling_params: dict | list[dict] | None = None,
+        lora: dict | None = None,
        **kwargs,
    ):
-        logger.info("Uncaught kwargs: %s", kwargs)
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
        logger.debug("Got sampling params: %s", sampling_params)
        validate_model_and_sampling_params_types(model, sampling_params)
        if image is None and mask is not None:
@@ -86,7 +95,10 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):

        # Prefer DALL-E compatible API for simple (one-stage) diffusion models
        if (spec is None or spec["stages"] == ["diffusion"]) and not is_bagel:
-            sampling_params = cast(dict | None, sampling_params)
+            # The number of sampling parameter groups should have been validated.
+            # Now, simply convert single-item list to dict.
+            if isinstance(sampling_params, list):
+                sampling_params = sampling_params[0]
            if audio is None and image is None and video is None:
                # No multimodal input --- use DALL-E image generation
                logger.info("Using DALL-E image generation endpoint")
@@ -97,6 +109,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                    height=height,
                    negative_prompt=negative_prompt,
                    sampling_params=sampling_params,
+                    lora=lora,
                )
                return (output,)
            elif image is not None and audio is None and video is None:
@@ -111,6 +124,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
                    negative_prompt=negative_prompt,
                    mask=mask,
                    sampling_params=sampling_params,
+                    lora=lora,
                )
                return (output,)

@@ -128,12 +142,92 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
            audio=audio,
            video=video,
            sampling_params=sampling_params,
+            lora=lora,
        )

        return (output,)


-class VLLMOmniComprehension(_VLLMOmniGenerateBase):
+class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
+                "width": ("INT", {"default": 832, "min": 1}),
+                "height": ("INT", {"default": 480, "min": 1}),
+                "fps": ("INT", {"default": 16, "min": 1}),
+                "num_frames": ("INT", {"default": 41, "min": 1}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+                "lora": ("REMOTE_LORA",),
+                "model_params": ("VIDEO_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("VIDEO",)
+    RETURN_NAMES = ("video",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        fps: int,
+        num_frames: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        model_params: dict | None = None,
+        lora: dict | None = None,
+        **kwargs,
+    ):
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
+        logger.debug("Got sampling params: %s", sampling_params)
+        logger.debug("Got model params: %s", model_params)
+        validate_model_and_sampling_params_types(model, sampling_params)
+
+        # Currently, all video generation models are single-stage diffusion models
+        if isinstance(sampling_params, list):
+            if len(sampling_params) != 1:
+                raise ValueError(
+                    "Video generation expects a single sampling params group. "
+                    "Please provide one Diffusion sampling node."
+                )
+            sampling_params = sampling_params[0]
+
+        if sampling_params is not None:
+            sampling_params.pop("type", None)  # internal fields
+        if model_params is not None:
+            model_params.pop("type", None)  # internal fields
+
+        client = VLLMOmniClient(url)
+        output = await client.generate_video(
+            model=model,
+            prompt=prompt,
+            image=image,  # image present => i2v, absent => t2v
+            width=width,
+            height=height,
+            num_frames=num_frames,
+            fps=fps,
+            negative_prompt=negative_prompt,
+            sampling_params=sampling_params,
+            lora=lora,
+            model_params=model_params,
+        )
+        return (output,)
+
+
+class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
    @classmethod
    def INPUT_TYPES(cls):
        return {
@@ -180,7 +274,8 @@ class VLLMOmniComprehension(_VLLMOmniGenerateBase):
        use_audio_in_video: bool = True,
        **kwargs,
    ) -> tuple[str, AudioInput]:
-        logger.info("Uncaught kwargs: %s", kwargs)
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
        logger.debug("Got sampling params: %s", sampling_params)
        validate_model_and_sampling_params_types(model, sampling_params)

@@ -197,7 +292,7 @@ class VLLMOmniComprehension(_VLLMOmniGenerateBase):
            (
                text_response,
                _,
-            ) = await client.generate_comprehension_chat_completion(
+            ) = await client.generate_understanding_chat_completion(
                model=model,
                prompt=prompt,
                image=image,
@@ -221,7 +316,7 @@ class VLLMOmniComprehension(_VLLMOmniGenerateBase):
            (
                text_response,
                audio,
-            ) = await client.generate_comprehension_chat_completion(
+            ) = await client.generate_understanding_chat_completion(
                model=model,
                prompt=prompt,
                image=image,
@@ -287,15 +382,13 @@ class VLLMOmniTTS(_VLLMOmniGenerateBase):
        logger.info("Got extra kwargs in TTS: %s", kwargs)

        is_qwen_tts = "qwen3-tts" in model.lower()
-        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
-        if not is_qwen_tts and extra_params_type == "qwen-tts":
+        if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
            raise ValueError(
                "You have provided Qwen-specific TTS params."
                "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
            )

        combined_params = {**kwargs, **(model_specific_params or {})}
-        combined_params.pop("type", None)  # Internal fields in model_specific_params

        client = VLLMOmniClient(url)

@@ -352,8 +445,7 @@ class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
        **kwargs,
    ):
        is_qwen_tts = "qwen3-tts" in model.lower()
-        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
-        if not is_qwen_tts and extra_params_type == "qwen-tts":
+        if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
            raise ValueError(
                "You have provided Qwen-specific TTS params."
                "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
@@ -366,7 +458,6 @@ class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
            **kwargs,
            **(model_specific_params or {}),
        }
-        combined_params.pop("type", None)  # Internal fields in model_specific_params

        client = VLLMOmniClient(url)

@@ -419,10 +510,7 @@ class VLLMOmniARSampling:
    CATEGORY = "vLLM-Omni/Sampling Params"

    def get_params(self, seed, **kwargs):
-        params = {
-            "type": "autoregression",  # for internal use, removed before sending the request
-            **kwargs,
-        }
+        params = AutoregressionSamplingParams(kwargs)
        if seed >= 0:
            params["seed"] = seed
        return (params,)
@@ -479,6 +567,13 @@ class VLLMOmniDiffusionSampling:
                        "tooltip": "Enable VAE slicing for reduced memory usage (slight quality trade-off)",
                    },
                ),
+                "vae_use_tiling": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Enable VAE tiling for reduced memory usage (slight quality trade-off)",
+                    },
+                ),
                # === Put seed at last. ===
                # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
                "seed": (
@@ -499,10 +594,7 @@ class VLLMOmniDiffusionSampling:
    CATEGORY = "vLLM-Omni/Sampling Params"

    def get_params(self, seed, **kwargs):
-        params = {
-            "type": "diffusion",  # for internal use, removed before sending the request
-            **kwargs,
-        }
+        params = DiffusionSamplingParams(kwargs)
        if seed >= 0:
            params["seed"] = seed
        return (params,)
@@ -542,6 +634,52 @@ class VLLMOmniSamplingParamsList:
        return (params,)


+class VLLMOmniRemoteLoRA:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "local_path": ("STRING", {"default": ""}),
+                "name": ("STRING", {"default": ""}),
+                "scale": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
+                ),
+                "int_id": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "step": 1,
+                        "tooltip": "0 means it is not set and the server can derive it.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("REMOTE_LORA",)
+    RETURN_NAMES = ("lora",)
+    FUNCTION = "get_lora"
+    CATEGORY = "vLLM-Omni"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]:
+        if not local_path.strip() or not name.strip():
+            return "Both local_path and name must be provided."
+        return True
+
+    def get_lora(self, local_path: str, name: str, scale: float, int_id: int):
+        local_path = local_path.strip()
+        name = name.strip()
+        lora = {
+            "local_path": local_path or None,
+            "name": name or None,
+            "scale": float(scale),
+            "int_id": int(int_id) if int_id > 0 else None,
+        }
+        return (lora,)
+
+
 class VLLMOmniQwenTTSParams:
    @classmethod
    def INPUT_TYPES(cls):
@@ -566,4 +704,33 @@ class VLLMOmniQwenTTSParams:
    CATEGORY = "vLLM-Omni/TTS Params"

    def get_params(self, **kwargs):
-        return ({"type": "qwen-tts", **kwargs},)
+        return (QwenTTSModelSpecificParams(kwargs),)
+
+
+class VLLMOmniWanParams:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "guidance_scale_2": (
+                    "FLOAT",
+                    {"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1},
+                ),
+                "boundary_ratio": (
+                    "FLOAT",
+                    {"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001},
+                ),
+                "flow_shift": (
+                    "FLOAT",
+                    {"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("VIDEO_PARAMS",)
+    RETURN_NAMES = ("Wan video params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Video Params"
+
+    def get_params(self, **kwargs):
+        return (WanModelSpecificParams(kwargs),)
@@ -6,6 +6,8 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
 Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
 """

+import asyncio
+import json
 from typing import Any

 import aiohttp
@@ -18,6 +20,7 @@ from .format import (
    base64_to_audio,
    base64_to_image_tensor,
    bytes_to_audio,
+    bytes_to_video,
    image_tensor_to_base64,
    image_tensor_to_png_bytes,
    video_to_base64,
@@ -29,10 +32,43 @@ from .types import AudioFormat
 logger = get_logger(__name__)


+async def url_json(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> dict[str, Any]:
+    try:
+        async with getattr(session, verb)(url, **kwargs) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise (ValueError if response.status < 500 else RuntimeError)(
+                    f"vLLM-Omni API returned status {response.status}: {error_text}"
+                )
+            try:
+                return await response.json()
+            except aiohttp.ContentTypeError as e:
+                raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+    except aiohttp.ClientError as e:
+        raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+
+async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> bytes:
+    try:
+        async with getattr(session, verb)(url, **kwargs) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise (ValueError if response.status < 500 else RuntimeError)(
+                    f"vLLM-Omni API returned status {response.status}: {error_text}"
+                )
+            return await response.read()
+    except aiohttp.ClientError as e:
+        raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+
 class VLLMOmniClient:
-    def __init__(self, base_url: str, timeout: float = 300.0):
+    def __init__(
+        self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
+    ):
        self.base_url = base_url
        self.timeout = aiohttp.ClientTimeout(total=timeout)
+        self.poll_interval = poll_interval
+        self.max_poll_duration = max_poll_duration

    async def generate_image(
        self,
@@ -43,12 +79,13 @@ class VLLMOmniClient:
        height: int,
        negative_prompt: str | None = None,
        sampling_params: dict | None = None,
+        lora: dict | None = None,
    ) -> torch.Tensor:
        """Run text-to-image generation via DALLE API"""
        await self._check_model_exist(model)

        size = f"{width}x{height}"
-        payload = {
+        payload: dict[str, Any] = {
            "model": model,
            "prompt": prompt,
            "size": size,
@@ -57,16 +94,9 @@ class VLLMOmniClient:
        if negative_prompt:
            payload["negative_prompt"] = negative_prompt
        if sampling_params is not None:
-            # Only select specific sampling params
-            for k in (
-                "n",
-                "num_inference_steps",
-                "guidance_scale",
-                "true_cfg_scale",
-                "vae_use_slicing",
-            ):
-                if k in sampling_params and sampling_params[k] is not None:
-                    payload[k] = sampling_params[k]
+            payload.update(sampling_params)
+        if lora is not None:
+            payload["lora"] = lora
        logger.debug("img gen payload: %s", payload)

        url = self.base_url + "/images/generations"
@@ -119,6 +149,7 @@ class VLLMOmniClient:
        negative_prompt: str | None = None,
        mask: torch.Tensor | None = None,
        sampling_params: dict | None = None,
+        lora: dict | None = None,
    ) -> torch.Tensor:
        """Run image editing via DALLE API"""
        await self._check_model_exist(model)
@@ -138,10 +169,10 @@ class VLLMOmniClient:
        if negative_prompt:
            form.add_field("negative_prompt", negative_prompt)
        if sampling_params is not None:
-            # Only select specific sampling params
-            for k in ("n", "num_inference_steps", "guidance_scale", "true_cfg_scale"):
-                if k in sampling_params and sampling_params[k] is not None:
-                    form.add_field(k, str(sampling_params[k]))
+            for k, v in sampling_params.items():
+                form.add_field(k, str(v))
+        if lora is not None:
+            form.add_field("lora", json.dumps(lora, ensure_ascii=False))
        if mask is not None:
            mask_filename = "mask.png"
            form.add_field(
@@ -194,6 +225,7 @@ class VLLMOmniClient:
        audio: AudioInput | None = None,
        video: VideoInput | None = None,
        sampling_params: dict | list[dict] | None = None,
+        lora: dict | None = None,
    ) -> torch.Tensor:
        payload = VLLMOmniClient._prepare_chat_completion_messages(
            model=model,
@@ -204,6 +236,8 @@ class VLLMOmniClient:
            video=video,
            sampling_params=sampling_params,
            modalities=["image"],
+            # === below are additional `extra_body` fields, handled by **kwargs ===
+            lora=lora,
        )
        choices = await self._generate_base_chat_completion(model, payload)

@@ -217,7 +251,88 @@ class VLLMOmniClient:

        return torch.stack(image_tensors, dim=0)

-    async def generate_comprehension_chat_completion(
+    async def generate_video(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        num_frames: int,
+        fps: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        sampling_params: dict | None = None,
+        model_params: dict | None = None,
+        lora: dict | None = None,
+        **extra_body,
+    ) -> VideoInput:
+        form = aiohttp.FormData()
+        form.add_field("model", model)
+        form.add_field("prompt", prompt)
+        form.add_field("width", str(width))
+        form.add_field("height", str(height))
+        form.add_field("num_frames", str(num_frames))
+        form.add_field("fps", str(fps))
+        if negative_prompt:
+            form.add_field("negative_prompt", negative_prompt)
+        if sampling_params is not None:
+            for k, v in sampling_params.items():
+                form.add_field(k, str(v))
+        if model_params is not None:
+            for k, v in model_params.items():
+                form.add_field(k, str(v))
+        if lora is not None:
+            form.add_field("lora", json.dumps(lora, ensure_ascii=False))
+        if extra_body:
+            form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
+
+        if image is not None:
+            image_filename = "image.png"  # Required for multipart form
+            form.add_field(
+                "input_reference",
+                image_tensor_to_png_bytes(image, image_filename),
+                filename=image_filename,
+                content_type="image/png",
+            )
+
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            # Start the video generation job
+            url = f"{self.base_url}/videos"
+            data = await url_json(session, url, "post", data=form)
+            if (job_id := data.get("id", None)) is None:
+                raise RuntimeError("API response missing job 'id' field - expected OpenAI compliant format")
+            if (job_status := data.get("status", None)) is None:
+                raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
+
+            # Poll for video generation job completion
+            deadline = asyncio.get_running_loop().time() + self.max_poll_duration
+            url = f"{self.base_url}/videos/{job_id}"
+            while job_status not in {"completed", "failed"}:
+                await asyncio.sleep(self.poll_interval)
+
+                data = await url_json(session, url)
+                if (job_status := data.get("status", None)) is None:
+                    raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
+                if asyncio.get_running_loop().time() >= deadline:
+                    raise RuntimeError(f"Timed out waiting for video job {job_id} to complete")
+
+            if job_status == "failed":
+                raise RuntimeError(f"Video job failed: {data}")
+
+            # Retrieve completed content
+            video_bytes = await url_bytes(session, f"{url}/content")
+
+            # Decode video and make a best effort at cleaning up server resources
+            try:
+                return bytes_to_video(video_bytes)
+            finally:
+                try:
+                    await url_json(session, url, "delete")
+                except Exception as exc:
+                    logger.warning("Failed to clean up video job %s: %s", job_id, exc)
+
+    async def generate_understanding_chat_completion(
        self,
        *,
        model: str,
@@ -427,30 +542,28 @@ class VLLMOmniClient:
            message_content.append({"type": "video_url", "video_url": {"url": video_to_base64(video)}})
        messages = [{"role": "user", "content": message_content}]

+        payload: dict[str, Any] = {"messages": messages, "model": model}
+        if modalities:
+            payload["modalities"] = modalities
+
        combined_extra_body: dict[str, Any] = {}
        if sampling_params is not None:
            spec, _ = lookup_model_spec(model)
            is_single_sampling_param = isinstance(sampling_params, dict) or len(sampling_params) == 1

-            # Exclude internal key
-            if isinstance(sampling_params, dict):
-                sampling_params = {k: v for k, v in sampling_params.items() if k != "type"}
-            else:
-                sampling_params = [{k: v for k, v in sp.items() if k != "type"} for sp in sampling_params]
-
            if (spec is None and is_single_sampling_param) or (spec is not None and spec["stages"] == ["diffusion"]):
                # Diffusion format: extra_body directly contains sampling params.
-                # Validation should have taken care of matching sampling params' types.
+                # Validation has already taken care of matching sampling params' types and length. Safe to take [0].
                # * Use this mode if the model is a simple one-stage diffusion model.
                # * Fallback to this mode if model is not registered and a single sampling param is provided.
                sampling_params = sampling_params if isinstance(sampling_params, dict) else sampling_params[0]
-                combined_extra_body: dict[str, Any] = {**sampling_params}
+                combined_extra_body: dict[str, Any] = sampling_params.copy()
                if "n" in combined_extra_body:
-                    combined_extra_body["num_outputs_per_prompt"] = combined_extra_body["n"]
-                    del combined_extra_body["n"]
+                    combined_extra_body["num_outputs_per_prompt"] = combined_extra_body.pop("n")
            else:
-                # Use AR style payload, extra_body has a sampling_params_list field
-                combined_extra_body: dict[str, Any] = {"sampling_params_list": sampling_params}
+                # AR format: the payload has a sampling_params_list field, containing a list.
+                sampling_params_list = sampling_params if isinstance(sampling_params, list) else [sampling_params]
+                payload["sampling_params_list"] = sampling_params_list

        if negative_prompt:
            combined_extra_body["negative_prompt"] = negative_prompt
@@ -458,12 +571,11 @@ class VLLMOmniClient:
        if extra_body:
            combined_extra_body.update(extra_body)

-        payload: dict[str, Any] = {"messages": messages, "model": model}
+        # Add extra_body only if it has any content.
        if combined_extra_body:
            payload["extra_body"] = combined_extra_body
-        if modalities:
-            payload["modalities"] = modalities

+        # Place to inject any model-specific payload adjustment
        spec, _ = lookup_model_spec(model)
        if spec:
            preprocessor = spec.get("payload_preprocessor", None)
@@ -6,12 +6,17 @@ Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed und

 import base64
 import mimetypes
+from fractions import Fraction
 from io import BytesIO

 import av
 import numpy as np
 import torch
+from av.audio.frame import AudioFrame
+from av.audio.resampler import AudioResampler
+from av.video.frame import VideoFrame
 from comfy_api.input import AudioInput, VideoInput
+from comfy_api.latest import InputImpl, Types
 from comfy_extras import nodes_audio
 from PIL import Image

@@ -138,6 +143,75 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
    return f"data:{mime_type};base64,{base64_str}"


+def bytes_to_video(video_bytes: bytes) -> VideoInput:
+    video_buffer = BytesIO(video_bytes)
+
+    try:
+        with av.open(video_buffer, mode="r") as container:
+            video_stream = next((s for s in container.streams if s.type == "video"), None)
+            if video_stream is None:
+                raise ValueError("No video stream found in decoded payload.")
+
+            frames = []
+            for frame in container.decode(video_stream):
+                if not isinstance(frame, VideoFrame):
+                    continue
+                image = frame.to_ndarray(format="rgb24")
+                frames.append(torch.from_numpy(image).float() / 255.0)
+
+            if len(frames) == 0:
+                raise ValueError("No video frames found in decoded payload.")
+
+            images = torch.stack(frames, dim=0)
+            frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
+
+            audio: AudioInput | None = None
+            if len(container.streams.audio):
+                audio_stream = container.streams.audio[-1]
+                audio_frames = []
+                resampler = AudioResampler(format="fltp")
+                for frame in container.decode(audio_stream):
+                    if not isinstance(frame, AudioFrame):
+                        continue
+                    resampled = resampler.resample(frame)
+                    if not isinstance(resampled, list):
+                        resampled = [resampled]
+                    for audio_frame in resampled:
+                        if audio_frame is not None:
+                            audio_frames.append(audio_frame.to_ndarray())
+
+                if len(audio_frames) > 0:
+                    audio_data = np.concatenate(audio_frames, axis=1)
+                    sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1
+                    audio = {
+                        "waveform": torch.from_numpy(audio_data).unsqueeze(0),
+                        "sample_rate": sample_rate,
+                    }
+
+            components = Types.VideoComponents(
+                images=images,
+                frame_rate=frame_rate,
+                audio=audio,
+                metadata=container.metadata if container.metadata else None,
+            )
+    except Exception as e:
+        raise RuntimeError(f"Failed to decode video: {e}")
+
+    return InputImpl.VideoFromComponents(components)
+
+
+def base64_to_video(base64_str: str) -> VideoInput:
+    if base64_str.startswith("data:video"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        video_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    return bytes_to_video(video_bytes)
+
+
 def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
    waveform = audio["waveform"][0]  # Shape: (C, T)
    sample_rate = audio["sample_rate"]
@@ -11,8 +11,7 @@ def _bagel_payload_preprocessor(payload: dict) -> dict:
                    content["text"] = "<|im_start|>" + content["text"] + "<|im_end|>"
    except (KeyError, TypeError):
        raise RuntimeError("Internal Error: malformatted BAGEL payload")
-    extra_body = payload.pop("extra_body", {})
-    return {**payload, **extra_body}
+    return payload


 def _qwen25_payload_preprocessor(payload: dict) -> dict:
@@ -37,7 +36,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
        ],
        "modes": [
            {
-                "mode": ModelMode.COMPREHENSION,
+                "mode": ModelMode.UNDERSTANDING,
                "input_modalities": [Modality.TEXT, Modality.IMAGE],
            }
        ],
@@ -48,7 +47,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
        "payload_preprocessor": _qwen25_payload_preprocessor,
        "modes": [
            {
-                "mode": ModelMode.COMPREHENSION,
+                "mode": ModelMode.UNDERSTANDING,
                "input_modalities": [
                    Modality.TEXT,
                    Modality.IMAGE,
@@ -62,7 +61,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
        "stages": ["autoregression", "autoregression", "autoregression"],
        "modes": [
            {
-                "mode": ModelMode.COMPREHENSION,
+                "mode": ModelMode.UNDERSTANDING,
                "input_modalities": [
                    Modality.TEXT,
                    Modality.IMAGE,
@@ -11,11 +11,27 @@ from typing import (
 AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]


+class AutoregressionSamplingParams(dict):
+    pass
+
+
+class DiffusionSamplingParams(dict):
+    pass
+
+
+class QwenTTSModelSpecificParams(dict):
+    pass
+
+
+class WanModelSpecificParams(dict):
+    pass
+
+
 class ModelMode(Enum):
    IMAGE_GENERATION = auto()
    VIDEO_GENERATION = auto()
    AUDIO_GENERATION = auto()
-    COMPREHENSION = auto()
+    UNDERSTANDING = auto()


 class Modality(Enum):
@@ -1,5 +1,6 @@
 from .logger import get_logger
 from .models import lookup_model_spec
+from .types import AutoregressionSamplingParams, DiffusionSamplingParams

 logger = get_logger(__name__)

@@ -15,7 +16,7 @@ def validate_model_and_sampling_params_types(
    # Skip if no spec or no sampling params
    pipeline_spec, _ = lookup_model_spec(model_name)
    if pipeline_spec is None:
-        logger.info("skipping sampling params check because spec is not found")
+        logger.info(f"skipping sampling params check because spec for {model_name} is not found")
        return
    if sampling_param_list is None:
        return
@@ -31,24 +32,20 @@ def validate_model_and_sampling_params_types(
            )
        # Check that each stage's type match
        for i, sp in enumerate(sampling_param_list):
-            if "type" not in sp:
-                raise RuntimeError("Internal error: unknown sampling parameter type")
-            if sp["type"] != stages[i]:
+            if not _check_sampling_param_matches_stage(sp, stages[i]):
                raise ValueError(
-                    f"Sampling parameter type ({sp['type']}) does not match "
+                    f"Sampling parameter type ({sp.__class__.__name__}) does not match "
                    f"stage type ({stages[i]}) at index {i} for model {model_name}."
                )
    elif isinstance(sampling_param_list, dict):
-        if "type" not in sampling_param_list:
-            raise RuntimeError("Internal error: unknown sampling parameter type")
        # Check that the provided single sampling param matches all stages
-        elif any(stage != sampling_param_list["type"] for stage in stages):
-            raise ValueError(
-                f"When passing a single sampling parameter node, all stages of the model must match "
-                f"the provided sampling parameter's type. "
-                f"However, the stages of model {model_name} are: {stages}. "
-                f"The provided sampling parameter is {sampling_param_list['type']}"
-            )
+        for i, stage in enumerate(stages):
+            if not _check_sampling_param_matches_stage(sampling_param_list, stage):
+                raise ValueError(
+                    f"Provided single sampling parameter type ({sampling_param_list.__class__.__name__}) must match "
+                    f"the types of all stages of the model. "
+                    f"However, stage {i} of model {model_name} is of type {stage}."
+                )


 def add_sampling_parameters_to_stage(
@@ -73,12 +70,24 @@ def add_sampling_parameters_to_stage(

    stages = pipeline_spec["stages"]
    if isinstance(sampling_param_list, dict):
+        sampling_param_list = sampling_param_list.__class__(sampling_param_list)
        sampling_param_list.update(params_to_add)
    elif sampling_param_list is None:
        sampling_param_list = params_to_add.copy()
    else:
        for i, stage in enumerate(stages):
            if stage == stage_type:
-                sampling_param_list[i].update(params_to_add)
+                stage_param = sampling_param_list[i]
+                stage_param = stage_param.__class__(stage_param)
+                stage_param.update(params_to_add)
+                sampling_param_list[i] = stage_param

    return sampling_param_list
+
+
+def _check_sampling_param_matches_stage(sampling_param: dict, stage_type: str) -> bool:
+    if stage_type == "autoregression":
+        return isinstance(sampling_param, AutoregressionSamplingParams)
+    if stage_type == "diffusion":
+        return isinstance(sampling_param, DiffusionSamplingParams)
+    raise RuntimeError(f"Internal error: unknown stage type {stage_type}.")
@@ -1 +0,0 @@
-{"id":"1c99f525-0a37-45ba-a28a-7df7c3af66b4","revision":0,"last_node_id":12,"last_link_id":14,"nodes":[{"id":1,"type":"VLLMOmniComprehension","pos":[1191.2177053682556,144.66829928181377],"size":[400,268],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":1},{"localized_name":"video","name":"video","shape":7,"type":"VIDEO","link":2},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":3},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":14},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"output_text","name":"output_text","type":"BOOLEAN","widget":{"name":"output_text"},"link":null},{"localized_name":"output_audio","name":"output_audio","type":"BOOLEAN","widget":{"name":"output_audio"},"link":null},{"localized_name":"use_audio_in_video","name":"use_audio_in_video","type":"BOOLEAN","widget":{"name":"use_audio_in_video"},"link":null}],"outputs":[{"localized_name":"text_response","name":"text_response","type":"STRING","links":[8]},{"localized_name":"audio_response","name":"audio_response","type":"AUDIO","links":[9]}],"properties":{"Node name for S&R":"VLLMOmniComprehension"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen2.5-Omni-7B","",true,true,true]},{"id":3,"type":"LoadVideo","pos":[729.5984141255855,-198.631920454299],"size":[282.798828125,233.0743408203125],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[2]}],"properties":{"Node name for S&R":"LoadVideo"},"widgets_values":["draw.mp4","image"]},{"id":4,"type":"LoadAudio","pos":[729.8037086965753,99.86963519703949],"size":[282.798828125,136],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":5,"type":"VLLMOmniARSampling","pos":[510.3517536642828,658.073751009259],"size":[270,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[12]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":7,"type":"VLLMOmniARSampling","pos":[503.33235181647115,419.34158016181806],"size":[270,178],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[5,13]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":8,"type":"VLLMOmniSamplingParamsList","pos":[820.6056617389042,426.38372037182273],"size":[263.066015625,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"param1","name":"param1","type":"SAMPLING_PARAMS","link":5},{"localized_name":"param2","name":"param2","shape":7,"type":"SAMPLING_PARAMS","link":12},{"localized_name":"param3","name":"param3","shape":7,"type":"SAMPLING_PARAMS","link":13}],"outputs":[{"localized_name":"param list","name":"param list"Line truncated
@@ -0,0 +1,552 @@
+{
+    "id": "6643e5fd-fa2a-4f25-935a-483173a8097c",
+    "revision": 0,
+    "last_node_id": 8,
+    "last_link_id": 6,
+    "nodes": [
+        {
+            "id": 1,
+            "type": "PreviewImage",
+            "pos": [
+                1446.2005205859512,
+                -316.7359686049902
+            ],
+            "size": [
+                305.6628685610174,
+                307.0336884172169
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "images",
+                    "name": "images",
+                    "type": "IMAGE",
+                    "link": 1
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewImage"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 8,
+            "type": "PreviewImage",
+            "pos": [
+                1010.1848219273943,
+                -311.54216706114363
+            ],
+            "size": [
+                305.6628685610174,
+                307.0336884172169
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "images",
+                    "name": "images",
+                    "type": "IMAGE",
+                    "link": 6
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewImage"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 7,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                234.61958293819296,
+                9.405681270043619
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "n",
+                    "name": "n",
+                    "type": "INT",
+                    "widget": {
+                        "name": "n"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_inference_steps",
+                    "name": "num_inference_steps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_inference_steps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "guidance_scale",
+                    "name": "guidance_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "true_cfg_scale",
+                    "name": "true_cfg_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "true_cfg_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_slicing",
+                    "name": "vae_use_slicing",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_slicing"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "diffusion sampling params",
+                    "name": "diffusion sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        5
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                1,
+                50,
+                1,
+                1,
+                false,
+                false,
+                1525,
+                "randomize"
+            ]
+        },
+        {
+            "id": 4,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                666.8380026154548,
+                268.86271068330126
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 1,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "n",
+                    "name": "n",
+                    "type": "INT",
+                    "widget": {
+                        "name": "n"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_inference_steps",
+                    "name": "num_inference_steps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_inference_steps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "guidance_scale",
+                    "name": "guidance_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "true_cfg_scale",
+                    "name": "true_cfg_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "true_cfg_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_slicing",
+                    "name": "vae_use_slicing",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_slicing"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "diffusion sampling params",
+                    "name": "diffusion sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        2
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                4,
+                50,
+                7,
+                1,
+                false,
+                false,
+                42,
+                "fixed"
+            ]
+        },
+        {
+            "id": 5,
+            "type": "VLLMOmniGenerateImage",
+            "pos": [
+                984.723613585788,
+                63.376900027553276
+            ],
+            "size": [
+                416.56628685610167,
+                372.1662621294205
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 4
+                },
+                {
+                    "localized_name": "mask",
+                    "name": "mask",
+                    "shape": 7,
+                    "type": "MASK",
+                    "link": null
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 2
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "negative_prompt",
+                    "name": "negative_prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "negative_prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "width",
+                    "name": "width",
+                    "type": "INT",
+                    "widget": {
+                        "name": "width"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "height",
+                    "name": "height",
+                    "type": "INT",
+                    "widget": {
+                        "name": "height"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "IMAGE",
+                    "links": [
+                        1
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateImage"
+            },
+            "widgets_values": [
+                "http://localhost:8001/v1",
+                "/home/models/Qwen/Qwen-Image-Edit",
+                "A high-quality, high contrast, stylized portrait of the object in the uploaded reference image. Pop art and doodle style, with abundant scribbling patterns, such as a teal crown, orange lightning bolts, colorful handwritten scripts.",
+                "Realistic",
+                800,
+                800
+            ]
+        },
+        {
+            "id": 6,
+            "type": "VLLMOmniGenerateImage",
+            "pos": [
+                541.2702717429818,
+                -160.70392744708548
+            ],
+            "size": [
+                416.56628685610167,
+                372.1662621294205
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": null
+                },
+                {
+                    "localized_name": "mask",
+                    "name": "mask",
+                    "shape": 7,
+                    "type": "MASK",
+                    "link": null
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 5
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "negative_prompt",
+                    "name": "negative_prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "negative_prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "width",
+                    "name": "width",
+                    "type": "INT",
+                    "widget": {
+                        "name": "width"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "height",
+                    "name": "height",
+                    "type": "INT",
+                    "widget": {
+                        "name": "height"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "IMAGE",
+                    "links": [
+                        4,
+                        6
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateImage"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "/home/models/Tongyi-MAI/Z-Image-Turbo",
+                "A headshot of a cute Siamese kitty. Blurred background due to wide aperture. Close-up look. Realistic.",
+                "Cartoonish.",
+                800,
+                800
+            ]
+        }
+    ],
+    "links": [
+        [
+            1,
+            5,
+            0,
+            1,
+            0,
+            "IMAGE"
+        ],
+        [
+            2,
+            4,
+            0,
+            5,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            4,
+            6,
+            0,
+            5,
+            0,
+            "IMAGE"
+        ],
+        [
+            5,
+            7,
+            0,
+            6,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            6,
+            6,
+            0,
+            8,
+            0,
+            "IMAGE"
+        ]
+    ],
+    "groups": [],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 1.0426432563169903,
+            "offset": [
+                -90.4071374799346,
+                575.2948920069742
+            ]
+        }
+    },
+    "version": 0.4
+}
@@ -1 +1 @@
-{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,202],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,42,"fixed"]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[276.8427259980847,549.4885226647582]}},"version":0.4}
+{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"图像","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"图像","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"选择文件上传","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"图像","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"遮罩","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,226],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"vae_use_tiling","name":"vae_use_tiling","type":"BOOLEAN","widget":{"name":"vae_use_tiling"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,false,42,"randomize"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[-55.18047513099182,220.2553505195962]}},"version":0.4}
@@ -0,0 +1,761 @@
+{
+    "id": "1c99f525-0a37-45ba-a28a-7df7c3af66b4",
+    "revision": 0,
+    "last_node_id": 12,
+    "last_link_id": 14,
+    "nodes": [
+        {
+            "id": 1,
+            "type": "VLLMOmniUnderstanding",
+            "pos": [
+                1191.2177053682556,
+                144.66829928181377
+            ],
+            "size": [
+                400,
+                268
+            ],
+            "flags": {},
+            "order": 8,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 1
+                },
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "shape": 7,
+                    "type": "VIDEO",
+                    "link": 2
+                },
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "shape": 7,
+                    "type": "AUDIO",
+                    "link": 3
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 14
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "output_text",
+                    "name": "output_text",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "output_text"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "output_audio",
+                    "name": "output_audio",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "output_audio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "use_audio_in_video",
+                    "name": "use_audio_in_video",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "use_audio_in_video"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "text_response",
+                    "name": "text_response",
+                    "type": "STRING",
+                    "links": [
+                        8
+                    ]
+                },
+                {
+                    "localized_name": "audio_response",
+                    "name": "audio_response",
+                    "type": "AUDIO",
+                    "links": [
+                        9
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniUnderstanding"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "Qwen/Qwen2.5-Omni-7B",
+                "",
+                true,
+                true,
+                true
+            ]
+        },
+        {
+            "id": 3,
+            "type": "LoadVideo",
+            "pos": [
+                729.5984141255855,
+                -198.631920454299
+            ],
+            "size": [
+                282.798828125,
+                233.0743408203125
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "file",
+                    "name": "file",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "file"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "IMAGEUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "VIDEO",
+                    "name": "VIDEO",
+                    "type": "VIDEO",
+                    "links": [
+                        2
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadVideo"
+            },
+            "widgets_values": [
+                "draw.mp4",
+                "image"
+            ]
+        },
+        {
+            "id": 4,
+            "type": "LoadAudio",
+            "pos": [
+                729.8037086965753,
+                99.86963519703949
+            ],
+            "size": [
+                282.798828125,
+                136
+            ],
+            "flags": {},
+            "order": 1,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "audio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "audioUI",
+                    "name": "audioUI",
+                    "type": "AUDIO_UI",
+                    "widget": {
+                        "name": "audioUI"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "AUDIOUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AUDIO",
+                    "name": "AUDIO",
+                    "type": "AUDIO",
+                    "links": [
+                        3
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadAudio"
+            },
+            "widgets_values": [
+                "Megan-Fox.mp3",
+                null,
+                null
+            ]
+        },
+        {
+            "id": 5,
+            "type": "VLLMOmniARSampling",
+            "pos": [
+                510.3517536642828,
+                658.073751009259
+            ],
+            "size": [
+                270,
+                178
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "max_tokens",
+                    "name": "max_tokens",
+                    "type": "INT",
+                    "widget": {
+                        "name": "max_tokens"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "temperature",
+                    "name": "temperature",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "temperature"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "top_p",
+                    "name": "top_p",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "top_p"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "repetition_penalty",
+                    "name": "repetition_penalty",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "repetition_penalty"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AR sampling params",
+                    "name": "AR sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        12
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniARSampling"
+            },
+            "widgets_values": [
+                100,
+                1,
+                1,
+                1,
+                -1,
+                "randomize"
+            ]
+        },
+        {
+            "id": 7,
+            "type": "VLLMOmniARSampling",
+            "pos": [
+                503.33235181647115,
+                419.34158016181806
+            ],
+            "size": [
+                270,
+                178
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "max_tokens",
+                    "name": "max_tokens",
+                    "type": "INT",
+                    "widget": {
+                        "name": "max_tokens"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "temperature",
+                    "name": "temperature",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "temperature"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "top_p",
+                    "name": "top_p",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "top_p"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "repetition_penalty",
+                    "name": "repetition_penalty",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "repetition_penalty"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "AR sampling params",
+                    "name": "AR sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        5,
+                        13
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniARSampling"
+            },
+            "widgets_values": [
+                100,
+                1,
+                1,
+                1,
+                -1,
+                "randomize"
+            ]
+        },
+        {
+            "id": 8,
+            "type": "VLLMOmniSamplingParamsList",
+            "pos": [
+                820.6056617389042,
+                426.38372037182273
+            ],
+            "size": [
+                263.066015625,
+                66
+            ],
+            "flags": {},
+            "order": 7,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "param1",
+                    "name": "param1",
+                    "type": "SAMPLING_PARAMS",
+                    "link": 5
+                },
+                {
+                    "localized_name": "param2",
+                    "name": "param2",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 12
+                },
+                {
+                    "localized_name": "param3",
+                    "name": "param3",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 13
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "param list",
+                    "name": "param list",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        14
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniSamplingParamsList"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 11,
+            "type": "MarkdownNote",
+            "pos": [
+                826.2328280438272,
+                569.1890318701705
+            ],
+            "size": [
+                333.8220435590464,
+                261.63596728060656
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Sampling Parameters",
+            "properties": {},
+            "widgets_values": [
+                "## Sampling Parameter Types\n\nThere are two types of sampling parameters: one for autoregression and one for diffusion.\nYou should ensure that you have chosen the correct type of sampling parameters for the model you request.\n\n## Stages & Shorthand\n\nFor multi-stage models such as Qwen Omni, you can either\n- connect one sampling parameter node, which is applied to all stages.\n- connect exactly the same number of sampling parameter nodes to a \"Multi-Stage Sampling Parameter List\" node, then connect this node to the primary request node.\n\nNote that this shorthand is intended to stay consistent with the [online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/qwen2_5_omni/)"
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 12,
+            "type": "MarkdownNote",
+            "pos": [
+                378.9866207003777,
+                152.59550252215752
+            ],
+            "size": [
+                319.7287574247016,
+                107.15904081906785
+            ],
+            "flags": {},
+            "order": 6,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Input",
+            "properties": {},
+            "widgets_values": [
+                "Note that not all models support every modality as input. For example, `ByteDance-Seed/BAGEL-7B-MoT` in Multimodality Understanding mode only support text and image input.\n\nYou should ensure that the input are supported by the model. You can check the corresponding [online serving documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/) for confirmation."
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 2,
+            "type": "LoadImage",
+            "pos": [
+                394.4674804308822,
+                -207.6987397548834
+            ],
+            "size": [
+                282.798828125,
+                314
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "image"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "IMAGEUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "IMAGE",
+                    "name": "IMAGE",
+                    "type": "IMAGE",
+                    "links": [
+                        1
+                    ]
+                },
+                {
+                    "localized_name": "MASK",
+                    "name": "MASK",
+                    "type": "MASK",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadImage"
+            },
+            "widgets_values": [
+                "example.png",
+                "image"
+            ]
+        },
+        {
+            "id": 10,
+            "type": "PreviewAudio",
+            "pos": [
+                1664.548345556043,
+                297.5921292054054
+            ],
+            "size": [
+                270,
+                88
+            ],
+            "flags": {},
+            "order": 10,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "type": "AUDIO",
+                    "link": 9
+                },
+                {
+                    "localized_name": "audioUI",
+                    "name": "audioUI",
+                    "type": "AUDIO_UI",
+                    "widget": {
+                        "name": "audioUI"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewAudio"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 9,
+            "type": "ShowText|pysssss",
+            "pos": [
+                1649.2506875091847,
+                66.22823888292349
+            ],
+            "size": [
+                318.7188464232943,
+                173.38502269972975
+            ],
+            "flags": {},
+            "order": 9,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "text",
+                    "name": "text",
+                    "type": "STRING",
+                    "link": 8
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "STRING",
+                    "name": "STRING",
+                    "shape": 6,
+                    "type": "STRING",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "ShowText|pysssss"
+            },
+            "widgets_values": []
+        }
+    ],
+    "links": [
+        [
+            1,
+            2,
+            0,
+            1,
+            0,
+            "IMAGE"
+        ],
+        [
+            2,
+            3,
+            0,
+            1,
+            1,
+            "VIDEO"
+        ],
+        [
+            3,
+            4,
+            0,
+            1,
+            2,
+            "AUDIO"
+        ],
+        [
+            5,
+            7,
+            0,
+            8,
+            0,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            8,
+            1,
+            0,
+            9,
+            0,
+            "STRING"
+        ],
+        [
+            9,
+            1,
+            1,
+            10,
+            0,
+            "AUDIO"
+        ],
+        [
+            12,
+            5,
+            0,
+            8,
+            1,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            13,
+            7,
+            0,
+            8,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            14,
+            8,
+            0,
+            1,
+            3,
+            "SAMPLING_PARAMS"
+        ]
+    ],
+    "groups": [
+        {
+            "id": 1,
+            "title": "Sampling Parameters",
+            "bounding": [
+                480.1649301181556,
+                341.08402513937995,
+                692.4113568972277,
+                510.48648853403665
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        },
+        {
+            "id": 2,
+            "title": "Input",
+            "bounding": [
+                344.2364817813528,
+                -287.5850484183313,
+                704.8832238768634,
+                559.5124009832894
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        }
+    ],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 0.9478575057427204,
+            "offset": [
+                33.81037136029557,
+                307.1974296197726
+            ]
+        }
+    },
+    "version": 0.4
+}
@@ -0,0 +1,513 @@
+{
+    "id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
+    "revision": 0,
+    "last_node_id": 14,
+    "last_link_id": 28,
+    "nodes": [
+        {
+            "id": 10,
+            "type": "MarkdownNote",
+            "pos": [
+                227.99306819575278,
+                -231.24143306069843
+            ],
+            "size": [
+                240.3326478922474,
+                136.3505820791881
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Task and Input",
+            "properties": {},
+            "widgets_values": [
+                "vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 4,
+            "type": "LoadImage",
+            "pos": [
+                496.31859627609606,
+                -229.71277089860084
+            ],
+            "size": [
+                270,
+                314
+            ],
+            "flags": {},
+            "order": 1,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "image"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "choose file to upload",
+                    "name": "upload",
+                    "type": "IMAGEUPLOAD",
+                    "widget": {
+                        "name": "upload"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "IMAGE",
+                    "name": "IMAGE",
+                    "type": "IMAGE",
+                    "links": [
+                        25
+                    ]
+                },
+                {
+                    "localized_name": "MASK",
+                    "name": "MASK",
+                    "type": "MASK",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadImage"
+            },
+            "widgets_values": [
+                "cute-cat.jpg",
+                "image"
+            ]
+        },
+        {
+            "id": 8,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                465.5140218220927,
+                148.80072646828987
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "n",
+                    "name": "n",
+                    "type": "INT",
+                    "widget": {
+                        "name": "n"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_inference_steps",
+                    "name": "num_inference_steps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_inference_steps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "guidance_scale",
+                    "name": "guidance_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "true_cfg_scale",
+                    "name": "true_cfg_scale",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "true_cfg_scale"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_slicing",
+                    "name": "vae_use_slicing",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_slicing"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "seed",
+                    "name": "seed",
+                    "type": "INT",
+                    "widget": {
+                        "name": "seed"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "diffusion sampling params",
+                    "name": "diffusion sampling params",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        27
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                1,
+                50,
+                1,
+                1,
+                false,
+                false,
+                1992,
+                "randomize"
+            ]
+        },
+        {
+            "id": 14,
+            "type": "VLLMOmniWanParams",
+            "pos": [
+                518.6005521845747,
+                428.4455663719664
+            ],
+            "size": [
+                270,
+                106
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "guidance_scale_2",
+                    "name": "guidance_scale_2",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale_2"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "boundary_ratio",
+                    "name": "boundary_ratio",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "boundary_ratio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "flow_shift",
+                    "name": "flow_shift",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "flow_shift"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "Wan video params",
+                    "name": "Wan video params",
+                    "type": "VIDEO_PARAMS",
+                    "links": [
+                        28
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniWanParams"
+            },
+            "widgets_values": [
+                4,
+                0.875,
+                5
+            ]
+        },
+        {
+            "id": 12,
+            "type": "VLLMOmniGenerateVideo",
+            "pos": [
+                827.2566336087867,
+                -73.77449831827545
+            ],
+            "size": [
+                400,
+                346
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 25
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 27
+                },
+                {
+                    "localized_name": "lora",
+                    "name": "lora",
+                    "shape": 7,
+                    "type": "REMOTE_LORA",
+                    "link": null
+                },
+                {
+                    "localized_name": "model_params",
+                    "name": "model_params",
+                    "shape": 7,
+                    "type": "VIDEO_PARAMS",
+                    "link": 28
+                },
+                {
+                    "localized_name": "url",
+                    "name": "url",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "url"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "model",
+                    "name": "model",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "model"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "prompt",
+                    "name": "prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "negative_prompt",
+                    "name": "negative_prompt",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "negative_prompt"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "width",
+                    "name": "width",
+                    "type": "INT",
+                    "widget": {
+                        "name": "width"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "height",
+                    "name": "height",
+                    "type": "INT",
+                    "widget": {
+                        "name": "height"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "fps",
+                    "name": "fps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "fps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_frames",
+                    "name": "num_frames",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_frames"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "type": "VIDEO",
+                    "links": [
+                        26
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateVideo"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+                "Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
+                "Cartoonish",
+                900,
+                1200,
+                12,
+                41
+            ]
+        },
+        {
+            "id": 13,
+            "type": "SaveVideo",
+            "pos": [
+                1257.5441369501657,
+                -44.12956394419863
+            ],
+            "size": [
+                374,
+                372.34896826553796
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "type": "VIDEO",
+                    "link": 26
+                },
+                {
+                    "localized_name": "filename_prefix",
+                    "name": "filename_prefix",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "filename_prefix"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "format",
+                    "name": "format",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "format"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "codec",
+                    "name": "codec",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "codec"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [],
+            "properties": {},
+            "widgets_values": [
+                "video/ComfyUI",
+                "auto",
+                "auto"
+            ]
+        }
+    ],
+    "links": [
+        [
+            25,
+            4,
+            0,
+            12,
+            0,
+            "IMAGE"
+        ],
+        [
+            26,
+            12,
+            0,
+            13,
+            0,
+            "VIDEO"
+        ],
+        [
+            27,
+            8,
+            0,
+            12,
+            1,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            28,
+            14,
+            0,
+            12,
+            3,
+            "VIDEO_PARAMS"
+        ]
+    ],
+    "groups": [
+        {
+            "id": 1,
+            "title": "Input",
+            "bounding": [
+                213.1706010087147,
+                -313.20095750667554,
+                560.8246471437027,
+                407.4246532472182
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        }
+    ],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 1.1469075819486894,
+            "offset": [
+                101.7632950847089,
+                492.29121889347806
+            ]
+        }
+    },
+    "version": 0.4
+}
@@ -56,7 +56,7 @@ python extract_tts_prompts.py -i meta.lst -o top100.txt -n 100
 **Options:**
 - `-i, --input`: Input metadata file (default: `meta.lst`)
 - `-o, --output`: Output prompts file (default: `prompts.txt`)
- `-n, --num_lines`: Number of prompts to extract (required)
+- `-n, --num-lines`: Number of prompts to extract (required)

 ### 7. Clean Up (Optional)

@@ -56,7 +56,7 @@ def main():
        "-o", "--output", type=str, default="prompts.txt", help="Output txt file path (default: prompts.txt)"
    )
    parser.add_argument(
-        "-n", "--num_lines", type=int, required=True, help="Number of lines to extract from the beginning"
+        "-n", "--num-lines", type=int, required=True, help="Number of lines to extract from the beginning"
    )

    args = parser.parse_args()
@@ -43,6 +43,27 @@ The benchmark supports three dataset modes via `--dataset`:

 ### VBench dataset

+`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry
+per-request generation fields. In this mode, all requests share CLI values:
+`--width --height --num-frames --fps --num-inference-steps`
+(pass `--width` and `--height` together).
+
+Example (`t2v`):
+
+```bash
+python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+	--base-url http://localhost:8099 \
+	--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+	--task t2v \
+	--dataset vbench \
+	--num-prompts 50 \
+	--width 640 --height 480 \
+	--num-frames 81 --fps 16 \
+	--num-inference-steps 40
+```
+
+Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
+
 If you use i2v/i2i bench datasets and need auto-download support, you may need:

 ```bash
@@ -90,7 +111,7 @@ Common optional flags:
 Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.

 - For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
+- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.

 Precedence rules for `trace` (i.e., what actually gets sent):

@@ -0,0 +1,198 @@
+import base64
+import mimetypes
+import os
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any
+
+import aiohttp
+from tqdm import tqdm
+
+
+@dataclass
+class RequestFuncInput:
+    prompt: str
+    api_url: str
+    model: str
+    width: int | None = None
+    height: int | None = None
+    num_frames: int | None = None
+    num_inference_steps: int | None = None
+    seed: int | None = None
+    fps: int | None = None
+    timestamp: float | None = None
+    slo_ms: float | None = None
+    extra_body: dict[str, Any] = field(default_factory=dict)
+    image_paths: list[str] | None = None
+    request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+
+
+@dataclass
+class RequestFuncOutput:
+    success: bool = False
+    latency: float = 0.0
+    error: str = ""
+    start_time: float = 0.0
+    response_body: dict[str, Any] = field(default_factory=dict)
+    peak_memory_mb: float = 0.0
+    slo_achieved: bool | None = None
+
+
+def _guess_mime_type(path: str) -> str:
+    mime, _ = mimetypes.guess_type(path)
+    return mime or "application/octet-stream"
+
+
+def _encode_image_as_data_url(path: str) -> str:
+    with open(path, "rb") as f:
+        encoded = base64.b64encode(f.read()).decode("utf-8")
+    mime = _guess_mime_type(path)
+    return f"data:{mime};base64,{encoded}"
+
+
+async def async_request_chat_completions(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput:
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    extra_body = dict(input.extra_body)
+    if input.width and input.height:
+        extra_body.setdefault("height", input.height)
+        extra_body.setdefault("width", input.width)
+    if input.num_frames:
+        extra_body.setdefault("num_frames", input.num_frames)
+    if input.num_inference_steps:
+        extra_body.setdefault("num_inference_steps", input.num_inference_steps)
+    if input.seed is not None:
+        extra_body.setdefault("seed", input.seed)
+    if input.fps:
+        extra_body.setdefault("fps", input.fps)
+
+    if input.image_paths and len(input.image_paths) > 0:
+        content = []
+        if input.prompt:
+            content.append({"type": "text", "text": input.prompt})
+        for img_path in input.image_paths:
+            if not os.path.exists(img_path):
+                output.error = f"Image file not found: {img_path}"
+                output.success = False
+                if pbar:
+                    pbar.update(1)
+                return output
+            content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": _encode_image_as_data_url(img_path)},
+                }
+            )
+        messages = [{"role": "user", "content": content}]
+    else:
+        messages = [{"role": "user", "content": input.prompt}]
+
+    payload = {
+        "model": input.model,
+        "messages": messages,
+    }
+    if extra_body:
+        payload["extra_body"] = extra_body
+
+    try:
+        async with session.post(input.api_url, json=payload) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+                if "peak_memory_mb" in resp_json:
+                    output.peak_memory_mb = resp_json["peak_memory_mb"]
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+async def async_request_openai_images(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput:
+    """
+    Send request to OpenAI's /v1/images/generations endpoint.
+    """
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    # Build size string from width/height
+    width = input.width or 1024
+    height = input.height or 1024
+    size = f"{width}x{height}"
+
+    payload: dict[str, Any] = {
+        "model": input.model,
+        "prompt": input.prompt,
+        "n": 1,
+        "size": size,
+        "response_format": "b64_json",
+    }
+
+    # Add optional parameters
+    if input.seed is not None:
+        payload["seed"] = input.seed
+    if input.num_inference_steps is not None:
+        payload["num_inference_steps"] = input.num_inference_steps
+
+    # Add any extra body parameters
+    if input.extra_body:
+        for key, value in input.extra_body.items():
+            if key not in payload:
+                payload[key] = value
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": "Bearer EMPTY",
+    }
+
+    try:
+        async with session.post(input.api_url, json=payload, headers=headers) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+                # Check for usage/memory info if available
+                if "usage" in resp_json and "peak_memory_mb" in resp_json.get("usage", {}):
+                    output.peak_memory_mb = resp_json["usage"]["peak_memory_mb"]
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+backends_function_mapping = {
+    "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
+    "openai": (async_request_openai_images, "/v1/images/generations"),
+}
@@ -6,80 +6,60 @@
 Benchmark online serving for diffusion models (Image/Video Generation).
 If you want to use i2v, i2i dataset, you should `uv pip install gdown` first

+Supports multiple backends:
+    - vllm-omni: Uses /v1/chat/completions endpoint (default)
+    - openai: Uses /v1/images/generations endpoint
+
 Usage:
-    # Video
+    # Video (vllm-omni backend)
    t2v:
    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --dataset vbench --task t2v --num-prompts 10 \
+        --backend vllm-omni --dataset vbench --task t2v --num-prompts 10 \
        --height 480 --width 640 --fps 16 --num-frames 80

    i2v:
    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --dataset vbench --task i2v --num-prompts 10
+        --backend vllm-omni --dataset vbench --task i2v --num-prompts 10


-    # Image
+    # Image (vllm-omni backend)
    t2i:
    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --dataset vbench --task t2i --num-prompts 10 \
+        --backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
        --height 1024 --width 1024

    i2i:
    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --dataset vbench --task i2i --num-prompts 10
+        --backend vllm-omni --dataset vbench --task i2i --num-prompts 10
+
+    # Image (openai backend)
+    t2i:
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+        --backend openai --dataset vbench --task t2i --num-prompts 10 \
+        --height 1024 --width 1024 --port 3000

 """

 import argparse
 import ast
 import asyncio
-import base64
 import glob
 import json
-import mimetypes
 import os
 import time
 import uuid
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator
-from dataclasses import dataclass, field, replace
+from dataclasses import replace
 from typing import Any

 import aiohttp
 import numpy as np
 import requests
+from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
 from tqdm.asyncio import tqdm


-@dataclass
-class RequestFuncInput:
-    prompt: str
-    api_url: str
-    model: str
-    width: int | None = None
-    height: int | None = None
-    num_frames: int | None = None
-    num_inference_steps: int | None = None
-    seed: int | None = None
-    fps: int | None = None
-    timestamp: float | None = None
-    slo_ms: float | None = None
-    extra_body: dict[str, Any] = field(default_factory=dict)
-    image_paths: list[str] | None = None
-    request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
-
-
-@dataclass
-class RequestFuncOutput:
-    success: bool = False
-    latency: float = 0.0
-    error: str = ""
-    start_time: float = 0.0
-    response_body: dict[str, Any] = field(default_factory=dict)
-    peak_memory_mb: float = 0.0
-    slo_achieved: bool | None = None
-
-
 class BaseDataset(ABC):
    def __init__(self, args, api_url: str, model: str):
        self.args = args
@@ -684,92 +664,6 @@ async def iter_requests(
        yield req


-def _guess_mime_type(path: str) -> str:
-    mime, _ = mimetypes.guess_type(path)
-    return mime or "application/octet-stream"
-
-
-def _encode_image_as_data_url(path: str) -> str:
-    with open(path, "rb") as f:
-        encoded = base64.b64encode(f.read()).decode("utf-8")
-    mime = _guess_mime_type(path)
-    return f"data:{mime};base64,{encoded}"
-
-
-async def async_request_chat_completions(
-    input: RequestFuncInput,
-    session: aiohttp.ClientSession,
-    pbar: tqdm | None = None,
-) -> RequestFuncOutput:
-    output = RequestFuncOutput()
-    output.start_time = time.perf_counter()
-
-    extra_body = dict(input.extra_body)
-    if input.width and input.height:
-        extra_body.setdefault("height", input.height)
-        extra_body.setdefault("width", input.width)
-    if input.num_frames:
-        extra_body.setdefault("num_frames", input.num_frames)
-    if input.num_inference_steps:
-        extra_body.setdefault("num_inference_steps", input.num_inference_steps)
-    if input.seed is not None:
-        extra_body.setdefault("seed", input.seed)
-    if input.fps:
-        extra_body.setdefault("fps", input.fps)
-
-    if input.image_paths and len(input.image_paths) > 0:
-        content = []
-        if input.prompt:
-            content.append({"type": "text", "text": input.prompt})
-        for img_path in input.image_paths:
-            if not os.path.exists(img_path):
-                output.error = f"Image file not found: {img_path}"
-                output.success = False
-                if pbar:
-                    pbar.update(1)
-                return output
-            content.append(
-                {
-                    "type": "image_url",
-                    "image_url": {"url": _encode_image_as_data_url(img_path)},
-                }
-            )
-        messages = [{"role": "user", "content": content}]
-    else:
-        messages = [{"role": "user", "content": input.prompt}]
-
-    payload = {
-        "model": input.model,
-        "messages": messages,
-    }
-    if extra_body:
-        payload["extra_body"] = extra_body
-
-    try:
-        async with session.post(input.api_url, json=payload) as response:
-            if response.status == 200:
-                resp_json = await response.json()
-                output.response_body = resp_json
-                output.success = True
-                if "peak_memory_mb" in resp_json:
-                    output.peak_memory_mb = resp_json["peak_memory_mb"]
-            else:
-                output.error = f"HTTP {response.status}: {await response.text()}"
-                output.success = False
-    except Exception as e:
-        output.error = str(e)
-        output.success = False
-
-    output.latency = time.perf_counter() - output.start_time
-
-    if output.success and input.slo_ms is not None:
-        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
-
-    if pbar:
-        pbar.update(1)
-    return output
-
-
 def calculate_metrics(
    outputs: list[RequestFuncOutput],
    total_duration: float,
@@ -848,9 +742,9 @@ async def benchmark(args):
    if args.base_url is None:
        args.base_url = f"http://{args.host}:{args.port}"

-    # Setup dataset (vLLM-Omni supports diffusion via /v1/chat/completions)
-    api_url = f"{args.base_url}/v1/chat/completions"
-    request_func = async_request_chat_completions
+    # Setup API URL and request function based on backend
+    request_func, api_url = backends_function_mapping[args.backend]
+    api_url = f"{args.base_url}{api_url}"

    if args.dataset == "vbench":
        dataset = VBenchDataset(args, api_url, args.model)
@@ -920,9 +814,16 @@ async def benchmark(args):
    # Calculate metrics
    metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)

+    # Add configuration info to metrics for JSON output
+    metrics["backend"] = args.backend
+    metrics["model"] = args.model
+    metrics["dataset"] = args.dataset
+    metrics["task"] = args.task
+
    print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))

    # Section 1: Configuration
+    print("{:<40} {:<15}".format("Backend:", args.backend))
    print("{:<40} {:<15}".format("Model:", args.model))
    print("{:<40} {:<15}".format("Dataset:", args.dataset))
    print("{:<40} {:<15}".format("Task:", args.task))
@@ -978,6 +879,13 @@ if __name__ == "__main__":
    parser.add_argument("--host", type=str, default="localhost", help="Server host.")
    parser.add_argument("--port", type=int, default=8091, help="Server port.")
    parser.add_argument("--model", type=str, default="default", help="Model name.")
+    parser.add_argument(
+        "--backend",
+        type=str,
+        default="vllm-omni",
+        choices=["vllm-omni", "openai"],
+        help="Backend to target the benchmark to.",
+    )
    parser.add_argument(
        "--dataset",
        type=str,
@@ -0,0 +1,397 @@
+# RDMA Test Configuration Guide
+
+This document explains how to configure the RDMA environment and run tests for `MooncakeTransferEngineConnector`.
+
+## Table of Contents
+
+- [Docker Container Permissions](#docker-container-permissions)
+- [Single-Node Testing](#single-node-testing)
+- [Multi-Node Testing](#multi-node-testing)
+- [Running Tests](#running-tests)
+- [Cross-Node Testing](#cross-node-testing)
+- [Troubleshooting](#troubleshooting)
+
+---
+
+## Docker Container Permissions
+
+RDMA tests require access to InfiniBand/RoCE devices and system topology. Add the following permissions when running `docker run`.
+
+### Option 1: Minimal Permissions (Recommended)
+
+```bash
+docker run -it \
+    --cap-add=SYS_PTRACE \
+    --cap-add=IPC_LOCK \
+    --security-opt seccomp=unconfined \
+    --network=host \
+    --device=/dev/infiniband \
+    -v /sys/class/infiniband:/sys/class/infiniband:ro \
+    your-image:tag
+```
+
+Parameter explanation:
+- `--cap-add=SYS_PTRACE`: Allow reading system topology information
+- `--cap-add=IPC_LOCK`: Allow memory locking (required for RDMA memory registration)
+- `--security-opt seccomp=unconfined`: Disable seccomp restrictions
+- `--network=host`: Use host network (required for RDMA)
+- `--device=/dev/infiniband`: Mount InfiniBand devices
+- `-v /sys/class/infiniband`: Mount IB device info (read-only)
+
+### Option 2: Full Permissions (Quick but not recommended for production)
+
+```bash
+docker run -it \
+    --privileged \
+    --network=host \
+    your-image:tag
+```
+
+`--privileged` grants full host permissions. Suitable for quick testing but not recommended for production.
+
+---
+
+## Single-Node Testing
+
+When running single-node tests (producer and consumer on the same machine), ensure they use the **same RDMA device**.
+
+### Problem Background
+
+InfiniBand devices use LID (Local Identifier) for routing. Different devices have different LIDs and cannot communicate directly. If no device is specified, Mooncake may assign different devices to connectors, causing handshake failures.
+
+Common error:
+```
+[Handshake] Failed to modify QP to RTR, check mtu, gid, peer lid, peer qp num: Invalid argument [22]
+```
+
+### Solution
+
+**Method 1: Set Environment Variable (Recommended)**
+
+```bash
+# List available RDMA devices
+ibstat
+
+# Select a device (e.g., mlx5_0)
+export RDMA_DEVICE_NAME='mlx5_0'
+
+# Run tests
+pytest test_mooncake_transfer_engine_rdma.py -v -s
+```
+
+**Method 2: Use RoCE Devices**
+
+If the system has RoCE devices (using IPv4 routing), the test code will automatically detect and prefer them. RoCE device GIDs start with `00:00:00:00:00:00:00:00:00:00:ff:ff` (IPv4-mapped).
+
+**Method 3: Ensure MTU Consistency**
+
+Make sure both endpoints use the same MTU:
+
+```bash
+# Check device MTU
+ibstatus mlx5_0
+```
+
+---
+
+## Multi-Node Testing
+
+For multi-node tests, producer and consumer run on different machines connected via InfiniBand switch.
+
+### Prerequisites
+
+1. Both machines have Mooncake and RDMA drivers installed
+2. Both machines are in the same InfiniBand subnet
+3. Switch is properly configured
+
+### Configuration
+
+**Machine A (Producer):**
+
+```bash
+# Set RDMA host IP (InfiniBand interface IP)
+export RDMA_TEST_HOST='10.0.0.1'
+
+# Optional: Specify device
+export RDMA_DEVICE_NAME='mlx5_0'
+```
+
+**Machine B (Consumer):**
+
+```bash
+# Set RDMA host IP
+export RDMA_TEST_HOST='10.0.0.2'
+
+# Optional: Specify device
+export RDMA_DEVICE_NAME='mlx5_0'
+```
+
+### Verify Connectivity
+
+```bash
+# Ping IB interface
+ping 10.0.0.2
+
+# Test RDMA connectivity with ibping
+# On Machine B (server)
+ibping -S
+
+# On Machine A (client)
+ibping -G <Machine_B_GID>
+```
+
+---
+
+## Running Tests
+
+### Run All RDMA Tests (Single-Node, fast suite)
+
+Slow tests (large payloads, stress, concurrency integrity) are marked `@pytest.mark.slow`. Use `-m "not slow"` to skip them in quick CI or local fast iteration.
+
+```bash
+cd tests/distributed/omni_connectors
+
+# Fast suite only (excludes slow/stress tests)
+pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
+```
+
+### Run Including Slow Tests
+
+```bash
+# Run ALL tests including slow/stress tests
+pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s
+
+# Run ONLY the slow/stress tests
+pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m slow
+```
+
+### Run Buffer Management Tests
+
+```bash
+# Fast only
+pytest test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
+
+# Including allocator invariant tests (double-free, overlap, merge)
+pytest test_mooncake_transfer_engine_buffer.py -v -s
+```
+
+### Run Specific Test Classes
+
+```bash
+# Basic connector tests
+pytest test_mooncake_transfer_engine_rdma.py::TestBasicConnector -v -s
+
+# End-to-end RDMA transfer tests
+pytest test_mooncake_transfer_engine_rdma.py::TestEndToEnd -v -s
+
+# Lifecycle & resource management tests
+pytest test_mooncake_transfer_engine_rdma.py::TestLifecycle -v -s
+
+# GPU memory pool tests (requires CUDA)
+pytest test_mooncake_transfer_engine_rdma.py::TestGPUPool -v -s
+
+# Stress / correctness tests (slow)
+pytest test_mooncake_transfer_engine_rdma.py::TestStressCorrectness -v -s
+```
+
+### RDMA Environment Diagnostics
+
+For quick diagnostics (device status, Mooncake availability, env vars, etc.),
+see the [Troubleshooting section](../../../docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md#troubleshooting)
+in the connector documentation.
+
+---
+
+## Cross-Node Testing
+
+The `cross_node_mooncake_transfer_engine.py` script enables testing RDMA transfers between two separate physical machines. This script is **not** auto-discovered by `pytest` (it does not start with `test_`) — it must be run manually on each node.
+
+### Prerequisites
+
+1. Both machines have Mooncake installed
+2. Both machines are connected via InfiniBand/RoCE switch
+3. Firewall allows ZMQ ports (default: 15500, 15501)
+4. Same RDMA device name on both nodes (if multiple devices exist)
+
+### Running Cross-Node Tests
+
+**On Machine A (Producer) — start first:**
+
+```bash
+cd benchmarks/distributed/omni_connectors/
+
+# Optional: specify device if multiple exist
+export RDMA_DEVICE_NAME='mlx5_0'
+
+python cross_node_mooncake_transfer_engine.py \
+    --role producer \
+    --local-host <PRODUCER_IP> \
+    --remote-host <CONSUMER_IP> \
+    --tensor-size-mb 100 \
+    --num-transfers 3
+```
+
+**On Machine B (Consumer) — start after producer:**
+
+```bash
+cd benchmarks/distributed/omni_connectors/
+
+export RDMA_DEVICE_NAME='mlx5_0'
+
+python cross_node_mooncake_transfer_engine.py \
+    --role consumer \
+    --local-host <CONSUMER_IP> \
+    --remote-host <PRODUCER_IP> \
+    --tensor-size-mb 100 \
+    --num-transfers 3
+```
+
+### Transfer Modes
+
+| Mode | Description | Example |
+|------|-------------|---------|
+| `copy` | Normal path — tensor copied to RDMA pool (default) | `--mode copy` |
+| `zerocopy` | Zero-copy path — data created directly in RDMA pool | `--mode zerocopy` |
+| `gpu` | GPU transfer — RDMA pool on GPU, uses GPUDirect | `--mode gpu --gpu-id 0` |
+
+### Benchmark Mode
+
+Skip MD5 verification and measure pure RDMA throughput:
+
+```bash
+# Producer
+python cross_node_mooncake_transfer_engine.py \
+    --role producer \
+    --local-host <PRODUCER_IP> \
+    --remote-host <CONSUMER_IP> \
+    --tensor-size-mb 1024 \
+    --num-transfers 20 \
+    --benchmark
+
+# Consumer
+python cross_node_mooncake_transfer_engine.py \
+    --role consumer \
+    --local-host <CONSUMER_IP> \
+    --remote-host <PRODUCER_IP> \
+    --tensor-size-mb 1024 \
+    --num-transfers 20 \
+    --benchmark
+```
+
+### Cross-Node Test Options
+
+| Option | Description | Default |
+|--------|-------------|---------|
+| `--role` | `producer` or `consumer` | Required |
+| `--local-host` | Local RDMA IP address | Required |
+| `--remote-host` | Remote RDMA IP address | Required |
+| `--local-port` | Local ZMQ port for RDMA data | 15500 |
+| `--remote-port` | Remote ZMQ port for RDMA data | 15500 |
+| `--ctrl-port` | Control channel port | 15501 |
+| `--tensor-size-mb` | Tensor size in MB | 100 |
+| `--num-transfers` | Number of transfers | 3 |
+| `--mode` | `copy`, `zerocopy`, or `gpu` | `copy` |
+| `--gpu-id` | GPU ID for GPU mode | 0 |
+| `--benchmark` | Skip MD5, pure performance test | off |
+
+---
+
+## Troubleshooting
+
+### 1. "Failed to modify QP to RTR" Error
+
+**Cause**: QP handshake failed, usually due to device configuration mismatch.
+
+**Solution**:
+```bash
+# Force using the same device
+export RDMA_DEVICE_NAME='mlx5_0'
+```
+
+### 2. "Mooncake TransferEngine is not available"
+
+**Cause**: Mooncake not installed or import failed.
+
+**Solution**:
+```bash
+# Check Mooncake installation
+python -c "from mooncake.engine import TransferEngine; print('OK')"
+
+# Reinstall if needed
+pip install mooncake-transfer-engine
+# Or using uv
+uv pip install mooncake-transfer-engine
+
+```
+
+### 3. "Permission denied" accessing /dev/infiniband
+
+**Cause**: Container lacks IB device access permissions.
+
+**Solution**:
+```bash
+docker run --device=/dev/infiniband --cap-add=IPC_LOCK ...
+```
+
+### 4. Test Timeout
+
+**Cause**: RDMA connection establishment failed or network latency.
+
+**Solution**:
+```bash
+# Check network status
+ibstat
+ibstatus
+```
+
+### 5. GPU Test Failed "CUDA is not available"
+
+**Cause**: CUDA environment not configured or GPU unavailable.
+
+**Solution**:
+```bash
+# Check CUDA
+python -c "import torch; print(torch.cuda.is_available())"
+
+# Docker needs NVIDIA runtime
+docker run --gpus all ...
+```
+
+---
+
+## Environment Variables Reference
+
+| Variable | Description | Example |
+|----------|-------------|---------|
+| `RDMA_DEVICE_NAME` | Specify RDMA device name | `mlx5_0` |
+| `RDMA_TEST_HOST` | Specify test host IP | `10.0.0.1` |
+| `MC_TE_METRIC` | Enable Mooncake metrics | `1` |
+| `MC_IB_PCI_RELAXED_ORDERING` | Enable PCIe relaxed ordering | `1` |
+
+---
+
+## Test Files Overview
+
+| File | Description | Auto-discovered by pytest |
+|------|-------------|--------------------------|
+| `test_mooncake_transfer_engine_rdma.py` | Integration tests for MooncakeTransferEngineConnector (basic, E2E, lifecycle, GPU) | Yes |
+| `test_mooncake_transfer_engine_buffer.py` | Memory pool and buffer management unit tests | Yes |
+| `cross_node_mooncake_transfer_engine.py` | Cross-node (multi-machine) testing script — run manually | No (filename does not start with `test_`) |
+
+### test_mooncake_transfer_engine_rdma.py — Test Classes
+
+| Test Class | Memory Pool | Marker | Description |
+|------------|-------------|--------|-------------|
+| `TestBasicConnector` | CPU | — | Initialization, put tensor/bytes/object, cleanup, pool exhaustion |
+| `TestEndToEnd` | CPU | — | E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload (100MB), mixed types, concurrency |
+| `TestLifecycle` | CPU | — | Close, context manager, double-close safety |
+| `TestGPUPool` | GPU | — | GPU pool init, put CPU/GPU tensor, GPU E2E transfer |
+| `TestStressCorrectness` | CPU | `slow` | Concurrent put+get with MD5 integrity, bidirectional concurrency, edge cases (1-element tensor, empty bytes), 500MB payload, rapid alloc/free cycles |
+
+### test_mooncake_transfer_engine_buffer.py — Test Classes
+
+| Test Class | Marker | Description |
+|------------|--------|-------------|
+| `TestBufferAllocator` | — | Basic alloc/free, alignment, exhaustion/recovery, thread safety |
+| `TestAllocatorInvariants` | `slow` | Double-free safety, overlap corruption detection, adjacent-block merging, fragmentation/defrag |
+| `TestManagedBuffer` | — | Tensor views, context manager |
@@ -0,0 +1,644 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Cross-Node RDMA Test Script (Automated Version)
+
+This script enables testing RDMA transfers between two separate machines.
+Supports three transfer modes:
+  - copy:      Normal path - tensor copied to RDMA pool (default)
+  - zerocopy:  Zero-copy path - data created directly in RDMA pool
+  - gpu:       GPU transfer - RDMA pool on GPU, uses GPUDirect
+
+Usage:
+    # On Machine A (Producer) - start first:
+    python cross_node_mooncake_transfer_engine.py --role producer --local-host hostname_A --remote-host hostname_B
+
+    # On Machine B (Consumer) - start after producer:
+    python cross_node_mooncake_transfer_engine.py --role consumer --local-host hostname_B --remote-host hostname_A
+
+    # Zero-copy mode:
+    python cross_node_mooncake_transfer_engine.py --role producer ... --mode zerocopy
+
+    # GPU mode (requires GPUDirect RDMA support):
+    python cross_node_mooncake_transfer_engine.py --role producer ... --mode gpu --gpu-id 0
+
+    # Benchmark mode (skip random data generation and MD5 verification,
+    # measures pure RDMA throughput):
+    python cross_node_mooncake_transfer_engine.py --role producer ... --benchmark
+
+Environment Variables:
+    RDMA_DEVICE_NAME:              Specify RDMA device (e.g., mlx5_0)
+    MC_IB_PCI_RELAXED_ORDERING:    Set to 1 to enable PCIe relaxed ordering
+                                   for higher RDMA throughput
+"""
+
+import argparse
+import hashlib
+import os
+import sys
+import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+import msgspec
+import torch
+import zmq
+
+# Add parent path for imports
+sys.path.insert(
+    0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
+)
+
+from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
+    ManagedBuffer,
+    MooncakeTransferEngineConnector,
+    TransferEngine,
+)
+
+
+def compute_md5(tensor: torch.Tensor) -> str:
+    """Compute MD5 checksum of a tensor."""
+    if tensor.is_cuda:
+        tensor = tensor.cpu()
+    data = tensor.contiguous().view(torch.uint8).numpy().tobytes()
+    return hashlib.md5(data).hexdigest()
+
+
+# Control channel message types
+class CtrlMsg(msgspec.Struct):
+    """Control channel message."""
+
+    msg_type: str  # "READY", "TRANSFER", "ACK", "DONE", "ERROR"
+    request_id: str = ""
+    md5: str = ""
+    data_size: int = 0
+    error: str = ""
+
+
+@dataclass
+class TransferConfig:
+    """Configuration for cross-node transfer test."""
+
+    local_host: str
+    remote_host: str
+    local_port: int
+    remote_port: int
+    ctrl_port: int
+    num_transfers: int
+    tensor_size_mb: int
+    mode: str  # "copy", "zerocopy", "gpu"
+    gpu_id: int = 0
+    benchmark: bool = False  # Skip MD5 verification for pure performance test
+
+
+@dataclass
+class TransferStats:
+    """Statistics for transfer operations."""
+
+    success_count: int = 0
+    fail_count: int = 0
+    total_bytes: int = 0
+    elapsed_time: float = 0.0
+
+    @property
+    def throughput_mbps(self) -> float:
+        if self.elapsed_time > 0:
+            return (self.total_bytes / (1024 * 1024)) / self.elapsed_time
+        return 0.0
+
+    def print_summary(self, role: str):
+        print(f"\n{'=' * 60}")
+        print(f" {role.upper()} SUMMARY")
+        print(f"  Successful: {self.success_count}/{self.success_count + self.fail_count}")
+        print(f"  Failed:     {self.fail_count}/{self.success_count + self.fail_count}")
+        print(f"  Total:      {self.total_bytes / (1024 * 1024):.2f} MB")
+        print(f"  Time:       {self.elapsed_time:.2f} s")
+        print(f"  Throughput: {self.throughput_mbps:.2f} MB/s")
+        print(f"{'=' * 60}")
+
+
+class CrossNodeTester(ABC):
+    """Abstract base class for cross-node RDMA testing."""
+
+    def __init__(self, config: TransferConfig):
+        self.config = config
+        self.connector: MooncakeTransferEngineConnector | None = None
+        self.zmq_ctx: zmq.Context | None = None
+        self.ctrl_socket: zmq.Socket | None = None
+        self.stats = TransferStats()
+
+    def get_connector_config(self) -> dict:
+        """Get connector configuration based on mode."""
+        pool_size = int(self.config.tensor_size_mb * 1.5) * 1024 * 1024
+        pool_size = max(pool_size, 128 * 1024 * 1024)
+
+        conn_config = {
+            "host": self.config.local_host,
+            "zmq_port": self.config.local_port,
+            "protocol": "rdma",
+            "memory_pool_size": pool_size,
+        }
+
+        # Set device based on mode
+        if self.config.mode == "gpu":
+            conn_config["memory_pool_device"] = f"cuda:{self.config.gpu_id}"
+        else:
+            conn_config["memory_pool_device"] = "cpu"
+
+        # RDMA device name from environment
+        device_name = os.environ.get("RDMA_DEVICE_NAME")
+        if device_name:
+            conn_config["device_name"] = device_name
+            print(f"[CONFIG] Using RDMA device: {device_name}")
+
+        return conn_config
+
+    def initialize(self):
+        """Initialize connector and ZMQ context."""
+        print(f"[{self.role}] Initializing connector...")
+        conn_config = self.get_connector_config()
+        self.connector = MooncakeTransferEngineConnector(conn_config)
+        self.zmq_ctx = zmq.Context()
+        print(f"[{self.role}] Ready at {self.config.local_host}:{self.config.local_port}")
+
+    def cleanup(self):
+        """Cleanup resources."""
+        if self.ctrl_socket:
+            self.ctrl_socket.close()
+        if self.zmq_ctx:
+            self.zmq_ctx.term()
+        if self.connector:
+            self.connector.close()
+        print(f"[{self.role}] Closed.")
+
+    @property
+    @abstractmethod
+    def role(self) -> str:
+        pass
+
+    @abstractmethod
+    def run(self):
+        pass
+
+
+class Producer(CrossNodeTester):
+    """Producer node - sends data to consumer."""
+
+    @property
+    def role(self) -> str:
+        return "PRODUCER"
+
+    def print_header(self):
+        print(f"\n{'=' * 60}")
+        print(f" PRODUCER MODE ({self.config.mode.upper()})")
+        print(f" Local:  {self.config.local_host}:{self.config.local_port}")
+        print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
+        print(f" Control Port: {self.config.ctrl_port}")
+        print(f" Transfer Mode: {self.config.mode}")
+        if self.config.mode == "gpu":
+            print(f" GPU ID: {self.config.gpu_id}")
+        print(f"{'=' * 60}\n")
+
+    def setup_control_channel(self):
+        """Setup ZMQ control channel as server (REP socket)."""
+        self.ctrl_socket = self.zmq_ctx.socket(zmq.REP)
+        self.ctrl_socket.bind(f"tcp://*:{self.config.ctrl_port}")
+        print(f"[PRODUCER] Control channel listening on port {self.config.ctrl_port}")
+        print("[PRODUCER] Waiting for consumer to connect...")
+
+    def wait_for_consumer(self) -> bool:
+        """Wait for consumer READY signal."""
+        msg_data = self.ctrl_socket.recv()
+        msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
+        if msg.msg_type != "READY":
+            print(f"[PRODUCER] Unexpected message: {msg.msg_type}")
+            return False
+        print("[PRODUCER] Consumer connected!")
+        self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
+        return True
+
+    def create_test_data(self, transfer_idx: int) -> tuple[Any, str, int]:
+        """
+        Create test data based on transfer mode.
+
+        Returns:
+            (data, md5, size) tuple
+        """
+        num_elements = (self.config.tensor_size_mb * 1024 * 1024) // 4
+        data_size = num_elements * 4
+
+        if self.config.mode == "zerocopy":
+            # Zero-Copy Path: Allocate directly from connector's pool
+            offset = self.connector.allocator.alloc(data_size)
+            managed_buf = ManagedBuffer(self.connector.allocator, offset, data_size, self.connector.pool)
+
+            if self.config.benchmark:
+                # In benchmark mode, skip random data generation (use uninitialized memory)
+                return managed_buf, "", data_size
+            else:
+                # Fill buffer with random data using tensor view
+                tensor_view = managed_buf.as_tensor(dtype=torch.float32, shape=(num_elements,))
+                random_data = torch.randn(num_elements, dtype=torch.float32)
+                if tensor_view.is_cuda:
+                    tensor_view.copy_(random_data.to(tensor_view.device))
+                else:
+                    tensor_view.copy_(random_data)
+                md5 = compute_md5(tensor_view)
+                return managed_buf, md5, data_size
+
+        elif self.config.mode == "gpu":
+            # GPU Path: Create tensor on GPU
+            device = f"cuda:{self.config.gpu_id}"
+            if self.config.benchmark:
+                # In benchmark mode, use empty tensor (no random generation)
+                gpu_tensor = torch.empty(num_elements, dtype=torch.float32, device=device)
+                return gpu_tensor, "", data_size
+            else:
+                cpu_tensor = torch.randn(num_elements, dtype=torch.float32)
+                md5 = compute_md5(cpu_tensor)
+                gpu_tensor = cpu_tensor.to(device)
+                return gpu_tensor, md5, data_size
+
+        else:
+            # Copy Path (default): Create regular CPU tensor
+            if self.config.benchmark:
+                # In benchmark mode, use empty tensor (no random generation)
+                tensor = torch.empty(num_elements, dtype=torch.float32)
+                return tensor, "", data_size
+            else:
+                tensor = torch.randn(num_elements, dtype=torch.float32)
+                md5 = compute_md5(tensor)
+                return tensor, md5, data_size
+
+    def do_transfer(self, transfer_idx: int) -> bool:
+        """Perform a single transfer."""
+        req_id = f"cross_node_transfer_{transfer_idx}"
+
+        if not self.config.benchmark:
+            print(f"\n[PRODUCER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
+
+        # Create test data
+        t0 = time.time()
+        data, md5, data_size = self.create_test_data(transfer_idx)
+        t_create = time.time() - t0
+
+        if not self.config.benchmark:
+            print(f"  Mode: {self.config.mode}")
+            print(f"  Size: {self.config.tensor_size_mb} MB")
+            if md5:
+                print(f"  MD5:  {md5[:16]}...")
+            print(f"  Create time: {t_create * 1000:.1f} ms")
+
+        # Put data
+        t1 = time.time()
+        success, size, metadata = self.connector.put("producer", "consumer", req_id, data)
+        t_put = time.time() - t1
+
+        if not success:
+            print("  [FAIL] Put failed")
+            return False
+
+        if not self.config.benchmark:
+            print(f"  [OK] Put successful, {size} bytes ({t_put * 1000:.1f} ms)")
+
+        # Wait for consumer to request transfer info
+        msg_data = self.ctrl_socket.recv()
+        msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
+
+        if msg.msg_type != "READY":
+            print(f"  [ERROR] Unexpected message: {msg.msg_type}")
+            return False
+
+        # Send transfer metadata to consumer
+        transfer_msg = CtrlMsg(
+            msg_type="TRANSFER",
+            request_id=req_id,
+            md5=md5,
+            data_size=data_size,
+        )
+        self.ctrl_socket.send(msgspec.msgpack.encode(transfer_msg))
+
+        # Wait for consumer ACK (this includes RDMA transfer time)
+        t2 = time.time()
+        msg_data = self.ctrl_socket.recv()
+        t_rdma = time.time() - t2
+        msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
+
+        success = msg.msg_type == "ACK"
+        if success:
+            if not self.config.benchmark:
+                print(f"  [OK] RDMA transfer complete ({t_rdma * 1000:.1f} ms)")
+            self.stats.success_count += 1
+            self.stats.total_bytes += size
+        else:
+            print(f"  [WARN] Consumer reported error: {msg.error}")
+            self.stats.fail_count += 1
+
+        # Send ACK to allow consumer to continue
+        self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
+
+        # Cleanup buffer
+        self.connector.cleanup(req_id)
+
+        return success
+
+    def run(self):
+        """Run the producer."""
+        self.print_header()
+        self.initialize()
+        self.setup_control_channel()
+
+        try:
+            if not self.wait_for_consumer():
+                return
+
+            if self.config.benchmark:
+                print(
+                    f"[BENCHMARK] Running {self.config.num_transfers} "
+                    f"transfers of {self.config.tensor_size_mb} MB each..."
+                )
+
+            start_time = time.time()
+
+            for i in range(self.config.num_transfers):
+                self.do_transfer(i)
+                if self.config.benchmark and (i + 1) % 10 == 0:
+                    elapsed = time.time() - start_time
+                    current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
+                    print(f"  Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
+
+            self.stats.elapsed_time = time.time() - start_time
+            self.stats.print_summary("PRODUCER")
+
+            # Wait for final consumer message and send DONE
+            self.ctrl_socket.recv()
+            self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="DONE")))
+
+        finally:
+            self.cleanup()
+
+
+class Consumer(CrossNodeTester):
+    """Consumer node - receives data from producer."""
+
+    @property
+    def role(self) -> str:
+        return "CONSUMER"
+
+    def print_header(self):
+        print(f"\n{'=' * 60}")
+        print(f" CONSUMER MODE ({self.config.mode.upper()})")
+        print(f" Local:  {self.config.local_host}:{self.config.local_port}")
+        print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
+        print(f" Control Port: {self.config.ctrl_port}")
+        print(f" Transfer Mode: {self.config.mode}")
+        if self.config.mode == "gpu":
+            print(f" GPU ID: {self.config.gpu_id}")
+        print(f"{'=' * 60}\n")
+
+    def setup_control_channel(self):
+        """Setup ZMQ control channel as client (REQ socket)."""
+        self.ctrl_socket = self.zmq_ctx.socket(zmq.REQ)
+        ctrl_addr = f"tcp://{self.config.remote_host}:{self.config.ctrl_port}"
+        print(f"[CONSUMER] Connecting to producer control channel at {ctrl_addr}...")
+        self.ctrl_socket.connect(ctrl_addr)
+
+    def connect_to_producer(self) -> bool:
+        """Connect to producer and send READY signal."""
+        self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
+        msg_data = self.ctrl_socket.recv()
+        msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
+        if msg.msg_type != "ACK":
+            print(f"[CONSUMER] Unexpected response: {msg.msg_type}")
+            return False
+        print("[CONSUMER] Connected to producer! Starting transfers...")
+        return True
+
+    def do_transfer(self, transfer_idx: int) -> bool:
+        """Perform a single transfer."""
+        if not self.config.benchmark:
+            print(f"\n[CONSUMER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
+
+        # Request next transfer info
+        self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
+        msg_data = self.ctrl_socket.recv()
+        msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
+
+        if msg.msg_type == "DONE":
+            print("[CONSUMER] Producer signaled completion")
+            return False
+
+        if msg.msg_type != "TRANSFER":
+            print(f"[CONSUMER] Unexpected message: {msg.msg_type}")
+            return False
+
+        req_id = msg.request_id
+        expected_md5 = msg.md5
+        data_size = msg.data_size
+        num_elements = data_size // 4
+
+        if not self.config.benchmark:
+            print(f"  Request ID: {req_id}")
+            if expected_md5:
+                print(f"  Expected MD5: {expected_md5[:16]}...")
+            print(f"  Data Size: {data_size / (1024 * 1024):.2f} MB")
+
+        # Build metadata for get
+        metadata = {
+            "request_id": req_id,
+            "source_host": self.config.remote_host,
+            "source_port": self.config.remote_port,
+            "data_size": data_size,
+            "dtype": "float32",
+            "shape": [num_elements],
+            "is_fast_path": True,
+        }
+
+        if not self.config.benchmark:
+            print(f"  [INFO] Requesting from {self.config.remote_host}:{self.config.remote_port}")
+
+        # Get data with timing
+        t0 = time.time()
+        result = self.connector.get("producer", "consumer", req_id, metadata)
+        t_get = time.time() - t0
+
+        response_msg = CtrlMsg(msg_type="ERROR", error="Get failed")
+
+        if result is not None:
+            recv_buffer, recv_size = result
+            if not self.config.benchmark:
+                print(f"  [OK] Get successful, {recv_size} bytes ({t_get * 1000:.1f} ms)")
+
+            if isinstance(recv_buffer, ManagedBuffer):
+                # In benchmark mode, skip MD5 verification
+                if self.config.benchmark or not expected_md5:
+                    response_msg = CtrlMsg(msg_type="ACK")
+                    self.stats.success_count += 1
+                    self.stats.total_bytes += recv_size
+                else:
+                    # Verify data
+                    t1 = time.time()
+                    reconstructed = recv_buffer.as_tensor(dtype=torch.float32, shape=(num_elements,))
+                    recv_md5 = compute_md5(reconstructed)
+                    t_md5 = time.time() - t1
+                    print(f"  MD5: {recv_md5[:16]}... ({t_md5 * 1000:.1f} ms)")
+
+                    if recv_md5 == expected_md5:
+                        print("  [PASS] MD5 checksum verified!")
+                        response_msg = CtrlMsg(msg_type="ACK")
+                        self.stats.success_count += 1
+                        self.stats.total_bytes += recv_size
+                    else:
+                        print("  [FAIL] MD5 mismatch!")
+                        response_msg = CtrlMsg(msg_type="ERROR", error="MD5 mismatch")
+                        self.stats.fail_count += 1
+
+                recv_buffer.release()
+            else:
+                response_msg = CtrlMsg(msg_type="ACK")
+                self.stats.success_count += 1
+                self.stats.total_bytes += recv_size
+        else:
+            print("  [FAIL] Get failed")
+            self.stats.fail_count += 1
+
+        # Send response to producer
+        self.ctrl_socket.send(msgspec.msgpack.encode(response_msg))
+        # Wait for ACK
+        self.ctrl_socket.recv()
+
+        return response_msg.msg_type == "ACK"
+
+    def run(self):
+        """Run the consumer."""
+        self.print_header()
+        self.initialize()
+        self.setup_control_channel()
+
+        try:
+            if not self.connect_to_producer():
+                return
+
+            if self.config.benchmark:
+                print(
+                    f"[BENCHMARK] Running {self.config.num_transfers} "
+                    f"transfers of {self.config.tensor_size_mb} MB each..."
+                )
+
+            start_time = time.time()
+
+            for i in range(self.config.num_transfers):
+                if not self.do_transfer(i):
+                    break
+                if self.config.benchmark and (i + 1) % 10 == 0:
+                    elapsed = time.time() - start_time
+                    current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
+                    print(f"  Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
+
+            self.stats.elapsed_time = time.time() - start_time
+            self.stats.print_summary("CONSUMER")
+
+            # Send final READY and wait for DONE
+            self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
+            self.ctrl_socket.recv()
+
+        finally:
+            self.cleanup()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Cross-Node RDMA Test (Automated)",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Transfer Modes:
+  copy      - Normal path: tensor copied to RDMA pool (default)
+  zerocopy  - Zero-copy path: data created directly in RDMA pool
+  gpu       - GPU transfer: RDMA pool on GPU, uses GPUDirect
+
+Examples:
+  # Copy mode (default):
+  python cross_node_mooncake_transfer_engine.py --role producer \
+  --local-host hostA --remote-host hostB
+
+  # Zero-copy mode:
+  python cross_node_mooncake_transfer_engine.py --role producer \
+  --local-host hostA --remote-host hostB --mode zerocopy
+
+  # GPU mode:
+  python cross_node_mooncake_transfer_engine.py --role producer \
+  --local-host hostA --remote-host hostB --mode gpu --gpu-id 0
+
+  # Benchmark mode (skip MD5, measure pure RDMA performance):
+  python cross_node_mooncake_transfer_engine.py --role producer \
+  --local-host hostA --remote-host hostB --benchmark
+
+  # With specific RDMA device:
+  RDMA_DEVICE_NAME=mlx5_0 python cross_node_mooncake_transfer_engine.py --role producer ...
+        """,
+    )
+
+    parser.add_argument(
+        "--role", required=True, choices=["producer", "consumer"], help="Role: producer (sends) or consumer (receives)"
+    )
+    parser.add_argument("--local-host", required=True, help="Local hostname or IP address")
+    parser.add_argument("--remote-host", required=True, help="Remote hostname or IP address")
+    parser.add_argument("--local-port", type=int, default=15500, help="Local ZMQ port for RDMA data (default: 15500)")
+    parser.add_argument("--remote-port", type=int, default=15500, help="Remote ZMQ port for RDMA data (default: 15500)")
+    parser.add_argument("--ctrl-port", type=int, default=15501, help="Control channel port (default: 15501)")
+    parser.add_argument("--num-transfers", type=int, default=20, help="Number of transfers to perform (default: 3)")
+    parser.add_argument("--tensor-size-mb", type=int, default=100, help="Tensor size in MB (default: 100)")
+    parser.add_argument(
+        "--mode",
+        choices=["copy", "zerocopy", "gpu"],
+        default="copy",
+        help="Transfer mode: copy, zerocopy, or gpu (default: copy)",
+    )
+    parser.add_argument("--gpu-id", type=int, default=0, help="GPU ID for GPU mode (default: 0)")
+    parser.add_argument(
+        "--benchmark", action="store_true", help="Benchmark mode: skip MD5 verification for pure performance test"
+    )
+
+    args = parser.parse_args()
+
+    # Check Mooncake
+    if TransferEngine is None:
+        print("[ERROR] Mooncake TransferEngine is not available.")
+        print("Install with: pip install mooncake")
+        sys.exit(1)
+
+    # Check CUDA for GPU mode
+    if args.mode == "gpu":
+        if not torch.cuda.is_available():
+            print("[ERROR] CUDA is not available but GPU mode was requested.")
+            sys.exit(1)
+        if args.gpu_id >= torch.cuda.device_count():
+            print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.cuda.device_count()} GPUs.")
+            sys.exit(1)
+        print(f"[INFO] Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")
+
+    config = TransferConfig(
+        local_host=args.local_host,
+        remote_host=args.remote_host,
+        local_port=args.local_port,
+        remote_port=args.remote_port,
+        ctrl_port=args.ctrl_port,
+        num_transfers=args.num_transfers,
+        tensor_size_mb=args.tensor_size_mb,
+        mode=args.mode,
+        gpu_id=args.gpu_id,
+        benchmark=args.benchmark,
+    )
+
+    if args.role == "producer":
+        tester = Producer(config)
+    else:
+        tester = Consumer(config)
+
+    tester.run()
+
+
+if __name__ == "__main__":
+    main()
@@ -36,7 +36,7 @@ bash benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
 ```

 What it does:
- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num_prompts 100`.
+- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num-prompts 100`.
 - Outputs to `benchmarks/qwen3-omni/transformers/benchmark_results/`:
  - `perf_stats.json` — aggregated & per-prompt TPS/latency (thinker/talker/code2wav/overall).
  - `results.json` — per-prompt outputs and audio paths.
@@ -53,7 +53,7 @@ bash benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
 ```

 What it does:
- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--enable-stats`.
+- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
 - Uses `benchmarks/build_dataset/top100.txt` and writes to:
  - Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
    - `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
@@ -16,7 +16,7 @@ if [[ ! -f "benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py" ]
 else
    cd benchmarks/qwen3-omni/transformers

-    python qwen3_omni_moe_transformers.py --prompts_file ../../build_dataset/top100.txt --num_prompts 100
+    python qwen3_omni_moe_transformers.py --prompts-file ../../build_dataset/top100.txt --num-prompts 100

    echo "Logs and outputs are saved to $(pwd)/benchmark_results:"
    echo "  - perf_stats.json    Aggregated/per-prompt TPS and latency (thinker/talker/code2wav/overall)"
@@ -227,17 +227,17 @@ def print_stats(stats: dict):
 def main():
    parser = argparse.ArgumentParser(description="Qwen3-Omni Benchmark Script")
    parser.add_argument(
-        "--prompts_file",
+        "--prompts-file",
        type=str,
        default="benchmark/build_dataset/top100.txt",
        help="Path to the prompts file (one prompt per line)",
    )
    parser.add_argument(
-        "--output_dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
+        "--output-dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
    )
-    parser.add_argument("--model_path", type=str, default=MODEL_PATH, help="Path to the model")
+    parser.add_argument("--model-path", type=str, default=MODEL_PATH, help="Path to the model")
    parser.add_argument("--speaker", type=str, default="Ethan", help="Speaker voice for audio output")
-    parser.add_argument("--num_prompts", type=int, default=None, help="Number of prompts to process (default: all)")
+    parser.add_argument("--num-prompts", type=int, default=None, help="Number of prompts to process (default: all)")
    args = parser.parse_args()

    # Load model and processor
@@ -23,7 +23,7 @@ else
    python $end2end_script_path --output-wav $outputs_dir \
                      --query-type text \
                      --txt-prompts $build_dataset_path \
-                      --enable-stats \
+                      --log-stats \
                      --log-dir $log_dir
    echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
    echo "  - omni_llm_pipeline_text                       run dir/base name"
@@ -0,0 +1,100 @@
+# Qwen3-TTS Benchmark
+
+Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
+
+## Prerequisites
+
+```bash
+pip install matplotlib aiohttp soundfile numpy tqdm
+pip install qwen_tts  # for HF baseline
+```
+
+## Quick Start
+
+Run the full benchmark (vllm-omni + HF baseline) with a single command:
+
+```bash
+cd benchmarks/qwen3-tts
+bash run_benchmark.sh
+```
+
+Results (JSON + PNG plots) are saved to `results/`.
+
+### Common options
+
+```bash
+# Only vllm-omni (skip HF baseline)
+bash run_benchmark.sh --async-only
+
+# Only HF baseline
+bash run_benchmark.sh --hf-only
+
+# Use a different model (e.g. 1.7B)
+MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
+
+# Use batch_size=4 config for higher throughput
+STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
+
+# Custom GPU, prompt count, concurrency levels
+GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
+```
+
+## Manual Steps
+
+### 1) Start the vLLM-Omni server
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
+    "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
+    --omni --host 127.0.0.1 --port 8000 \
+    --stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \
+    --trust-remote-code
+```
+
+### 2) Run online serving benchmark
+
+```bash
+python vllm_omni/bench_tts_serve.py \
+    --port 8000 \
+    --num-prompts 50 \
+    --max-concurrency 1 4 10 \
+    --config-name "async_chunk" \
+    --result-dir results/
+```
+
+### 3) Run HuggingFace baseline
+
+```bash
+python transformers/bench_tts_hf.py \
+    --model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
+    --num-prompts 50 \
+    --gpu-device 0 \
+    --result-dir results/
+```
+
+### 4) Generate comparison plots
+
+```bash
+python plot_results.py \
+    --results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
+    --labels "vllm-omni" "hf_transformers" \
+    --output results/comparison.png
+```
+
+## Stage Configs
+
+| Config | Batch Size | Description |
+|--------|:----------:|-------------|
+| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lower latency) |
+| `vllm_omni/configs/qwen3_tts_bs4.yaml` | 4 | Concurrent request processing (higher throughput) |
+
+Both configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
+
+The model is specified via the CLI `--model` flag (or `MODEL` env var in `run_benchmark.sh`), so the same configs work for both the 0.6B and 1.7B model variants.
+
+## Metrics
+
+- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
+- **E2E (End-to-End Latency)**: Total time from request to complete audio response
+- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
+- **Throughput**: Total audio seconds generated per wall-clock second
@@ -0,0 +1,254 @@
+"""Plot Qwen3-TTS benchmark results.
+
+Generates comparison bar charts similar to the async_chunk design doc:
+- TTFP (Time-to-First-Packet) across concurrency levels
+- E2E latency across concurrency levels
+- RTF (Real-Time Factor) across concurrency levels
+
+Usage:
+    # Compare two configs (async_chunk vs no_async_chunk):
+    python plot_results.py \
+        --results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
+        --labels "async_chunk" "no_async_chunk" \
+        --output results/qwen3_tts_benchmark.png
+
+    # Single config:
+    python plot_results.py \
+        --results results/bench_async_chunk_*.json \
+        --labels "async_chunk" \
+        --output results/qwen3_tts_benchmark.png
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def load_results(result_files: list[str]) -> list[list[dict]]:
+    """Load benchmark results from JSON files."""
+    all_results = []
+    for f in result_files:
+        with open(f) as fh:
+            data = json.load(fh)
+        all_results.append(data)
+    return all_results
+
+
+def plot_comparison(
+    all_results: list[list[dict]],
+    labels: list[str],
+    output_path: str,
+    title_prefix: str = "Qwen3-TTS",
+):
+    """Generate comparison bar charts."""
+    n_configs = len(all_results)
+
+    # Collect concurrency levels present in ALL configs (skip missing data)
+    all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
+    concurrencies = sorted(set.union(*all_concurrencies))
+
+    # Build data arrays, using None for missing concurrency levels
+    ttfp_data = {label: [] for label in labels}
+    e2e_data = {label: [] for label in labels}
+    rtf_data = {label: [] for label in labels}
+    throughput_data = {label: [] for label in labels}
+
+    for results, label in zip(all_results, labels):
+        conc_map = {r["concurrency"]: r for r in results}
+        for c in concurrencies:
+            r = conc_map.get(c)
+            ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
+            e2e_data[label].append(r["mean_e2e_ms"] if r else None)
+            rtf_data[label].append(r["mean_rtf"] if r else None)
+            throughput_data[label].append(r["audio_throughput"] if r else None)
+
+    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
+    fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
+
+    x = np.arange(len(concurrencies))
+    width = 0.35 if n_configs == 2 else 0.5
+    if n_configs > 1:
+        offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
+    else:
+        offsets = [0]
+
+    colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
+
+    def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
+        bars = []
+        for i, (label, values) in enumerate(data_dict.items()):
+            # Replace None with 0 for plotting, but track which are missing
+            plot_values = [v if v is not None else 0 for v in values]
+            color = colors[i % len(colors)]
+            bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
+            bars.append(bar)
+            # Add value labels on bars (skip None/missing data)
+            max_val = max((v for v in values if v is not None), default=1)
+            for rect, val in zip(bar, values):
+                if val is not None and val > 0:
+                    ax.text(
+                        rect.get_x() + rect.get_width() / 2,
+                        rect.get_height() + max_val * 0.02,
+                        f"{val:{fmt}}",
+                        ha="center",
+                        va="bottom",
+                        fontsize=9,
+                        fontweight="bold",
+                    )
+        ax.set_xlabel("Concurrency", fontsize=12)
+        ax.set_ylabel(ylabel, fontsize=12)
+        ax.set_title(title, fontsize=13, fontweight="bold")
+        ax.set_xticks(x)
+        ax.set_xticklabels([str(c) for c in concurrencies])
+        ax.legend(fontsize=10)
+        ax.grid(axis="y", alpha=0.3)
+        ax.set_axisbelow(True)
+
+    plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
+    plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
+    plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
+    plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
+
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def plot_single_summary(results: list[dict], label: str, output_path: str):
+    """Generate a single-config summary with percentile breakdown."""
+    concurrencies = [r["concurrency"] for r in results]
+
+    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
+    fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
+
+    # TTFP breakdown
+    ax = axes[0]
+    means = [r["mean_ttfp_ms"] for r in results]
+    medians = [r["median_ttfp_ms"] for r in results]
+    p90s = [r["p90_ttfp_ms"] for r in results]
+    p99s = [r["p99_ttfp_ms"] for r in results]
+    x = np.arange(len(concurrencies))
+    w = 0.2
+    ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
+    ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
+    ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
+    ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_xlabel("Concurrency")
+    ax.set_ylabel("TTFP (ms)")
+    ax.set_title("Time to First Audio Packet")
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", alpha=0.3)
+
+    # E2E breakdown
+    ax = axes[1]
+    means = [r["mean_e2e_ms"] for r in results]
+    medians = [r["median_e2e_ms"] for r in results]
+    p90s = [r["p90_e2e_ms"] for r in results]
+    p99s = [r["p99_e2e_ms"] for r in results]
+    ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
+    ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
+    ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
+    ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_xlabel("Concurrency")
+    ax.set_ylabel("E2E Latency (ms)")
+    ax.set_title("End-to-End Latency")
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", alpha=0.3)
+
+    # RTF
+    ax = axes[2]
+    means = [r["mean_rtf"] for r in results]
+    medians = [r["median_rtf"] for r in results]
+    ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
+    ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_xlabel("Concurrency")
+    ax.set_ylabel("RTF")
+    ax.set_title("Real-Time Factor")
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", alpha=0.3)
+
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
+    """Print a markdown-formatted comparison table."""
+    concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
+
+    print("\n## Benchmark Results\n")
+    header = "| Metric | Concurrency |"
+    sep = "| --- | --- |"
+    for label in labels:
+        header += f" {label} |"
+        sep += " --- |"
+    print(header)
+    print(sep)
+
+    for metric, key, fmt in [
+        ("TTFP (ms)", "mean_ttfp_ms", ".1f"),
+        ("E2E (ms)", "mean_e2e_ms", ".1f"),
+        ("RTF", "mean_rtf", ".3f"),
+        ("Throughput (audio-s/s)", "audio_throughput", ".2f"),
+    ]:
+        for c in concurrencies:
+            row = f"| {metric} | {c} |"
+            for results in all_results:
+                conc_map = {r["concurrency"]: r for r in results}
+                val = conc_map.get(c, {}).get(key, 0)
+                row += f" {val:{fmt}} |"
+            print(row)
+
+    # Improvement calculation (only if 2 configs)
+    if len(all_results) == 2:
+        print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
+        print("| Metric | Concurrency | Improvement |")
+        print("| --- | --- | --- |")
+        for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
+            for c in concurrencies:
+                m0 = {r["concurrency"]: r for r in all_results[0]}
+                m1 = {r["concurrency"]: r for r in all_results[1]}
+                v0 = m0.get(c, {}).get(key, 0)
+                v1 = m1.get(c, {}).get(key, 0)
+                if v1 > 0:
+                    pct = (v1 - v0) / v1 * 100
+                    print(f"| {metric} | {c} | {pct:+.1f}% |")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
+    parser.add_argument(
+        "--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
+    )
+    parser.add_argument(
+        "--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
+    )
+    parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
+    parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
+
+    all_results = load_results(args.results)
+    print_comparison_table(all_results, args.labels)
+
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+
+    if len(all_results) == 1:
+        plot_single_summary(all_results[0], args.labels[0], args.output)
+    else:
+        plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)
@@ -0,0 +1,3 @@
+# Benchmark results are machine-specific - do not commit
+*
+!.gitignore
@@ -0,0 +1,272 @@
+#!/bin/bash
+# Qwen3-TTS Benchmark Runner
+#
+# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
+# Produces JSON results and comparison plots.
+#
+# Usage:
+#   # Full comparison (vllm-omni + HF):
+#   bash run_benchmark.sh
+#
+#   # Only vllm-omni async_chunk config:
+#   bash run_benchmark.sh --async-only
+#
+#   # Only HuggingFace baseline:
+#   bash run_benchmark.sh --hf-only
+#
+#   # vllm-omni only (skip HF):
+#   bash run_benchmark.sh --skip-hf
+#
+#   # Custom settings:
+#   GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
+#
+#   # Use 1.7B model:
+#   MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
+#
+#   # Use batch_size=4 config:
+#   STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
+#
+# Environment variables:
+#   GPU_DEVICE       - GPU index to use (default: 0)
+#   NUM_PROMPTS      - Number of prompts per concurrency level (default: 50)
+#   CONCURRENCY      - Space-separated concurrency levels (default: "1 4 10")
+#   MODEL            - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
+#   PORT             - Server port (default: 8000)
+#   GPU_MEM_TALKER   - gpu_memory_utilization for talker stage (default: 0.3)
+#   GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
+#   STAGE_CONFIG     - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
+
+# Defaults
+GPU_DEVICE="${GPU_DEVICE:-0}"
+NUM_PROMPTS="${NUM_PROMPTS:-50}"
+CONCURRENCY="${CONCURRENCY:-1 4 10}"
+MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
+PORT="${PORT:-8000}"
+GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}"
+GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}"
+NUM_WARMUPS="${NUM_WARMUPS:-3}"
+STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
+RESULT_DIR="${SCRIPT_DIR}/results"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+
+# Parse args
+RUN_ASYNC=true
+RUN_HF=true
+for arg in "$@"; do
+    case "$arg" in
+        --async-only) RUN_HF=false ;;
+        --hf-only) RUN_ASYNC=false ;;
+        --skip-hf) RUN_HF=false ;;
+    esac
+done
+
+mkdir -p "${RESULT_DIR}"
+
+echo "============================================================"
+echo " Qwen3-TTS Benchmark"
+echo "============================================================"
+echo " GPU:          ${GPU_DEVICE}"
+echo " Model:        ${MODEL}"
+echo " Prompts:      ${NUM_PROMPTS}"
+echo " Concurrency:  ${CONCURRENCY}"
+echo " Port:         ${PORT}"
+echo " Stage config: ${STAGE_CONFIG}"
+echo " Results:      ${RESULT_DIR}"
+echo "============================================================"
+
+# Prepare stage config with correct GPU device and memory settings
+prepare_config() {
+    local config_template="$1"
+    local config_name="$2"
+    local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml"
+
+    # Use sed to patch GPU device and memory utilization
+    sed \
+        -e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \
+        -e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \
+        -e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \
+        "${config_template}" > "${output_path}"
+
+    echo "${output_path}"
+}
+
+# Start server and wait for it to be ready
+start_server() {
+    local stage_config="$1"
+    local config_name="$2"
+    local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
+
+    echo ""
+    echo "Starting server with config: ${config_name}"
+    echo "  Stage config: ${stage_config}"
+    echo "  Log file: ${log_file}"
+
+    VLLM_WORKER_MULTIPROC_METHOD=spawn \
+    CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
+    python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
+        --omni \
+        --host 127.0.0.1 \
+        --port "${PORT}" \
+        --stage-configs-path "${stage_config}" \
+        --stage-init-timeout 120 \
+        --trust-remote-code \
+        --disable-log-stats \
+        > "${log_file}" 2>&1 &
+
+    SERVER_PID=$!
+    echo "  Server PID: ${SERVER_PID}"
+
+    # Wait for server to be ready
+    echo "  Waiting for server to be ready..."
+    local max_wait=300
+    local waited=0
+    while [ ${waited} -lt ${max_wait} ]; do
+        if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
+            echo "  Server is ready! (waited ${waited}s)"
+            return 0
+        fi
+        # Check if process is still alive
+        if ! kill -0 ${SERVER_PID} 2>/dev/null; then
+            echo "  ERROR: Server process died. Check log: ${log_file}"
+            tail -20 "${log_file}"
+            return 1
+        fi
+        sleep 2
+        waited=$((waited + 2))
+    done
+
+    echo "  ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
+    kill ${SERVER_PID} 2>/dev/null || true
+    return 1
+}
+
+# Stop the server
+stop_server() {
+    if [ -n "${SERVER_PID:-}" ]; then
+        echo "  Stopping server (PID: ${SERVER_PID})..."
+        kill ${SERVER_PID} 2>/dev/null || true
+        wait ${SERVER_PID} 2>/dev/null || true
+        # Kill any remaining child processes on the port
+        local pids
+        pids=$(lsof -ti:${PORT} 2>/dev/null || true)
+        if [ -n "${pids}" ]; then
+            echo "  Cleaning up remaining processes on port ${PORT}..."
+            echo "${pids}" | xargs kill -9 2>/dev/null || true
+        fi
+        echo "  Server stopped."
+        SERVER_PID=""
+    fi
+}
+
+# Cleanup on exit
+trap 'stop_server' EXIT
+
+# Run benchmark for a given config
+run_bench() {
+    local config_name="$1"
+    local config_template="$2"
+
+    echo ""
+    echo "============================================================"
+    echo " Benchmarking: ${config_name}"
+    echo "============================================================"
+
+    local stage_config
+    stage_config=$(prepare_config "${config_template}" "${config_name}")
+
+    start_server "${stage_config}" "${config_name}"
+
+    # Convert concurrency string to args
+    local conc_args=""
+    for c in ${CONCURRENCY}; do
+        conc_args="${conc_args} ${c}"
+    done
+
+    cd "${PROJECT_ROOT}"
+    python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
+        --host 127.0.0.1 \
+        --port "${PORT}" \
+        --num-prompts "${NUM_PROMPTS}" \
+        --max-concurrency ${conc_args} \
+        --num-warmups "${NUM_WARMUPS}" \
+        --config-name "${config_name}" \
+        --result-dir "${RESULT_DIR}"
+
+    stop_server
+
+    # Allow GPU memory to settle
+    sleep 5
+}
+
+# Run vllm-omni benchmark
+if [ "${RUN_ASYNC}" = true ]; then
+    run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}"
+fi
+
+# Run HuggingFace baseline benchmark
+if [ "${RUN_HF}" = true ]; then
+    echo ""
+    echo "============================================================"
+    echo " Benchmarking: HuggingFace transformers (offline)"
+    echo "============================================================"
+
+    cd "${PROJECT_ROOT}"
+    python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
+        --model "${MODEL}" \
+        --num-prompts "${NUM_PROMPTS}" \
+        --num-warmups "${NUM_WARMUPS}" \
+        --gpu-device "${GPU_DEVICE}" \
+        --config-name "hf_transformers" \
+        --result-dir "${RESULT_DIR}"
+
+    # Allow GPU memory to settle
+    sleep 5
+fi
+
+# Plot results
+echo ""
+echo "============================================================"
+echo " Generating plots..."
+echo "============================================================"
+
+RESULT_FILES=""
+LABELS=""
+
+if [ "${RUN_ASYNC}" = true ]; then
+    ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
+    if [ -n "${ASYNC_FILE}" ]; then
+        RESULT_FILES="${ASYNC_FILE}"
+        LABELS="async_chunk"
+    fi
+fi
+
+if [ "${RUN_HF}" = true ]; then
+    HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
+    if [ -n "${HF_FILE}" ]; then
+        if [ -n "${RESULT_FILES}" ]; then
+            RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
+            LABELS="${LABELS} hf_transformers"
+        else
+            RESULT_FILES="${HF_FILE}"
+            LABELS="hf_transformers"
+        fi
+    fi
+fi
+
+if [ -n "${RESULT_FILES}" ]; then
+    python "${SCRIPT_DIR}/plot_results.py" \
+        --results ${RESULT_FILES} \
+        --labels ${LABELS} \
+        --output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
+fi
+
+echo ""
+echo "============================================================"
+echo " Benchmark complete!"
+echo " Results: ${RESULT_DIR}"
+echo "============================================================"
@@ -0,0 +1,281 @@
+"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
+
+Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
+Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
+
+Usage:
+    python bench_tts_hf.py \
+        --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
+        --num-prompts 50 \
+        --num-warmups 3 \
+        --gpu-device 0 \
+        --result-dir results/
+"""
+
+import argparse
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
+    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
+]
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 1  # always 1 for offline
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    # TTFP stats - not applicable for HF offline, set to E2E for compatibility
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    # E2E stats (ms)
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    # RTF stats
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    p99_rtf: float = 0.0
+    # Audio stats
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0
+    request_throughput: float = 0.0
+    # Per-request details
+    per_request: list = field(default_factory=list)
+
+
+def run_benchmark(args):
+    from qwen_tts import Qwen3TTSModel
+
+    device = f"cuda:{args.gpu_device}"
+    print(f"Loading model: {args.model} on {device}")
+    model = Qwen3TTSModel.from_pretrained(
+        args.model,
+        device_map=device,
+        dtype=torch.bfloat16,
+    )
+    print("Model loaded.")
+
+    # Build prompt list
+    prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
+
+    # Warmup
+    if args.num_warmups > 0:
+        print(f"Warming up with {args.num_warmups} requests...")
+        for i in range(args.num_warmups):
+            p = PROMPTS[i % len(PROMPTS)]
+            wavs, sr = model.generate_custom_voice(
+                text=p,
+                language=args.language,
+                speaker=args.voice,
+            )
+        # Sync GPU
+        torch.cuda.synchronize(device)
+        print("Warmup done.")
+
+    # Benchmark
+    print(f"Running {args.num_prompts} requests sequentially...")
+    e2e_times = []
+    rtfs = []
+    audio_durations = []
+    per_request = []
+    failed = 0
+
+    audio_dir = None
+    if args.save_audio:
+        audio_dir = Path(args.result_dir) / "audio_hf"
+        audio_dir.mkdir(parents=True, exist_ok=True)
+
+    total_start = time.perf_counter()
+
+    for i, prompt in enumerate(prompts):
+        try:
+            torch.cuda.synchronize(device)
+            st = time.perf_counter()
+
+            wavs, sr = model.generate_custom_voice(
+                text=prompt,
+                language=args.language,
+                speaker=args.voice,
+            )
+
+            torch.cuda.synchronize(device)
+            elapsed = time.perf_counter() - st
+
+            # Compute audio duration
+            audio_samples = wavs[0]
+            if isinstance(audio_samples, torch.Tensor):
+                audio_samples = audio_samples.cpu().numpy()
+            audio_dur = len(audio_samples) / sr
+
+            rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
+
+            e2e_times.append(elapsed)
+            rtfs.append(rtf)
+            audio_durations.append(audio_dur)
+            per_request.append(
+                {
+                    "e2e_ms": elapsed * 1000,
+                    "ttfp_ms": elapsed * 1000,  # no streaming, TTFP = E2E
+                    "rtf": rtf,
+                    "audio_duration_s": audio_dur,
+                    "prompt": prompt,
+                }
+            )
+
+            if audio_dir:
+                sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
+
+            if (i + 1) % 10 == 0 or i == 0:
+                print(
+                    f"  [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms  rtf={rtf:.3f}  audio={audio_dur:.2f}s"
+                )
+
+        except Exception as e:
+            print(f"  [{i + 1}/{args.num_prompts}] FAILED: {e}")
+            failed += 1
+
+    total_duration = time.perf_counter() - total_start
+    completed = len(e2e_times)
+
+    # Compute stats
+    result = BenchmarkResult(
+        config_name=args.config_name,
+        concurrency=1,
+        num_prompts=args.num_prompts,
+        completed=completed,
+        failed=failed,
+        duration_s=total_duration,
+    )
+
+    if e2e_times:
+        e2e_ms = [t * 1000 for t in e2e_times]
+
+        result.mean_e2e_ms = float(np.mean(e2e_ms))
+        result.median_e2e_ms = float(np.median(e2e_ms))
+        result.std_e2e_ms = float(np.std(e2e_ms))
+        result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
+        result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
+        result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
+
+        # For HF offline, TTFP = E2E (no streaming)
+        result.mean_ttfp_ms = result.mean_e2e_ms
+        result.median_ttfp_ms = result.median_e2e_ms
+        result.std_ttfp_ms = result.std_e2e_ms
+        result.p90_ttfp_ms = result.p90_e2e_ms
+        result.p95_ttfp_ms = result.p95_e2e_ms
+        result.p99_ttfp_ms = result.p99_e2e_ms
+
+        result.mean_rtf = float(np.mean(rtfs))
+        result.median_rtf = float(np.median(rtfs))
+        result.std_rtf = float(np.std(rtfs))
+        result.p99_rtf = float(np.percentile(rtfs, 99))
+
+        result.mean_audio_duration_s = float(np.mean(audio_durations))
+        result.total_audio_duration_s = float(np.sum(audio_durations))
+        result.audio_throughput = result.total_audio_duration_s / total_duration
+        result.request_throughput = completed / total_duration
+        result.per_request = per_request
+
+    # Print summary in standardized performance template
+    W = 50
+    print("")
+    print(f"{'=' * W}")
+    print(f"{'Serving Benchmark Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Successful requests:':<40}{completed:<10}")
+    print(f"{'Failed requests:':<40}{failed:<10}")
+    print(f"{'Maximum request concurrency:':<40}{1:<10}")
+    print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
+    print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'End-to-end Latency':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
+    print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
+    print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
+    print(f"{'=' * W}")
+    print(f"{'Audio Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
+    print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Time to First Packet':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
+    print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
+    print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Real Time Factor':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
+    print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
+    print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
+    print(f"{'=' * W}")
+    print("")
+
+    # Save results (as a list with single concurrency=1 entry, matching serve format)
+    result_dir = Path(args.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump([asdict(result)], f, indent=2)
+    print(f"Results saved to {result_file}")
+
+    return result
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
+    parser.add_argument(
+        "--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
+    )
+    parser.add_argument("--num-prompts", type=int, default=50)
+    parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--gpu-device", type=int, default=0)
+    parser.add_argument("--voice", type=str, default="Vivian")
+    parser.add_argument("--language", type=str, default="English")
+    parser.add_argument(
+        "--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
+    )
+    parser.add_argument("--result-dir", type=str, default="results")
+    parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_benchmark(args)
@@ -0,0 +1,301 @@
+"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
+
+Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
+concurrency levels for both async_chunk modes. Saves results as JSON.
+
+Usage:
+    # Run against a server already serving with a given config:
+    python bench_async_chunk.py \
+        --host 127.0.0.1 --port 8000 \
+        --config-name async_chunk_on \
+        --num-prompts 50 \
+        --max-concurrency 1 10 \
+        --result-dir results/
+"""
+
+import argparse
+import asyncio
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+from tqdm.asyncio import tqdm
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
+    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
+]
+
+
+@dataclass
+class RequestResult:
+    success: bool = False
+    ttfp: float = 0.0
+    e2e: float = 0.0
+    audio_bytes: int = 0
+    audio_duration: float = 0.0
+    rtf: float = 0.0
+    prompt: str = ""
+    error: str = ""
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 0
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0
+    request_throughput: float = 0.0
+    per_request: list = field(default_factory=list)
+
+
+def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
+    return num_bytes / sample_width / sample_rate
+
+
+async def send_tts_request(
+    session: aiohttp.ClientSession,
+    api_url: str,
+    prompt: str,
+    voice: str = "vivian",
+    language: str = "English",
+    stream: bool = True,
+    pbar: tqdm | None = None,
+) -> RequestResult:
+    payload = {
+        "input": prompt,
+        "voice": voice,
+        "language": language,
+        "stream": stream,
+        "response_format": "pcm",
+    }
+
+    result = RequestResult(prompt=prompt)
+    st = time.perf_counter()
+
+    try:
+        async with session.post(api_url, json=payload) as response:
+            if response.status != 200:
+                result.error = f"HTTP {response.status}: {await response.text()}"
+                return result
+
+            first_chunk = True
+            total_bytes = 0
+
+            async for chunk in response.content.iter_any():
+                if first_chunk and len(chunk) > 0:
+                    result.ttfp = time.perf_counter() - st
+                    first_chunk = False
+                total_bytes += len(chunk)
+
+            result.e2e = time.perf_counter() - st
+            result.audio_bytes = total_bytes
+            result.audio_duration = pcm_bytes_to_duration(total_bytes)
+            if result.audio_duration > 0:
+                result.rtf = result.e2e / result.audio_duration
+            result.success = True
+
+    except Exception as e:
+        result.error = str(e)
+        result.e2e = time.perf_counter() - st
+
+    if pbar:
+        pbar.update(1)
+    return result
+
+
+async def run_benchmark(
+    host: str,
+    port: int,
+    num_prompts: int,
+    max_concurrency: int,
+    num_warmups: int = 3,
+    voice: str = "vivian",
+    language: str = "English",
+    stream: bool = True,
+) -> BenchmarkResult:
+    api_url = f"http://{host}:{port}/v1/audio/speech"
+
+    connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
+    session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
+
+    if num_warmups > 0:
+        print(f"  Warming up with {num_warmups} requests...")
+        warmup_tasks = [
+            send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
+            for i in range(num_warmups)
+        ]
+        await asyncio.gather(*warmup_tasks)
+        print("  Warmup done.")
+
+    request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
+
+    print(f"  Running {num_prompts} requests with concurrency={max_concurrency}...")
+    semaphore = asyncio.Semaphore(max_concurrency)
+    pbar = tqdm(total=num_prompts, desc=f"  concurrency={max_concurrency}")
+
+    async def limited_request(prompt):
+        async with semaphore:
+            return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
+
+    start_time = time.perf_counter()
+    tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
+    results: list[RequestResult] = await asyncio.gather(*tasks)
+    duration = time.perf_counter() - start_time
+    pbar.close()
+
+    await session.close()
+
+    successful = [r for r in results if r.success]
+    failed = [r for r in results if not r.success]
+
+    bench = BenchmarkResult(
+        concurrency=max_concurrency,
+        num_prompts=num_prompts,
+        completed=len(successful),
+        failed=len(failed),
+        duration_s=duration,
+    )
+
+    if successful:
+        ttfps = [r.ttfp * 1000 for r in successful]
+        e2es = [r.e2e * 1000 for r in successful]
+        rtfs = [r.rtf for r in successful]
+        audio_durs = [r.audio_duration for r in successful]
+
+        bench.mean_ttfp_ms = float(np.mean(ttfps))
+        bench.median_ttfp_ms = float(np.median(ttfps))
+        bench.std_ttfp_ms = float(np.std(ttfps))
+        bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
+        bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
+        bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
+
+        bench.mean_e2e_ms = float(np.mean(e2es))
+        bench.median_e2e_ms = float(np.median(e2es))
+        bench.std_e2e_ms = float(np.std(e2es))
+        bench.p90_e2e_ms = float(np.percentile(e2es, 90))
+        bench.p95_e2e_ms = float(np.percentile(e2es, 95))
+        bench.p99_e2e_ms = float(np.percentile(e2es, 99))
+
+        bench.mean_rtf = float(np.mean(rtfs))
+        bench.median_rtf = float(np.median(rtfs))
+        bench.std_rtf = float(np.std(rtfs))
+
+        bench.mean_audio_duration_s = float(np.mean(audio_durs))
+        bench.total_audio_duration_s = float(np.sum(audio_durs))
+        bench.audio_throughput = bench.total_audio_duration_s / duration
+        bench.request_throughput = len(successful) / duration
+
+        bench.per_request = [
+            {
+                "ttfp_ms": r.ttfp * 1000,
+                "e2e_ms": r.e2e * 1000,
+                "rtf": r.rtf,
+                "audio_duration_s": r.audio_duration,
+                "prompt": r.prompt,
+            }
+            for r in successful
+        ]
+
+    print(f"\n{'=' * 60}")
+    print(f"  Concurrency: {max_concurrency}  |  Completed: {bench.completed}  |  Failed: {bench.failed}")
+    print(f"  Duration: {duration:.2f}s  |  Throughput: {bench.request_throughput:.2f} req/s")
+    print(
+        f"  TTFP (ms):  mean={bench.mean_ttfp_ms:.1f}  median={bench.median_ttfp_ms:.1f}"
+        f"  p90={bench.p90_ttfp_ms:.1f}  p99={bench.p99_ttfp_ms:.1f}"
+    )
+    print(
+        f"  E2E (ms):   mean={bench.mean_e2e_ms:.1f}  median={bench.median_e2e_ms:.1f}"
+        f"  p90={bench.p90_e2e_ms:.1f}  p99={bench.p99_e2e_ms:.1f}"
+    )
+    print(f"  RTF:        mean={bench.mean_rtf:.3f}  median={bench.median_rtf:.3f}")
+    print(f"  Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
+    print(f"{'=' * 60}\n")
+
+    if failed:
+        for r in failed[:3]:
+            print(f"  [ERROR] {r.error[:200]}")
+
+    return bench
+
+
+async def main(args):
+    all_results = []
+
+    for concurrency in args.max_concurrency:
+        result = await run_benchmark(
+            host=args.host,
+            port=args.port,
+            num_prompts=args.num_prompts,
+            max_concurrency=concurrency,
+            num_warmups=args.num_warmups,
+            voice=args.voice,
+            language=args.language,
+            stream=args.stream,
+        )
+        result.config_name = args.config_name
+        all_results.append(asdict(result))
+
+    result_dir = Path(args.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump(all_results, f, indent=2)
+    print(f"Results saved to {result_file}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--num-prompts", type=int, default=50)
+    parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
+    parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--voice", type=str, default="vivian")
+    parser.add_argument("--language", type=str, default="English")
+    parser.add_argument("--stream", action="store_true", default=True)
+    parser.add_argument("--no-stream", dest="stream", action="store_false")
+    parser.add_argument("--config-name", type=str, default="async_chunk_on")
+    parser.add_argument("--result-dir", type=str, default="results")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    asyncio.run(main(args))
@@ -0,0 +1,348 @@
+"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
+
+Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
+across configurable concurrency levels. Saves results as JSON for plotting.
+
+Usage:
+    python bench_tts_serve.py \
+        --host 127.0.0.1 --port 8000 \
+        --num-prompts 50 \
+        --max-concurrency 1 4 10 \
+        --result-dir results/
+"""
+
+import argparse
+import asyncio
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+from tqdm.asyncio import tqdm
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
+    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
+]
+
+
+@dataclass
+class RequestResult:
+    success: bool = False
+    ttfp: float = 0.0  # Time to first audio packet (seconds)
+    e2e: float = 0.0  # End-to-end latency (seconds)
+    audio_bytes: int = 0  # Total audio bytes received
+    audio_duration: float = 0.0  # Audio duration in seconds (estimated from PCM)
+    rtf: float = 0.0  # Real-time factor = e2e / audio_duration
+    prompt: str = ""
+    error: str = ""
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 0
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    # TTFP stats (ms)
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    # E2E stats (ms)
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    # RTF stats
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    p99_rtf: float = 0.0
+    # Audio stats
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0  # audio_duration / wall_time
+    request_throughput: float = 0.0  # requests / second
+    # Per-request details
+    per_request: list = field(default_factory=list)
+
+
+def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
+    """Convert raw PCM byte count to duration in seconds."""
+    num_samples = num_bytes / sample_width
+    return num_samples / sample_rate
+
+
+async def send_tts_request(
+    session: aiohttp.ClientSession,
+    api_url: str,
+    prompt: str,
+    voice: str = "vivian",
+    language: str = "English",
+    pbar: tqdm | None = None,
+) -> RequestResult:
+    """Send a streaming TTS request and measure latency metrics."""
+    payload = {
+        "input": prompt,
+        "voice": voice,
+        "language": language,
+        "stream": True,
+        "response_format": "pcm",
+    }
+
+    result = RequestResult(prompt=prompt)
+    st = time.perf_counter()
+
+    try:
+        async with session.post(api_url, json=payload) as response:
+            if response.status != 200:
+                result.error = f"HTTP {response.status}: {await response.text()}"
+                result.success = False
+                return result
+
+            first_chunk = True
+            total_bytes = 0
+
+            async for chunk in response.content.iter_any():
+                if first_chunk and len(chunk) > 0:
+                    result.ttfp = time.perf_counter() - st
+                    first_chunk = False
+                total_bytes += len(chunk)
+
+            result.e2e = time.perf_counter() - st
+            result.audio_bytes = total_bytes
+            result.audio_duration = pcm_bytes_to_duration(total_bytes)
+
+            if result.audio_duration > 0:
+                result.rtf = result.e2e / result.audio_duration
+            result.success = True
+
+    except Exception as e:
+        result.error = str(e)
+        result.success = False
+        result.e2e = time.perf_counter() - st
+
+    if pbar:
+        pbar.update(1)
+    return result
+
+
+async def run_benchmark(
+    host: str,
+    port: int,
+    num_prompts: int,
+    max_concurrency: int,
+    num_warmups: int = 3,
+    voice: str = "vivian",
+    language: str = "English",
+) -> BenchmarkResult:
+    """Run benchmark at a given concurrency level."""
+    api_url = f"http://{host}:{port}/v1/audio/speech"
+
+    connector = aiohttp.TCPConnector(
+        limit=max_concurrency,
+        limit_per_host=max_concurrency,
+        keepalive_timeout=60,
+    )
+    session = aiohttp.ClientSession(
+        connector=connector,
+        timeout=aiohttp.ClientTimeout(total=600),
+    )
+
+    # Warmup
+    if num_warmups > 0:
+        print(f"  Warming up with {num_warmups} requests...")
+        warmup_tasks = []
+        for i in range(num_warmups):
+            prompt = PROMPTS[i % len(PROMPTS)]
+            warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
+        await asyncio.gather(*warmup_tasks)
+        print("  Warmup done.")
+
+    # Build request list
+    request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
+
+    # Run benchmark
+    print(f"  Running {num_prompts} requests with concurrency={max_concurrency}...")
+    semaphore = asyncio.Semaphore(max_concurrency)
+    pbar = tqdm(total=num_prompts, desc=f"  concurrency={max_concurrency}")
+
+    async def limited_request(prompt):
+        async with semaphore:
+            return await send_tts_request(session, api_url, prompt, voice, language, pbar)
+
+    start_time = time.perf_counter()
+    tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
+    results: list[RequestResult] = await asyncio.gather(*tasks)
+    duration = time.perf_counter() - start_time
+    pbar.close()
+
+    await session.close()
+
+    # Compute stats
+    successful = [r for r in results if r.success]
+    failed = [r for r in results if not r.success]
+
+    bench = BenchmarkResult(
+        concurrency=max_concurrency,
+        num_prompts=num_prompts,
+        completed=len(successful),
+        failed=len(failed),
+        duration_s=duration,
+    )
+
+    if successful:
+        ttfps = [r.ttfp * 1000 for r in successful]  # convert to ms
+        e2es = [r.e2e * 1000 for r in successful]
+        rtfs = [r.rtf for r in successful]
+        audio_durs = [r.audio_duration for r in successful]
+
+        bench.mean_ttfp_ms = float(np.mean(ttfps))
+        bench.median_ttfp_ms = float(np.median(ttfps))
+        bench.std_ttfp_ms = float(np.std(ttfps))
+        bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
+        bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
+        bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
+
+        bench.mean_e2e_ms = float(np.mean(e2es))
+        bench.median_e2e_ms = float(np.median(e2es))
+        bench.std_e2e_ms = float(np.std(e2es))
+        bench.p90_e2e_ms = float(np.percentile(e2es, 90))
+        bench.p95_e2e_ms = float(np.percentile(e2es, 95))
+        bench.p99_e2e_ms = float(np.percentile(e2es, 99))
+
+        bench.mean_rtf = float(np.mean(rtfs))
+        bench.median_rtf = float(np.median(rtfs))
+        bench.std_rtf = float(np.std(rtfs))
+        bench.p99_rtf = float(np.percentile(rtfs, 99))
+
+        bench.mean_audio_duration_s = float(np.mean(audio_durs))
+        bench.total_audio_duration_s = float(np.sum(audio_durs))
+        bench.audio_throughput = bench.total_audio_duration_s / duration
+        bench.request_throughput = len(successful) / duration
+
+        bench.per_request = [
+            {
+                "ttfp_ms": r.ttfp * 1000,
+                "e2e_ms": r.e2e * 1000,
+                "rtf": r.rtf,
+                "audio_duration_s": r.audio_duration,
+                "prompt": r.prompt,
+            }
+            for r in successful
+        ]
+
+    # Print summary in standardized performance template
+    W = 50
+    print("")
+    print(f"{'=' * W}")
+    print(f"{'Serving Benchmark Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Successful requests:':<40}{bench.completed:<10}")
+    print(f"{'Failed requests:':<40}{bench.failed:<10}")
+    print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
+    print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
+    print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'End-to-end Latency':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
+    print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
+    print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
+    print(f"{'=' * W}")
+    print(f"{'Audio Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
+    print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Time to First Packet':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
+    print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
+    print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Real Time Factor':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
+    print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
+    print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
+    print(f"{'=' * W}")
+    print("")
+
+    if failed:
+        for r in failed[:3]:
+            print(f"  [ERROR] {r.error[:200]}")
+
+    return bench
+
+
+async def main(args):
+    all_results = []
+
+    for concurrency in args.max_concurrency:
+        result = await run_benchmark(
+            host=args.host,
+            port=args.port,
+            num_prompts=args.num_prompts,
+            max_concurrency=concurrency,
+            num_warmups=args.num_warmups,
+            voice=args.voice,
+            language=args.language,
+        )
+        result.config_name = args.config_name
+        all_results.append(asdict(result))
+
+    # Save results
+    result_dir = Path(args.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump(all_results, f, indent=2)
+    print(f"Results saved to {result_file}")
+
+    return all_results
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
+    parser.add_argument(  # noqa: E501
+        "--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
+    )
+    parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--voice", type=str, default="vivian")
+    parser.add_argument("--language", type=str, default="English")
+    parser.add_argument(
+        "--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
+    )
+    parser.add_argument("--result-dir", type=str, default="results")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    asyncio.run(main(args))
@@ -0,0 +1,96 @@
+# Qwen3-TTS batch_size=1 config (streaming with async_chunk)
+# 2-stage pipeline: Talker -> Code2Wav
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: qwen3_tts
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
+      hf_overrides:
+        architectures: [Qwen3TTSTalkerForConditionalGeneration]
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      enforce_eager: false
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: latent
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 512
+      max_model_len: 4096
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
+      hf_overrides:
+        architectures: [Qwen3TTSCode2Wav]
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: audio
+      gpu_memory_utilization: 0.2
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 8192
+      max_model_len: 32768
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    tts_args:
+      max_instructions_length: 500
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.0
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        codec_streaming: true
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
@@ -0,0 +1,97 @@
+# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
+# Enables concurrent request processing with max_inflight=4
+# 2-stage pipeline: Talker -> Code2Wav
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 4
+    engine_args:
+      model_stage: qwen3_tts
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
+      hf_overrides:
+        architectures: [Qwen3TTSTalkerForConditionalGeneration]
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      enforce_eager: false
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: latent
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 512
+      max_model_len: 4096
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 4
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
+      hf_overrides:
+        architectures: [Qwen3TTSCode2Wav]
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: audio
+      gpu_memory_utilization: 0.2
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 8192
+      max_model_len: 32768
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    tts_args:
+      max_instructions_length: 500
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.0
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 4
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        codec_streaming: true
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
@@ -0,0 +1,249 @@
+"""Plot TTFP comparison: async_chunk off vs on.
+
+Generates a bar chart with improvement arrows, matching the Qwen3-Omni
+async_chunk benchmark figure style.
+
+Usage:
+    python plot_async_chunk.py \
+        --off results/bench_async_chunk_off_*.json \
+        --on results/bench_async_chunk_on_*.json \
+        --output results/qwen3_tts_async_chunk_ttfp.png
+
+    # Also supports E2E and RTF metrics:
+    python plot_async_chunk.py \
+        --off results/bench_async_chunk_off_*.json \
+        --on results/bench_async_chunk_on_*.json \
+        --metric e2e \
+        --output results/qwen3_tts_async_chunk_e2e.png
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+METRIC_CONFIG = {
+    "ttfp": {
+        "key": "mean_ttfp_ms",
+        "ylabel": "TTFP (s)",
+        "title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
+        "to_seconds": True,
+    },
+    "e2e": {
+        "key": "mean_e2e_ms",
+        "ylabel": "E2E (s)",
+        "title": "E2E Latency - Qwen3-TTS, by concurrency",
+        "to_seconds": True,
+    },
+    "rtf": {
+        "key": "mean_rtf",
+        "ylabel": "RTF",
+        "title": "Real-Time Factor - Qwen3-TTS, by concurrency",
+        "to_seconds": False,
+    },
+}
+
+
+def load_results(path: str) -> list[dict]:
+    with open(path) as f:
+        return json.load(f)
+
+
+def plot_ttfp_comparison(
+    off_results: list[dict],
+    on_results: list[dict],
+    metric: str,
+    output_path: str,
+    title_override: str | None = None,
+):
+    cfg = METRIC_CONFIG[metric]
+    key = cfg["key"]
+    to_seconds = cfg["to_seconds"]
+
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    off_vals = []
+    on_vals = []
+    for c in concurrencies:
+        v_off = off_map[c][key]
+        v_on = on_map[c][key]
+        if to_seconds:
+            v_off /= 1000.0
+            v_on /= 1000.0
+        off_vals.append(v_off)
+        on_vals.append(v_on)
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+
+    x = np.arange(len(concurrencies))
+    width = 0.3
+
+    ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
+    ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
+
+    # Draw improvement arrows and labels
+    for i in range(len(concurrencies)):
+        v_off = off_vals[i]
+        v_on = on_vals[i]
+        if v_on > 0:
+            improvement = v_off / v_on
+        else:
+            improvement = float("inf")
+
+        # Arrow from top of off-bar to top of on-bar
+        arrow_start_x = x[i] - width / 2
+        arrow_start_y = v_off * 0.95
+        arrow_end_x = x[i] + width / 2
+        arrow_end_y = v_on * 1.05
+
+        ax.annotate(
+            "",
+            xy=(arrow_end_x, arrow_end_y),
+            xytext=(arrow_start_x, arrow_start_y),
+            arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
+        )
+
+        # Improvement label
+        label_x = (arrow_start_x + arrow_end_x) / 2
+        label_y = arrow_start_y + (v_off - v_on) * 0.15
+        ax.text(
+            label_x,
+            label_y,
+            f"{improvement:.1f}x improvement",
+            ha="center",
+            va="bottom",
+            fontsize=10,
+            color="red",
+            fontweight="bold",
+        )
+
+    title = title_override or cfg["title"]
+    ax.set_title(title, fontsize=13, fontweight="bold")
+    ax.set_ylabel(cfg["ylabel"], fontsize=12)
+    ax.set_xlabel("Max concurrency", fontsize=12)
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_yscale("log")
+    ax.legend(loc="upper left", fontsize=11)
+    ax.grid(axis="y", alpha=0.3, linestyle="--")
+    ax.set_axisbelow(True)
+
+    plt.tight_layout()
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
+    """Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+    fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
+
+    for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
+        cfg = METRIC_CONFIG[metric]
+        key = cfg["key"]
+        to_seconds = cfg["to_seconds"]
+
+        off_vals = []
+        on_vals = []
+        for c in concurrencies:
+            v_off = off_map[c][key]
+            v_on = on_map[c][key]
+            if to_seconds:
+                v_off /= 1000.0
+                v_on /= 1000.0
+            off_vals.append(v_off)
+            on_vals.append(v_on)
+
+        x = np.arange(len(concurrencies))
+        width = 0.3
+        ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
+        ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
+
+        for i in range(len(concurrencies)):
+            if on_vals[i] > 0:
+                improvement = off_vals[i] / on_vals[i]
+                ax.annotate(
+                    "",
+                    xy=(x[i] + width / 2, on_vals[i] * 1.05),
+                    xytext=(x[i] - width / 2, off_vals[i] * 0.95),
+                    arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
+                )
+                label_y = off_vals[i] * 0.85
+                ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
+
+        ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
+        ax.set_ylabel(cfg["ylabel"], fontsize=11)
+        ax.set_xlabel("Max concurrency", fontsize=11)
+        ax.set_xticks(x)
+        ax.set_xticklabels([str(c) for c in concurrencies])
+        if metric != "rtf":
+            ax.set_yscale("log")
+        ax.legend(fontsize=9)
+        ax.grid(axis="y", alpha=0.3, linestyle="--")
+        ax.set_axisbelow(True)
+
+    plt.tight_layout()
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def print_table(off_results: list[dict], on_results: list[dict]):
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    print("\n## Benchmark Results: async_chunk off vs on\n")
+    print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
+    print("| --- | --- | --- | --- | --- |")
+
+    for name, key, fmt in [
+        ("TTFP (ms)", "mean_ttfp_ms", ".1f"),
+        ("E2E (ms)", "mean_e2e_ms", ".1f"),
+        ("RTF", "mean_rtf", ".3f"),
+        ("Throughput", "audio_throughput", ".2f"),
+    ]:
+        for c in concurrencies:
+            v_off = off_map[c].get(key, 0)
+            v_on = on_map[c].get(key, 0)
+            if v_on > 0 and key != "audio_throughput":
+                ratio = f"{v_off / v_on:.1f}x"
+            elif v_off > 0 and key == "audio_throughput":
+                ratio = f"{v_on / v_off:.1f}x"
+            else:
+                ratio = "N/A"
+            print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
+    parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
+    parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
+    parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
+    parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
+    parser.add_argument("--title", type=str, default=None, help="Custom title override")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    off_results = load_results(args.off)
+    on_results = load_results(args.on)
+
+    print_table(off_results, on_results)
+
+    if args.metric == "all":
+        plot_all_metrics(off_results, on_results, args.output)
+    else:
+        plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)
@@ -0,0 +1,167 @@
+#!/bin/bash
+# Qwen3-TTS async_chunk on vs off Benchmark
+#
+# Starts two servers (async_chunk on and off), benchmarks both,
+# and generates comparison plots.
+#
+# Usage:
+#   bash run_async_chunk_benchmark.sh
+#
+# Environment variables:
+#   GPU_DEVICE       - GPU index (default: 0)
+#   MODEL            - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
+#   NUM_PROMPTS      - Prompts per concurrency level (default: 50)
+#   CONCURRENCY      - Space-separated concurrency levels (default: "1 10")
+#   PORT_ON          - Port for async_chunk on server (default: 8000)
+#   PORT_OFF         - Port for async_chunk off server (default: 8001)
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+cd "$PROJECT_ROOT"
+
+GPU_DEVICE="${GPU_DEVICE:-0}"
+MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
+NUM_PROMPTS="${NUM_PROMPTS:-50}"
+CONCURRENCY="${CONCURRENCY:-1 10}"
+NUM_WARMUPS="${NUM_WARMUPS:-3}"
+PORT_ON="${PORT_ON:-8000}"
+PORT_OFF="${PORT_OFF:-8001}"
+RESULT_DIR="${SCRIPT_DIR}/results"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+
+STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
+STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
+
+mkdir -p "${RESULT_DIR}"
+
+echo "============================================================"
+echo " Qwen3-TTS async_chunk Benchmark"
+echo "============================================================"
+echo " GPU:            ${GPU_DEVICE}"
+echo " Model:          ${MODEL}"
+echo " Prompts:        ${NUM_PROMPTS}"
+echo " Concurrency:    ${CONCURRENCY}"
+echo " Port (on/off):  ${PORT_ON} / ${PORT_OFF}"
+echo " Results:        ${RESULT_DIR}"
+echo "============================================================"
+
+cleanup() {
+    echo "Cleaning up servers..."
+    kill "$PID_ON" 2>/dev/null || true
+    kill "$PID_OFF" 2>/dev/null || true
+    wait "$PID_ON" 2>/dev/null || true
+    wait "$PID_OFF" 2>/dev/null || true
+}
+trap cleanup EXIT
+
+wait_for_server() {
+    local port=$1
+    local name=$2
+    local max_wait=300
+    local elapsed=0
+    echo "Waiting for ${name} server on port ${port}..."
+    while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
+        sleep 5
+        elapsed=$((elapsed + 5))
+        if [ $elapsed -ge $max_wait ]; then
+            echo "ERROR: ${name} server failed to start within ${max_wait}s"
+            exit 1
+        fi
+    done
+    echo "${name} server ready (${elapsed}s)"
+}
+
+# ---- Phase 1: Start async_chunk ON server ----
+echo ""
+echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
+CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
+    --stage-configs-path "${STAGE_CONFIG_ON}" \
+    --host 0.0.0.0 --port "${PORT_ON}" \
+    --trust-remote-code --enforce-eager --omni \
+    > "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
+PID_ON=$!
+
+wait_for_server "${PORT_ON}" "async_chunk_on"
+
+echo "[Phase 1] Benchmarking async_chunk ON..."
+# shellcheck disable=SC2086
+python "${SCRIPT_DIR}/bench_async_chunk.py" \
+    --host 127.0.0.1 --port "${PORT_ON}" \
+    --config-name "async_chunk_on" \
+    --num-prompts "${NUM_PROMPTS}" \
+    --max-concurrency ${CONCURRENCY} \
+    --num-warmups "${NUM_WARMUPS}" \
+    --result-dir "${RESULT_DIR}"
+
+echo "[Phase 1] Stopping async_chunk ON server..."
+kill "$PID_ON" 2>/dev/null || true
+wait "$PID_ON" 2>/dev/null || true
+sleep 5
+
+# ---- Phase 2: Start async_chunk OFF server ----
+echo ""
+echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
+CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
+    --stage-configs-path "${STAGE_CONFIG_OFF}" \
+    --host 0.0.0.0 --port "${PORT_OFF}" \
+    --trust-remote-code --enforce-eager --omni \
+    > "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
+PID_OFF=$!
+
+wait_for_server "${PORT_OFF}" "async_chunk_off"
+
+echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
+# shellcheck disable=SC2086
+python "${SCRIPT_DIR}/bench_async_chunk.py" \
+    --host 127.0.0.1 --port "${PORT_OFF}" \
+    --config-name "async_chunk_off" \
+    --num-prompts "${NUM_PROMPTS}" \
+    --max-concurrency ${CONCURRENCY} \
+    --num-warmups "${NUM_WARMUPS}" \
+    --no-stream \
+    --result-dir "${RESULT_DIR}"
+
+echo "[Phase 2] Stopping async_chunk OFF server..."
+kill "$PID_OFF" 2>/dev/null || true
+wait "$PID_OFF" 2>/dev/null || true
+
+# ---- Phase 3: Plot results ----
+echo ""
+echo "[Phase 3] Generating plots..."
+
+# Find the latest result files
+RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
+RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
+
+if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
+    echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
+    exit 1
+fi
+
+echo "  ON results:  ${RESULT_ON}"
+echo "  OFF results: ${RESULT_OFF}"
+
+# TTFP comparison (main figure)
+python "${SCRIPT_DIR}/plot_async_chunk.py" \
+    --off "${RESULT_OFF}" \
+    --on "${RESULT_ON}" \
+    --metric ttfp \
+    --output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
+
+# All metrics comparison
+python "${SCRIPT_DIR}/plot_async_chunk.py" \
+    --off "${RESULT_OFF}" \
+    --on "${RESULT_ON}" \
+    --metric all \
+    --output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
+
+echo ""
+echo "============================================================"
+echo " Benchmark complete!"
+echo " Results: ${RESULT_DIR}/"
+echo " Plots:"
+echo "   - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
+echo "   - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
+echo "============================================================"
@@ -1,19 +1,19 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.15.0
+ARG VLLM_BASE_TAG=v0.17.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
-
 COPY . .

 # Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg && \
+    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

 # Install vllm-omni into the same uv-managed Python environment used by the base image.
-RUN uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]"
+# Use bash -c so that $(python3 -c ...) is expanded inside the container.
+RUN uv pip install --system --no-cache-dir ".[dev]"

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

@@ -2,13 +2,16 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
 ARG VLLM_ASCEND_TAG=v0.14.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}

+WORKDIR /vllm-workspace/vllm-ascend
+RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+RUN pip install -v -e .
+
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}

 COPY . .

-# Install vllm-omni with dev dependencies
-RUN pip install --no-cache-dir -e . --no-build-isolation
+RUN pip install -v -e .

 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

@@ -2,13 +2,16 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
 ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}

+WORKDIR /vllm-workspace/vllm-ascend
+RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+RUN pip install -v -e .
+
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}

 COPY . .

-# Install vllm-omni with dev dependencies
-RUN pip install --no-cache-dir -e . --no-build-isolation
+RUN pip install -v -e .

 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

@@ -1,28 +1,53 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.15.0
-FROM ${BASE_IMAGE} AS final
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0
+FROM ${BASE_IMAGE} AS base
+
+# Declare a variable to know if we want to use the nightly build or the stable build.
+# NOTE: REMINDER to vLLM-Omni rebase maintainer
+# Remember to set `USE_NIGHTLY_BUILD` to 0, when switching back to stable vLLM docker image
+# 1. If vLLM-omni maintainer is forced to use custom commits
+#    during rebasing, they can change the two variables
+# 2. Whenever vLLM upstream has released a stable version,
+#    we should swap over to use stable release ASAP.
+#    We should avoid relying on custom commits.
+ARG USE_NIGHTLY_BUILD=0
+ARG VLLM_VERSION_OR_COMMIT_HASH=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
+ARG ARG_PYTORCH_ROCM_ARCH
+ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}

 ARG COMMON_WORKDIR=/app
-
 WORKDIR ${COMMON_WORKDIR}

 # Step 1: Setup - Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg && \
+    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

-RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
+# Step 2: Conditionally reinstall vllm from source for nightly builds
+RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then \
+        python3 -m pip uninstall -y vllm && rm -rf vllm && \
+        git clone https://github.com/vllm-project/vllm.git && \
+        cd vllm && \
+        git checkout ${VLLM_VERSION_OR_COMMIT_HASH} && \
+        python3 -m pip install -r requirements/rocm.txt && \
+        python3 setup.py clean --all && \
+        PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} python3 setup.py develop && \
+        cd ../ && \
+        rm -rf vllm/.git; \
+    fi

-# Step 2: Copy vllm-omni code and install without uv
+# Step 3: Copy vllm-omni code and install without uv
+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]" --no-build-isolation

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

-CMD ["/bin/bash"]
+FROM base AS test

+CMD ["/bin/bash"]
 ENTRYPOINT []

-#Set entrypoint for vllm-openai official images
-FROM final AS vllm-openai
+# Set entrypoint for vllm-openai official images
+FROM base AS vllm-openai
 ENTRYPOINT ["vllm", "serve", "--omni"]
@@ -0,0 +1,152 @@
+# Argument to configure vllm base image if pre-built
+ARG VLLM_BASE=vllm-base
+
+FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
+
+WORKDIR /workspace/
+
+ARG PYTHON_VERSION=3.12
+ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
+
+RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
+    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
+
+RUN apt clean && apt-get update -y && \
+    apt-get install -y --no-install-recommends --fix-missing \
+    curl \
+    espeak-ng \
+    ffmpeg \
+    git \
+    libsndfile1 \
+    libsm6 \
+    libxext6 \
+    libgl1 \
+    lsb-release \
+    libaio-dev \
+    numactl \
+    wget \
+    vim \
+    python3.12 \
+    python3.12-dev \
+    python3-pip
+
+RUN apt update && apt upgrade -y && \
+    apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
+
+# Install UMD with fixed version
+RUN mkdir neo && \
+    cd neo && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
+    dpkg -i *.deb && \
+    cd .. && \
+    rm -rf neo
+
+ENV PATH="/root/.local/bin:$PATH"
+ENV VIRTUAL_ENV="/opt/venv"
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh
+RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+
+# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
+ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
+RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
+    bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
+    rm "${ONECCL_INSTALLER}" && \
+    echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
+    echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc
+RUN rm -f /opt/intel/oneapi/ccl/latest && \
+    ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest
+
+SHELL ["bash", "-c"]
+CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
+
+WORKDIR /workspace/
+ENV UV_HTTP_TIMEOUT=500
+
+# Configure package index for XPU
+ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+ENV UV_INDEX_STRATEGY="unsafe-best-match"
+ENV UV_LINK_MODE="copy"
+
+ARG VLLM_VERSION=v0.17.0
+RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
+WORKDIR /workspace/vllm
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --upgrade pip && \
+    uv pip install -r requirements/xpu.txt
+
+ # used for suffix method speculative decoding
+ # build deps for proto + nanobind-based extensions to set up the build environment
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install grpcio-tools protobuf nanobind
+ # arctic-inference is built from source which needs torch-xpu properly installed first
+RUN --mount=type=cache,target=/root/.cache/uv \
+    source /opt/intel/oneapi/setvars.sh --force && \
+    source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
+    export CMAKE_PREFIX_PATH="$(python -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
+    uv pip install --no-build-isolation arctic-inference==0.1.1
+
+ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
+
+ENV VLLM_TARGET_DEVICE=xpu
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --no-build-isolation .
+
+CMD ["/bin/bash"]
+
+FROM vllm-base AS vllm-openai
+
+# install additional dependencies for openai api server
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope
+
+# install development dependencies (for testing)
+RUN uv pip install -e tests/vllm_test_utils
+
+# install nixl from source code
+ENV NIXL_VERSION=0.7.0
+RUN python /workspace/vllm/tools/install_nixl_from_source_ubuntu.py
+
+# ensure vllm is properly installed
+RUN python -c "import vllm, inspect; print(vllm.__file__)"
+RUN uv pip show vllm
+
+CMD ["/bin/bash"]
+
+ENTRYPOINT []
+
+FROM ${VLLM_BASE} AS vllm-omni
+
+WORKDIR /workspace/vllm-omni
+COPY . .
+
+ENV VLLM_OMNI_TARGET_DEVICE=xpu
+RUN uv pip install --no-cache-dir ".[dev]" --no-build-isolation
+
+# FIX triton
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip uninstall triton triton-xpu && \
+    uv pip install triton-xpu==3.6.0
+
+# remove torch bundled oneccl to avoid conflicts
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip uninstall oneccl oneccl-devel
+
+FROM vllm-omni AS vllm-omni-openai
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT ["vllm", "serve", "--omni"]
@@ -8,27 +8,40 @@ nav:
    - OpenAI-Compatible API:
      - Image Generation: serving/image_generation_api.md
      - Image Edit: serving/image_edit_api.md
+      - Text to Speech: serving/speech_api.md
  - Examples:
    - examples/README.md
    - Offline Inference:
      - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
+      - CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
+      - Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
+      - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
+      - Helios Video Generation: user_guide/examples/offline_inference/helios.md
+      - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
      - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
      - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
-      - LoRA Inference Examples: user_guide/examples/offline_inference/lora_inference.md
+      - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
+      - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
      - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
-      - Qwen3-TTS Offline Inference: user_guide/examples/offline_inference/qwen3_tts.md
+      - Qwen3-TTS: user_guide/examples/offline_inference/qwen3_tts.md
      - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
      - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
      - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
+      - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
    - Online Serving:
      - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
+      - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
+      - Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
+      - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
      - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
-      - Online LoRA Inference (Diffusion): user_guide/examples/online_serving/lora_inference.md
+      - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
+      - Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
      - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
      - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
-      - Qwen3-TTS Online Serving: user_guide/examples/online_serving/qwen3_tts.md
+      - Qwen3-TTS: user_guide/examples/online_serving/qwen3_tts.md
      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
+      - Text-To-Video: user_guide/examples/online_serving/text_to_video.md
  - General:
    - usage/*
  - Configuration:
@@ -39,20 +52,29 @@ nav:
  - Features:
    - Sleep Mode: features/sleep_mode.md
    - Diffusion Features:
-      - Overview: user_guide/diffusion_acceleration.md
+      - Acceleration Overview: user_guide/diffusion_acceleration.md
      - TeaCache: user_guide/diffusion/teacache.md
      - Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
+      - Quantization:
+        - Overview: user_guide/diffusion/quantization/overview.md
+        - FP8: user_guide/diffusion/quantization/fp8.md
+        - GGUF: user_guide/diffusion/quantization/gguf.md
      - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
      - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
+      - LoRA: user_guide/diffusion/lora.md
+      - Hybrid Sharded Data Parallel: design/feature/hsdp.md
+      - Custom Pipeline: features/custom_pipeline.md
    - ComfyUI: features/comfyui.md
 - Developer Guide:
  - General:
    - contributing/README.md
+    - pr_reviewer.md
    - glob: contributing/*
      flatten_single_child_sections: true
  - Model Implementation:
    - contributing/model/README.md
    - contributing/model/adding_omni_model.md
+    - contributing/model/adding_tts_model.md
    - contributing/model/adding_diffusion_model.md
  - CI: contributing/ci
  - Design Documents:
@@ -62,6 +84,12 @@ nav:
      - design/feature/disaggregated_inference.md
      - design/feature/ray_based_execution.md
      - design/feature/omni_connectors/
+      - design/feature/cfg_parallel.md
+      - design/feature/sequence_parallel.md
+      - design/feature/tensor_parallel.md
+      - design/feature/cache_dit.md
+      - design/feature/teacache.md
+      - design/feature/async_chunk_design.md
    - Module Design:
      - design/module/ar_module.md
      - design/module/dit_module.md
@@ -72,6 +100,7 @@ nav:
  - api/vllm_omni
 - CLI Reference: cli
 - Community:
+  - Governance: community/governance.md
  - community/*
  - Slack: https://slack.vllm.ai
  - Blog: https://blog.vllm.ai
@@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in

 For more information, checkout the following:

- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
+- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
 - [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
@@ -7,23 +7,18 @@ Main entry points for vLLM-Omni inference and serving.
 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
 - [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
 - [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
- [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
- [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
 - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
 - [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
- [vllm_omni.entrypoints.log_utils.OrchestratorMetrics][]
- [vllm_omni.entrypoints.log_utils.StageRequestMetrics][]
- [vllm_omni.entrypoints.log_utils.StageStats][]
 - [vllm_omni.entrypoints.omni.Omni][]
 - [vllm_omni.entrypoints.omni.OmniBase][]
 - [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
 - [vllm_omni.entrypoints.omni_llm.OmniLLM][]
 - [vllm_omni.entrypoints.omni_stage.OmniStage][]
 - [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
+- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]

 ## Inputs

@@ -34,7 +29,6 @@ Input data structures for multi-modal inputs.
 - [vllm_omni.inputs.data.OmniTextPrompt][]
 - [vllm_omni.inputs.data.OmniTokenInputs][]
 - [vllm_omni.inputs.data.OmniTokensPrompt][]
- [vllm_omni.inputs.parse.parse_singleton_prompt_omni][]
 - [vllm_omni.inputs.preprocess.OmniInputPreprocessor][]

 ## Engine
@@ -42,6 +36,12 @@ Input data structures for multi-modal inputs.
 Engine classes for offline and online inference.

 - [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.BufferAllocator][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.ManagedBuffer][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.MooncakeAgentMetadata][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.MooncakeTransferEngineConnector][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.QueryRequest][]
+- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.QueryResponse][]
 - [vllm_omni.engine.AdditionalInformationEntry][]
 - [vllm_omni.engine.AdditionalInformationPayload][]
 - [vllm_omni.engine.OmniEngineCoreOutput][]
@@ -78,6 +78,8 @@ Configuration classes.
 - [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
 - [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
 - [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
+- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioConfig][]
+- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioTokenizerConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSSpeakerEncoderConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerCodePredictorConfig][]
@@ -95,8 +97,10 @@ Configuration classes.
 Worker classes and model runners for distributed inference.

 - [vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner][]
+- [vllm_omni.diffusion.worker.diffusion_worker.CustomPipelineWorkerExtension][]
 - [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
 - [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
+- [vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase][]
 - [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
 - [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
 - [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]
@@ -107,10 +111,12 @@ Worker classes and model runners for distributed inference.
 - [vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker][]
 - [vllm_omni.platforms.xpu.worker.xpu_generation_model_runner.XPUGenerationModelRunner][]
 - [vllm_omni.platforms.xpu.worker.xpu_generation_worker.XPUGenerationWorker][]
+- [vllm_omni.worker.base.OmniGPUWorkerBase][]
 - [vllm_omni.worker.gpu_ar_model_runner.ExecuteModelState][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
 - [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
 - [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
+- [vllm_omni.worker.gpu_memory_utils.parse_cuda_visible_devices][]
 - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
 - [vllm_omni.worker.mixins.OmniWorkerMixin][]
@@ -0,0 +1,9 @@
+# vllm-omni serve
+
+## JSON CLI Arguments
+
+--8<-- "docs/cli/json_tip.inc.md"
+
+## Arguments
+
+--8<-- "docs/generated/argparse_omni/omni_serve.inc.md"
@@ -0,0 +1,52 @@
+# Governance
+
+vLLM-Omni's governance is inspired by the [vLLM governance process](https://docs.vllm.ai/en/latest/governance/process/). We share the same commitment to open source community and meritocratic norms.
+
+## Values
+
+vLLM-Omni aims to be the fastest and easiest-to-use omni-modality inference and serving engine. Our values are aligned with [vLLM's values](https://docs.vllm.ai/en/latest/governance/process/#values):
+
+### Design Values
+
+1. **Top performance**: System performance is our top priority. We monitor overheads, optimize kernels, and publish benchmarks. We never leave performance on the table.
+2. **Ease of use**: vLLM-Omni must be simple to install, configure, and operate. We provide clear documentation, fast startup, clean logs, helpful error messages, and monitoring guides. Many users fork our code or study it deeply, so we keep it readable and modular.
+3. **Wide coverage**: vLLM-Omni supports frontier models and high-performance accelerators. We make it easy to add new models and hardware. vLLM-Omni + PyTorch form a simple interface that avoids complexity.
+4. **Production ready**: vLLM-Omni runs 24/7 in production. It must be easy to operate and monitor for health issues.
+5. **Extensibility**: vLLM-Omni serves as fundamental omni-modality infrastructure. Our codebase cannot cover every use case, so we design for easy forking and customization.
+
+### Collaboration Values
+
+1. **Tightly Knit and Fast-Moving**: Our maintainer team is aligned on vision, philosophy, and roadmap. We work closely to unblock each other and move quickly.
+2. **Individual Merit**: No one buys their way into governance. Committer status belongs to individuals, not companies. We reward contribution, maintenance, and project stewardship.
+
+## Project Maintainers
+
+### Lead Maintainers
+
+Lead maintainers are responsible for the overall direction and strategy of the project:
+
+- [@Gaohan123](https://github.com/Gaohan123)
+- [@hsliuustc0106](https://github.com/hsliuustc0106)
+- [@ywang96](https://github.com/ywang96)
+
+### Active Committers
+
+Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
+
+- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
+- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
+- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
+- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
+- [@princepride](https://github.com/princepride): Diffusion and Omni Support
+- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
+- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
+- [@wtomin](https://github.com/wtomin):
+- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
+
+## Meetings
+
+Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
+
+## Committer Nomination Process
+
+Every month, any active committer can nominate new committer(s) to the project. Up to **two new committers** will be admitted per month based on the quality and impact of their contributions.
@@ -2,7 +2,7 @@

 This section lists the most common options for running vLLM-Omni.

-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)

 Currently, the main options are maintained by stage configs for each model.

@@ -19,3 +19,4 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 - **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
 - **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
 - **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
+- **[CPU Offloading](../user_guide/diffusion/cpu_offload_diffusion.md)** - Enable CPU offloading (model-level and layerwise) for for DiT models
@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
 For offline (Assume necessary dependencies have ben imported):
 ```python
 model_name = "Qwen/Qwen2.5-Omni-7B"
-omni_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
+omni_llm = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
 ```

 For online serving:
@@ -135,6 +135,14 @@ Each stage in the `stage_args` list contains the following configuration options

 A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`).

+### `prompt_expand_func` (Optional)
+
+A custom Python function hook for the LLM stage (Stage 0) that expands a single incoming prompt object into multiple prompts. This is primarily used for multi-modal Classifier-Free Guidance (CFG), where it generates the necessary companion requests (like a negative text prompt) and tags them with internal roles (e.g., `cfg_text`). This ensures the upstream LLM generates the needed contextual hidden states for both the conditional and unconditional generations simultaneously.
+
+### `cfg_kv_collect_func` (Optional)
+
+A custom Python function hook for downstream diffusion stages (Stage 1+) to collect, map, and process the KV caches transferred from the companion requests fired by `prompt_expand_func`. It aggregates the hidden condition states cleanly (e.g., binding them as `cfg_text_past_key_values` and `cfg_text_kv_metadata`), allowing the diffusion runtime to perform CFG smoothly without redundantly evaluating text paths on the DiT workers.
+
 ### `runtime`

 Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed.
@@ -273,3 +281,15 @@ Default: `True`
 Penalty applied to tokens that have already appeared in the generated sequence. Values greater than 1.0 discourage repetition, while values less than 1.0 encourage it. A value of 1.0 applies no penalty.

 Default: `1.1`
+
+### `tts_args` (TTS stages only)
+
+Configuration for Text-to-Speech specific parameters. This section is only applicable to TTS model stages (e.g., `qwen3_tts`).
+
+#### `tts_args.max_instructions_length`
+
+Maximum character length for voice style/emotion instructions. Instructions exceeding this limit will be rejected with a validation error.
+
+Default: `500`
+
+This value can be overridden at runtime using the `--tts-max-instructions-length` CLI parameter when starting the server.
@@ -71,14 +71,7 @@ For additional features and advanced configurations, refer to the:
 ### Testing

 vLLM-Omni uses `pytest` to test the codebase.
-
-```bash
-# Run all tests
-pytest tests/
-
-# Run tests for a single test file with detailed output
-pytest -s -v tests/test_omni_llm.py
-```
+Please refer to the [test instructions](./ci/test_guide.md) for detailed testing information.

 !!! warning
    Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
@@ -123,6 +116,9 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 !!! note
    If the PR spans more than one category, please include all relevant prefixes.

+### Local Test
+Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
+
 ### Code Quality

 The PR needs to meet the following code quality standards:
@@ -0,0 +1,680 @@
+# Multi-Level Automated Testing System Documentation
+
+## Document Overview
+
+This testing system aims to build a complete, efficient, and well-structured quality assurance framework for the development, integration, and release of model services. It draws on the concept of the test pyramid from modern software engineering, progressively expanding testing activities from basic code logic verification to complex end-to-end (E2E) functionality, performance, accuracy, and even long-term stability validation.
+
+Through five levels (L1-L5) and common (Common) specifications, the system clarifies the testing objectives, scope, execution frequency, and required resources for different development stages (e.g., each commit, PR merge, daily build, pre-release). This ensures that models meet high standards for functionality, performance, and reliability across various deployment scenarios (online serving and offline inference).
+
+<table>
+  <thead>
+    <tr>
+      <th>Level</th>
+      <th>Scope & Focus</th>
+      <th>Time Cost</th>
+      <th>Test Dir</th>
+      <th>Doc</th>
+      <th>Frequency</th>
+      <th>Hardware</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td rowspan="2"><strong>Common</strong></td>
+      <td>Contribution Guideline & PR checklist</td>
+      <td>/</td>
+      <td>/</td>
+      <td>.github/PULL_REQUEST_TEMPLATE.md <a href="../tests_style/"> Test Style (PR Checklist)</a></td>
+      <td>/</td>
+      <td>/</td>
+    </tr>
+    <tr>
+      <td>CI Failure Description</td>
+      <td>/</td>
+      <td>/</td>
+      <td><a href="../failures/"> CI Failures</a></td>
+      <td>/</td>
+      <td>/</td>
+    </tr>
+    <tr>
+      <td><strong>L1</strong><br>(Unit & Logic)</td>
+      <td>Unit tests for components like entrypoints, models</td>
+      <td rowspan="2">&lt;15min</td>
+      <td>/tests/{component_name}/test_xxx</td>
+      <td>
+        <a href="#chapter-1-l1-l2-level-testing-unit-testing-and-basic-end-to-end-verification">Chapter 1</a><br>
+        Section 1 L1&amp;L2: Purpose, Test Content, Directory Location, Example
+      </td>
+      <td>PR with ready label (also can run locally)</td>
+      <td>CPU</td>
+    </tr>
+    <tr>
+      <td><strong>L2</strong><br>(E2E across models & GPU-required UT)</td>
+      <td>Online & Offline (basic deployment scenarios):<br>dummy, normal inference function (output format, stream), some instance startup UT</td>
+      <td>
+        /tests/e2e/online_serving/test_{model_name}.py<br>
+        /tests/e2e/offline_inference/test_{model_name}.py
+      </td>
+      <td>
+        <a href="#chapter-1-l1-l2-level-testing-unit-testing-and-basic-end-to-end-verification">Chapter 1</a><br>
+        L1&amp;L2: Purpose, Test Content, Directory Location, Example
+      </td>
+      <td>PR with ready label</td>
+      <td>GPU</td>
+    </tr>
+    <tr>
+      <td><strong>L3</strong><br>(Important Perf & Integration & Accuracy)</td>
+      <td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
+      <td>&lt;30min</td>
+      <td>
+        /tests/e2e/online_serving/test_{model_name}.py<br>
+        /tests/e2e/offline_inference/test_{model_name}.py
+      </td>
+      <td>
+        <a href="#chapter-2-l3-level-testing-core-integration-performance-and-accuracy-verification">Chapter 2</a><br>
+        L3: Purpose, Test Content, Directory Location, Example
+      </td>
+      <td>PR Merged (Also run L1&L2 Tests)</td>
+      <td>GPU</td>
+    </tr>
+    <tr>
+      <td><strong>L4</strong><br>(Perf & Integration & Accuracy)</td>
+      <td>Online & Offline: full functional scenarios + performance test + doc test</td>
+      <td>&lt;3 hour</td>
+      <td>
+        <strong>Full Function:</strong><br>
+        /tests/e2e/online_serving/test_{model_name}_expansion.py<br>
+        /tests/e2e/offline_inference/test_{model_name}_expansion.py<br>
+        <strong>Performance:</strong><br>
+        /tests/e2e/perf/nightly.json<br>
+        <strong>Doc Test:</strong><br>
+        tests/example/online_serving/test_{model_name}.py<br>
+        tests/example/offline_inference/test_{model_name}.py
+      </td>
+      <td>
+        <a href="#chapter-3-l4-level-testing-full-functionality-performance-and-documentation-testing">Chapter 3</a><br>
+        L4: Purpose, Test Content, Directory Location, Example
+      </td>
+      <td>Nightly</td>
+      <td>GPU</td>
+    </tr>
+    <tr>
+      <td><strong>L5</strong><br>(Stability & Reliability)</td>
+      <td>Online & Offline: long-term stability test + reliability test</td>
+      <td> Depends on reality </td>
+      <td>
+        <strong>Stability:</strong><br>
+        tests/e2e/stability/weekly.json<br>
+        <strong>Reliability:</strong><br>
+        tests/e2e/reliability/test_{model_name}.py
+      </td>
+      <td>
+        <a href="#chapter-4-l5-level-testing-stability-and-reliability-testing">Chapter 4</a><br>
+        L5: Purpose, Test Content, Directory Location, Example
+      </td>
+      <td>Weekly / Days before Release</td>
+      <td>GPU</td>
+    </tr>
+  </tbody>
+</table>
+
+
+---
+<details>
+<summary> The folder structure for tests file based on the 5 levels design</summary>
+Legend: `✅` = test exists, `⬜` = suggested to add.
+```
+vllm_omni/                                    tests/
+├── config/                             →     ├── config/
+│   ├── model.py                              │   └── test_model.py                    ⬜
+│   └── lora.py                               │   └── test_lora.py                      ⬜
+│
+├── core/                               →     ├── core/
+│   └── sched/                                 │   └── sched/
+│       ├── omni_ar_scheduler.py               │       ├── test_omni_ar_scheduler.py    ⬜
+│       ├── omni_generation_scheduler.py       │       ├── test_omni_generation_scheduler.py  ⬜
+│       └── output.py                          │       └── test_output.py               ✅ currently in entrypoints/test_omni_new_request_data.py (tests output.OmniNewRequestData)
+│
+├── diffusion/                          →     ├── diffusion/
+│   ├── diffusion_engine.py                    │   ├── test_diffusion_engine.py          ⬜
+│   ├── attention/                             │   ├── attention/
+│   │   ├── layer.py                            │   │   ├── test_attention_sp.py         ✅
+│   │   └── backends/                           │   │   └── test_flash_attn.py           ✅
+│   ├── distributed/                           │   ├── distributed/
+│   │   └── ...                                 │   │   ├── test_comm.py                 ✅
+│   │                                            │   │   ├── test_cfg_parallel.py        ✅
+│   │                                            │   │   └── test_sp_plan_hooks.py       ✅
+│   ├── lora/                                   │   ├── lora/
+│   │   └── ...                                 │   │   ├── test_base_linear.py          ✅
+│   │                                            │   │   └── test_lora_manager.py        ✅
+│   ├── models/                                 │   ├── models/
+│   │   ├── qwen_image/                         │   │   ├── qwen_image/                 (e2e coverage)
+│   │   ├── z_image/                            │   │   └── z_image/
+│   │   └── ...                                 │   │       └── test_zimage_tp_constraints.py  ✅
+│   └── worker/                                 │   └── worker/
+│       ├── diffusion_worker.py                 │       └── test_diffusion_worker.py   ✅ file at diffusion/test_diffusion_worker.py
+│       └── diffusion_model_runner.py            │
+│
+├── distributed/                         →     ├── distributed/
+│   └── omni_connectors/                         │   └── omni_connectors/
+│       ├── adapter.py                           │       ├── test_adapter_and_flow.py   ✅
+│       ├── kv_transfer_manager.py               │       ├── test_basic_connectors.py   ✅
+│       ├── connectors/                           │       ├── test_kv_flow.py             ✅
+│       └── utils/                               │       └── test_omni_connector_configs.py  ✅
+│
+├── engine/                             →     ├── engine/
+│   ├── input_processor.py                      │   ├── test_input_processor.py         ⬜  (no processor.py in source)
+│   ├── output_processor.py                     │   └── test_output_processor.py         ⬜
+│   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
+│
+├── entrypoints/                        →     ├── entrypoints/
+│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
+│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
+│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
+│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
+│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
+│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
+│   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
+│   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
+│   │   └── ...                                  │   │   └── test_*.py                  ⬜
+│   └── openai/                                  │   └── openai_api/                    # maps to entrypoints/openai/
+│       ├── api_server.py                        │       ├── test_api_server.py         ⬜  (e2e indirect coverage)
+│       ├── serving_chat.py                       │       ├── test_serving_chat_sampling_params.py  ✅
+│       ├── serving_speech.py                     │       ├── test_serving_speech.py     ✅
+│       └── image_api_utils.py                   │       └── test_image_server.py      ✅
+│
+├── inputs/                             →     ├── inputs/
+│   ├── data.py                                 │   ├── test_data.py                   ⬜
+│   ├── parse.py                                │   ├── test_parse.py                 ⬜
+│   └── preprocess.py                            │   └── test_preprocess.py            ✅ currently in entrypoints/test_omni_input_preprocessor.py
+│
+├── model_executor/                     →     ├── model_executor/
+│   ├── layers/                                  │   ├── layers/
+│   │   └── mrope.py                             │   │   └── test_mrope.py              ⬜
+│   ├── model_loader/                            │   ├── model_loader/
+│   │   └── weight_utils.py                      │   │   └── test_weight_utils.py      ⬜
+│   ├── models/                                  │   ├── models/
+│   │   ├── qwen2_5_omni/                         │   │   ├── qwen2_5_omni/
+│   │   │   ├── qwen2_5_omni_thinker.py           │   │   │   ├── test_audio_length.py  ✅
+│   │   │   ├── qwen2_5_omni_talker.py            │   │   │   ├── test_qwen2_5_omni_thinker.py  ⬜
+│   │   │   └── qwen2_5_omni_token2wav.py         │   │   │   ├── test_qwen2_5_omni_talker.py  ⬜
+│   │   └── qwen3_omni/                          │   │   │   └── test_qwen2_5_omni_token2wav.py  ⬜
+│   │       └── ...                               │   │   └── qwen3_omni/
+│   ├── stage_configs/                           │   │       └── test_*.py              ⬜
+│   │   └── *.yaml                               │   └── stage_configs/                 (used by e2e, test_*.py can be added)  ⬜
+│   └── stage_input_processors/                  │   └── stage_input_processors/
+│       └── ...                                  │       └── test_*.py                 ⬜
+│
+├── sample/                             →     ├── sample/
+│   └── __init__.py                             │   └── test_*.py                      ⬜
+│
+├── utils/                              →     ├── utils/
+│   └── __init__.py                             │   └── test_*.py                       ⬜  (no platform_utils.py currently)
+│
+├── worker/                             →     ├── worker/
+│   ├── gpu_ar_model_runner.py                  │   ├── test_gpu_ar_model_runner.py    ⬜
+│   ├── gpu_ar_worker.py                        │   ├── test_gpu_ar_worker.py           ⬜
+│   ├── gpu_generation_model_runner.py          │   ├── test_gpu_generation_model_runner.py  ✅
+│   ├── gpu_generation_worker.py                │   ├── test_gpu_generation_worker.py  ⬜
+│   ├── gpu_model_runner.py                     │   ├── test_omni_gpu_model_runner.py   ✅
+│   └── mixins.py                               │   └── (npu under platforms/npu/worker/)  # not worker/npu/
+│
+├── platforms/                          →     (no tests/platforms/, e2e and stage_configs provide indirect coverage)
+│   ├── cuda/
+│   ├── npu/worker/                             # NPU worker here, not vllm_omni/worker/npu/
+│   ├── rocm/
+│   └── xpu/worker/
+│
+├── outputs.py                          →     test_outputs.py                         ✅ (at tests root)
+├── (logger, patch, request, version)    →     (no corresponding unit test)
+│
+└── e2e (tests side only)               →     ├── e2e/
+                                               ├── online_serving/                     ✅ non-empty
+                                               │   ├── test_async_omni.py
+                                               │   ├── test_qwen3_omni.py
+                                               │   ├── test_qwen3_omni_expansion.py
+                                               │   ├── test_image_gen_edit.py
+                                               │   ├── test_images_generations_lora.py
+                                               │   └── stage_configs/
+                                               └── offline_inference/                  ✅
+                                                   ├── test_qwen2_5_omni.py
+                                                   ├── test_qwen3_omni.py
+                                                   ├── test_bagel_text2img.py
+                                                   ├── test_t2i_model.py
+                                                   ├── test_t2v_model.py
+                                                   ├── test_ovis_image.py
+                                                   ├── test_zimage_tensor_parallel.py
+                                                   ├── test_cache_dit.py
+                                                   ├── test_teacache.py
+                                                   ├── test_stable_audio_model.py
+                                                   ├── test_diffusion_cpu_offload.py
+                                                   ├── test_diffusion_layerwise_offload.py
+                                                   ├── test_diffusion_lora.py
+                                                   ├── test_sequence_parallel.py
+                                                   └── stage_configs/
+                                                       ├── qwen2_5_omni_ci.yaml
+                                                       ├── qwen3_omni_ci.yaml
+                                                       ├── bagel_*.yaml
+                                                       └── npu/, rocm/, etc.
+```
+
+
+</details>
+
+
+## Common Specifications
+
+Before entering specific testing levels, the project establishes two common specifications aimed at standardizing the development process and quickly locating issues.
+
+1.  ***PR Checklist ([Tests Style](../ci/tests_style.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance.
+2.  ***CI Failure Explanation ([CI Failures](../ci/failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency.
+
+## Chapter 1: L1 & L2 Level Testing - Unit Testing and Basic End-to-End Verification
+
+### 1.1 Testing Purpose
+
+L1 and L2 level testing form the foundation of the quality assurance system. L1 level testing focuses on verifying the internal logic correctness of code units (e.g., functions, classes), ensuring each independent component behaves as designed.
+
+L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle  .
+
+
+
+### 1.2 Testing Content and Scope
+
+-   ***L1 (Unit & Logic Testing)***:
+-   -   ***Scope***: Tests internal functions and methods of core components such as `entrypoints`, `models`.
+    -   ***Focus***: Branch coverage, exception handling, algorithm logic correctness. Does not involve external dependencies or the complete service stack.
+    -   ***Time Cost***: Execution time is controlled within ***15 minutes*** to ensure fast feedback.
+-   ***L2 (Basic End-to-End Testing)***:
+-   -   ***Scope***: Covers two basic deployment scenarios: `online` (serving) and `offline` (inference).
+    -   ***Focus***: Uses `dummy` models or lightweight real models to verify that the entire chain from request input to result output works normally, including output data structure, streaming (stream) support, etc. Also includes some unit tests that require launching independent service instances.
+    -   ***Characteristic***: Requires ***GPU*** resources to perform model computations.
+
+### 1.3 Test Directory and Execution Files
+
+A clear directory structure is key to managing test cases efficiently.
+
+-   ***L1 Test Directory***: `/tests/{component_name}/test_xxx.py`
+-   -   Here, `{component_name}` corresponds to modules in the source code, such as `distributed`, `entrypoints`, etc., and `test_xxx.py` is the specific test file.
+-   ***L2 Test Directory***:
+-   -   Online Serving: `/tests/e2e/online_serving/test_{model_name}.py`
+    -   Offline Inference: `/tests/e2e/offline_inference/test_{model_name}.py`
+
+### 1.4 Execution Method and Example
+
+-   ***Trigger Timing***: **`PR with ready label`**. That is, when a developer adds a "ready for review" or similar label to a PR on platforms like GitHub, L1 and L2 tests are automatically triggered.
+-   ***Execution Environment***: L1 uses ***CPU*** environment; L2 requires ***GPU*** environment.
+-   ***Script Example***:
+
+<details>
+<summary> L1 Test Examples</summary>
+
+Examples from `tests/model_executor/models/qwen2_5_omni/test_audio_length.py`
+```python
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+def test_resolve_max_mel_frames_default():
+    from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import resolve_max_mel_frames
+
+    assert resolve_max_mel_frames(None, default=30000) == 30000
+    assert resolve_max_mel_frames(None, default=6000) == 6000
+
+
+def test_resolve_max_mel_frames_explicit():
+    from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import resolve_max_mel_frames
+
+    # Explicit argument always wins over default
+    assert resolve_max_mel_frames(123, default=30000) == 123
+    assert resolve_max_mel_frames(6000, default=30000) == 6000
+    assert resolve_max_mel_frames(0, default=30000) == 0
+
+
+@pytest.mark.parametrize("repeats", [2, 4])
+@pytest.mark.parametrize("code_len", [0, 1, 32768])
+@pytest.mark.parametrize("max_mel_frames", [None, -1, 0, 1, 6000, 30000])
+def test_cap_and_align_mel_length_no_mismatch(repeats, code_len, max_mel_frames):
+    """Guard that any max_mel_frames yields a mel length aligned to repeats, and
+    consistent with the truncated code length (prevents concat mismatch).
+    """
+    from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import cap_and_align_mel_length
+
+    target_code_len, target_mel_len = cap_and_align_mel_length(
+        code_len=code_len,
+        repeats=repeats,
+        max_mel_frames=max_mel_frames,
+    )
+
+    assert isinstance(target_code_len, int)
+    assert isinstance(target_mel_len, int)
+
+    if code_len == 0:
+        assert target_code_len == 0
+        assert target_mel_len == 0
+        return
+
+    assert target_code_len >= 1
+    assert target_mel_len >= repeats
+    assert target_mel_len % repeats == 0
+    assert target_mel_len == target_code_len * repeats
+    assert target_code_len <= code_len
+
+    if max_mel_frames is not None and int(max_mel_frames) > 0 and int(max_mel_frames) >= repeats:
+        assert target_mel_len <= int(max_mel_frames)
+```
+</details>
+
+<details>
+<summary> L2 Test Examples</summary>
+You can refer to Test Examples in Chapter 2 to see example test cases that incorporate both L2 and L3 testing logic.
+</details>
+
+-   -   ***Run Command***:
+
+    `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py`
+    `pytest -s -v -m 'core_model and cpu' --run-level=core_model`
+
+## Chapter 2: L3 Level Testing - Core Integration, Performance, and Accuracy Verification
+
+### 2.1 Testing Purpose
+
+L3 level testing executes after code is merged into the main branch. Its core purpose is to verify the integration effect, key performance indicators, and output accuracy of ***real models*** in ***multiple deployment scenarios***
+
+. It acts as the "quality gatekeeper" for the main branch, ensuring that no merge breaks the core capabilities of the model service. Testing needs to provide clear conclusions within a relatively short time (<30min), balancing test depth with feedback speed.
+
+
+
+### 2.2 Testing Content and Scope
+
+-   ***Deployment Scenarios***: Covers richer `online` and `offline` deployment configurations, which may include different hardware configurations, batch sizes, concurrency levels, etc.
+-   ***Core Verification***:
+-   1.  ***Inference Functionality***: Ensures real models can perform forward computation normally and return results.
+    2.  ***Accuracy Compliance***: Verifies that the model's evaluation metrics (e.g., accuracy) meet the expected baseline, preventing code changes from introducing accuracy issues.
+    3.  ***Important Performance***: Verifies whether performance (e.g., P99 latency, throughput) in core scenarios meets preset thresholds.
+
+### 2.3 Test Directory and Execution Files
+
+-   ***Functional Testing***:
+-   -   Online Serving: `/tests/e2e/online_serving/test_{model_name}_expansion.py`
+    -   Offline Inference: `/tests/e2e/offline_inference/test_{model_name}_expansion.py`
+    -   (Note: `_expansion.py` likely means it contains more comprehensive scenario cases compared to L2 tests).
+
+### 2.4 Execution Method and Example
+
+-   ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch.
+-   ***Execution Environment***: ***GPU*** servers.
+-   ***Script Example***:
+
+???+ example "Test Examples"
+
+    **2.4.1 Mark Declaration Section**
+
+    ```python
+    @pytest.mark.advanced_model
+    @pytest.mark.core_model
+    @pytest.mark.parametrize("omni_server", test_params, indirect=True)
+    ```
+
+    **Explanation**:
+
+    @pytest.mark.advanced_model: Marks the test as L3 or L4 level, indicating that this test case performs deep validation, using real models for performance, integration, and accuracy testing. This forms a "basic-advanced" correspondence with the core_model mark at the L2 level.
+
+    @pytest.mark.core_model: Marks the test as L1 or L2 level, indicating that this test case validates the basic functionality of the core model. It uses mock weights and only checks if the relevant interface functions correctly.
+
+    @pytest.mark.parametrize: A parameterization decorator that allows abstracting test data into parameters, enabling reuse of the same test logic across different data configurations. indirect=True indicates that parameters will be passed to the fixture for processing.
+
+    **Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation run logic at the merge or nightly level, you can mark it only with @pytest.mark.advanced_model. If you believe the test case needs to accommodate both basic run and deep validation test logic, you should mark it with both @pytest.mark.core_model and @pytest.mark.advanced_model.
+
+    **2.4.2 Test Function Definition and Documentation**
+
+    ```python
+    def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
+        """
+        Test multi-modal input processing and text/audio output generation via OpenAI API.
+        Deploy Setting: default yaml
+        Input Modal: text + audio + video + image
+        Output Modal: text + audio
+        Input Setting: stream=True
+        Datasets: single request
+        """
+    ```
+
+    **Explanation**:
+
+    **Function Naming Convention**: Uses the test_ prefix, describes the test scenario mix_to_text_audio, and the number 001 indicates the first test case for this scenario.
+
+    **Parameter Explanation**:
+
+    omni_server: Omni server instance obtained via fixture, containing model information and configuration.
+
+    openai_client: Unified OpenAI client processing instance, encapsulating request sending and response validation logic.
+
+    Docstring: Describes the test purpose, deployment settings, input/output modalities, streaming settings, and dataset type in detail, providing clear context for test maintenance.
+
+    **2.4.3 Multimodal Data Preparation**
+
+    ```python
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
+    ```
+
+    **Explanation**:
+
+    **Data Generation Functions**: Use the generate_synthetic_* series of functions to generate synthetic test data, avoiding reliance on external resources and ensuring test reproducibility and stability.
+
+    **Parameter Explanation**:
+
+    Video: width, height, duration_frames
+
+    Image: width, height
+
+    Audio: duration_seconds, channels
+
+    **2.4.4 Request Configuration and Keyword Validation**
+
+    ```python
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "key_words": {
+            "audio": ["water", "cricket"],
+            "video": ["sphere", "globe", "circle", "round"],
+            "image": ["square", "quadrate"],
+            "text": ["beijing"]
+        },
+    }
+    ```
+
+    **Explanation**:
+
+    **Model Specification**: Uses omni_server.model to ensure the test aligns with the model configured on the server.
+
+    **Keyword Validation Mechanism**: This is an innovative design of the template to address the specific needs of multimodal testing:
+
+    Audio Keywords: Validate whether the generated text's description of audio content contains expected elements (e.g., "water" for water sounds, "cricket" for cricket sounds). If you provide multiple keywords, the validation is considered successful if at least one keyword is present.
+
+    **Video Keywords**: Validate whether the generated text's description of video content contains expected elements. If you provide multiple keywords, the validation is considered successful if at least one keyword is present.
+
+    Image Keywords: Validate whether the generated text's description of image content contains expected elements. If you provide multiple keywords, the validation is considered successful if at least one keyword is present.
+
+    Text Keywords: Validate whether the generated text contains expected elements. If you provide multiple keywords, the validation is considered successful if at least one keyword is present.
+
+    **2.4.5 Request Execution**
+
+    ```python
+    openai_client.send_omni_request(request_config, request_num=1)  # for omni-understanding models
+    # or
+    openai_client.send_diffusion_request(request_config, request_num=1)  # for diffusion models
+    ```
+
+    **Explanation**:
+
+    **Unified Client**: Uses the OpenAIClientHandler instance to send requests. This client encapsulates error handling, retry mechanisms, and response validation logic.
+
+    **Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
+
+    **Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
+
+-   ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
+
+## Chapter 3: L4 Level Testing - Full Functionality, Performance, and Documentation Testing
+
+### 3.1 Testing Purpose
+
+L4 level testing is a comprehensive quality audit before a version release. It expands upon L3, executing ***full*** functional scenarios, conducting systematic ***performance stress tests***, and simultaneously verifying the correctness of accompanying ***example documentation***. Its purpose is to perform deep validation of the system during off-peak nighttime hours, providing quality trend reports for daytime development and data support for release decisions.
+
+
+
+### 3.2 Testing Content and Scope
+
+-   ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
+-   ***Performance Testing***: Uses the `/tests/e2e/perf/nightly.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
+-   ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
+
+### 3.3 Test Directory and Execution Files
+
+-   ***Functional Testing***: Same directories as L3.
+-   ***Performance Test Configuration***: `/tests/e2e/perf/nightly.json`
+-   ***Documentation Example Tests***:
+-   -   `tests/example/online_serving/test_{model_name}.py`
+    -   `tests/example/offline_inference/test_{model_name}.py`
+
+### 3.4 Execution Method and Example
+
+-   ***Trigger Timing***: **`Nightly`**, automatically executed every night.
+-   ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing.
+-   ***Script Example***:
+<details>
+<summary> Test Examples</summary>
+When you want to add L4-level performance test cases, you can refer to the following format for case addition in tests/perf/tests/test.json:
+
+```JSON
+{
+    "test_name": "test_qwen3_omni",
+    "server_params": {
+        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "stage_config_name": "qwen3_omni.yaml"
+    },
+    "benchmark_params": [
+        {
+            "dataset_name": "random",
+            "num_prompts": [10, 20],
+            "request_rate": [0.5, 1],
+            "random_input_len": 2500,
+            "random_output_len": 900,
+            "ignore_eos": true,
+            "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+            "baseline": {
+                "mean_ttft_ms": 100000,
+                "mean_audio_ttfp_ms": 100000,
+                "mean_audio_rtf": 100000
+            }
+        }
+    ]
+}
+```
+
+#### Parameter Explanation
+
+***Overview***
+
+| Field            | Required | Description                                                     |
+| ---------------- | -------- | --------------------------------------------------------------- |
+| test_name        | Yes      | Unique identifier for the test case                             |
+| server_params    | Yes      | Server-side configuration parameters                            |
+| benchmark_params | Yes      | Benchmark running parameters (supports multiple configurations) |
+
+#### server_params Configuration
+
+##### Basic Parameters
+
+| Parameter         | Required | Example                            | Description                   |
+| ----------------- | -------- | ---------------------------------- | ----------------------------- |
+| model             | Yes      | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path            |
+| stage_config_name | Yes      | "qwen3_omni.yaml"                  | Stage configuration file name |
+
+##### Dynamic Configuration (update/delete)
+
+Supports incremental modifications based on the basic configuration:
+
+| Operation | Description                          |
+| --------- | ------------------------------------ |
+| update    | Update or add configuration items    |
+| delete    | Delete specified configuration items |
+
+***Example***:
+```
+"update": {
+    "async_chunk": true,  // Enable asynchronous chunk processing
+    "stage_args": {
+        "0": {
+            "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
+        }
+    }
+},
+"delete": {
+    "stage_args": {
+        "2": ["custom_process_input_func"]  // Delete this configuration for stage 2
+    }
+}
+```
+#### benchmark_params Configuration
+
+You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows:
+
+1.  Change the ---xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
+2.  For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
+3.  Add the baseline parameter to specify the required validation values, ensuring the validation metric names match those in the result.json generated by the benchmark.
+4.  The qps and concurrency modes are mutually exclusive. For detailed explanations, see the table below:
+
+| Parameter       | Type        | Required | Example/Values  | Description                                                                                                                                                                                                                                                          |
+| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
+| request_rate    | int / array | No       | 1, [1, 2, 3]    | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
+| max_concurrency | int / array | No       | 1, [1, 2, 3]    | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
+</details>
+
+-   -   ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
+
+## Chapter 4: L5 Level Testing - Stability and Reliability Testing
+
+### 4.1 Testing Purpose
+
+L5 level testing focuses on the performance of model services under ***long-running*** and ***abnormal fault*** scenarios. It aims to uncover deep-seated issues that only manifest under sustained pressure or extreme conditions, such as memory leaks, resource contention, gradual performance degradation, and lack of fault tolerance mechanisms. This is the final, yet crucial, line of defense for ensuring service high availability and production environment robustness.
+
+
+
+### 4.2 Testing Content and Scope
+
+-   ***Long-term Stability (Stability) Testing***: Uses the `tests/e2e/stability/weekly.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
+-   ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
+
+### 4.3 Test Directory and Execution Files
+
+-   ***Stability Test Configuration***: `tests/e2e/stability/weekly.json`
+-   ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
+
+### 4.4 Execution Method and Example
+
+-   ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower.
+-   ***Execution Environment***: ***GPU*** servers, requiring a stable and exclusive testing environment.
+-   ***Script Example***:
+<details>
+<summary> Test Examples</summary>
+```python
+# WIP
+```
+</details>
+
+-   -   ***Stability***: (Execution would be driven by the configuration in `weekly.json`).
+    -   ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
+
+## Summary
+
+This multi-level testing system achieves continuous, progressive validation of model service quality by tightly integrating testing activities with the development workflow (commit, review, merge, release). From rapid unit testing to comprehensive end-to-end testing, and further to in-depth performance, stability, and reliability verification, each level has clear objectives, collectively building a robust quality protection net. By following this system, teams can deliver high-quality, highly reliable model services more efficiently.
@@ -1,4 +1,10 @@
 # CI Failures

-What should I do when a CI job fails on my PR, but I don't think my PR caused
-the failure?
+What should I do when a CI job fails on my PR, but I don't think my PR caused the failure?
+
+
+## Common Case 1:
+
+🚨 Error: The command was interrupted by a signal: signal: terminated
+
+Reason: The test is terminated due to exceed the time limits. Sometimes rebuild the test will pass.
@@ -0,0 +1,75 @@
+# Test Guide
+## Setting Up the Test Environment
+### Creating a Container
+vLLM-Omni provides an official Docker image for deployment. These images are built upon vLLM Docker images and are available on [Docker Hub](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which vLLM release it is based on.
+For a local test environment, you can follow the steps below to create a container:
+## Installing Dependencies
+### vLLM & vLLM-Omni
+vLLM-Omni is built based on vLLM. You can follow [install guide](../../getting_started/installation/README.md) to build your local environment.
+
+### Test Case Dependencies
+When running test cases, you may need to install the following dependencies:
+
+```bash
+uv pip install ".[dev]"
+apt-get install -y ffmpeg
+```
+
+## Running Tests
+Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases.
+
+=== "L1 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "core_model and cpu"
+    ```
+    The latest test command is available in the "Simple Unit Test" step of this [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
+
+=== "L2 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "core_model and not cpu" --run-level=core_model
+    ```
+    If you only want to run a specific test case, you can use:
+    ```bash
+    pytest -s -v test_xxxx.py --run-level=core_model
+    ```
+    If you only want to run specific test cases on a particular platform, you can use:
+    ```bash
+    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
+    ```
+    The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
+
+=== "L3 level & L4 level"
+
+    ```bash
+    cd tests
+    pytest -s -v -m "advanced_model" --run-level=advanced_model
+    ```
+    If you only want to run L3 test case, you can use:
+    ```bash
+    pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model
+    ```
+    If you only want to run a specific test case, you can use:
+    ```bash
+    pytest -s -v test_xxxx.py --run-level=advanced_model
+    ```
+    If you only want to run specific test cases on a particular platform, you can use:
+    ```bash
+    pytest -s -v -m "core_model and distributed_cuda and L4"  --run-level=core_model
+    ```
+    Note: To run performance tests, use:
+    ```bash
+    pytest -s -v perf/scripts/run_benchmark.py
+    ```
+
+    The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
+
+    The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml).
+
+You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
+
+## Adding New Test Cases
+Please refer to the [L5 Layering Specification document](./CI_5levels.md).
@@ -5,33 +5,35 @@ By adding markers before test functions, tests can later be executed uniformly b
 ## Current Markers
 Defined in `pyproject.toml`:

-| Marker             | Description                                             |
-| ------------------ | ------------------------------------------------------- |
-| `core_model`       | Core model tests (run in each PR)                       |
-| `diffusion`        | Diffusion model tests                                   |
-| `omni`             | Omni model tests                                        |
-| `cache`            | Cache backend tests                                     |
-| `parallel`         | Parallelism/distributed tests                           |
-| `cpu`              | Tests that run on CPU                                   |
-| `gpu`              | Tests that run on GPU (auto-added)                      |
-| `cuda`             | Tests that run on CUDA (auto-added)                     |
-| `rocm`             | Tests that run on AMD/ROCm (auto-added)                 |
-| `npu`              | Tests that run on NPU/Ascend (auto-added)               |
-| `H100`             | Tests that require H100 GPU                             |
-| `L4`               | Tests that require L4 GPU                               |
-| `MI325`            | Tests that require MI325 GPU (AMD/ROCm)                 |
-| `A2`               | Tests that require A2 NPU                               |
-| `A3`               | Tests that require A3 NPU                               |
-| `distributed_cuda` | Tests that require multi cards on CUDA platform         |
-| `distributed_rocm` | Tests that require multi cards on ROCm platform         |
-| `distributed_npu`  | Tests that require multi cards on NPU platform          |
-| `skipif_cuda`      | Skip if the num of CUDA cards is less than the required |
-| `skipif_rocm`      | Skip if the num of ROCm cards is less than the required |
-| `skipif_npu`       | Skip if the num of NPU cards is less than the required  |
-| `slow`             | Slow tests (may skip in quick CI)                       |
-| `benchmark`        | Benchmark tests                                         |
+| Marker             | Description                                               |
+| ------------------ | --------------------------------------------------------- |
+| `core_model`       | L1&L2 tests (run in each PR)                              |
+| `advanced_model`   | L3&L4 level tests (run in each merge or nightly)          |
+| `diffusion`        | Diffusion model tests                                     |
+| `omni`             | Omni model tests                                          |
+| `cache`            | Cache backend tests                                       |
+| `parallel`         | Parallelism/distributed tests                             |
+| `cpu`              | Tests that run on CPU                                     |
+| `gpu`              | Tests that run on GPU *                                   |
+| `cuda`             | Tests that run on CUDA *                                  |
+| `rocm`             | Tests that run on AMD/ROCm *                              |
+| `xpu`              | Tests that run on Intel XPU *                             |
+| `npu`              | Tests that run on NPU/Ascend *                            |
+| `H100`             | Tests that require H100 GPU  *                            |
+| `L4`               | Tests that require L4 GPU *                               |
+| `MI325`            | Tests that require MI325 GPU (AMD/ROCm) *                 |
+| `A2`               | Tests that require A2 NPU *                               |
+| `A3`               | Tests that require A3 NPU *                               |
+| `distributed_cuda` | Tests that require multi cards on CUDA platform *         |
+| `distributed_rocm` | Tests that require multi cards on ROCm platform  *        |
+| `distributed_npu`  | Tests that require multi cards on NPU platform  *         |
+| `skipif_cuda`      | Skip if the num of CUDA cards is less than the required * |
+| `skipif_rocm`      | Skip if the num of ROCm cards is less than the required * |
+| `skipif_npu`       | Skip if the num of NPU cards is less than the required *  |
+| `slow`             | Slow tests (may skip in quick CI)                         |
+| `benchmark`        | Benchmark tests                                           |

-For those markers shown as auto-added, they will be added by the `@hardware_test` decorator.
+\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).

 ### Example usage for markers

@@ -48,12 +50,13 @@ from tests.utils import hardware_test
 def test_video_to_audio()
    ...
 ```
+
 ### Decorator: `@hardware_test`

 This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:

 1. **Applies platform and resource markers**  
-   Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `A2`, `A3`).
+   Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
   ```
   @pytest.mark.cuda
   @pytest.mark.L4
@@ -71,36 +74,50 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
   Support for `skipif_rocm` and `skipif_npu` will be implemented later.


-5. **Runs each test in a new process**  
-   Automatically wraps the distributed test with a decorator (`@create_new_process_for_each_test`) to ensure isolation and compatibility with multi-process hardware backends.
-
-6. **Works with pytest filtering**  
+5. **Works with pytest filtering**  
   Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "distributed_cuda and L4"`).

 #### Example usage for decorator
 - Single call for multiple platforms:
    ```python
    @hardware_test(
-        res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
-        num_cards={"cuda": 2, "rocm": 2, "npu": 2},
+        res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
+        num_cards={"cuda": 2, "rocm": 2, "xpu": 2, "npu": 2},
    )
    ```
    or
    ```python
    @hardware_test(
-        res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+        res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
        num_cards=2,
    )
    ```
 - `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), NPU (A2/A3)
 - `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing
- `hardware_test` automatically applies `@create_new_process_for_each_test` for distributed tests.
 - Distributed markers (`distributed_cuda`, `distributed_rocm`, `distributed_npu`) are auto-added for multi-card cases
 - Filtering examples:
    - CUDA only: `pytest -m "distributed_cuda and L4"`
    - ROCm only: `pytest -m "distributed_rocm and MI325"`
    - NPU only: `pytest -m "distributed_npu"`

+### Function: `hardware_marks`
+
+`hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function.
+
+```python
+from tests.utils import hardware_marks
+
+MULTI_CARD_MARKS = hardware_marks(
+    res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2
+)
+
+@pytest.mark.parametrize("omni_server", [
+    pytest.param(OmniServerParams(...), id="case_001", marks=MULTI_CARD_MARKS),
+], indirect=True)
+def test_feature(omni_server):
+    ...
+```
+
 ## Add Support for a New Platform

 If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
@@ -135,12 +152,13 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
           # Optionally: add skipif_tpu when implemented
           return [test_platform, test_resource, test_distributed]
   ```
-3. **Update `hardware_test` to recognize your new platform**:
-    In the relevant place (see the `hardware_test` implementation), add:
+3. **Update `hardware_marks` to recognize your new platform**:
+    In the relevant place (see the `hardware_marks` implementation), add:
    ```python
    if platform == "tpu":
        marks = tpu_marks(res=resource, num_cards=cards)
    ```
+    (`hardware_test` calls `hardware_marks` internally, so both will pick up the change.)
 4. **(Recommended) Add a test using your new markers**:
   ```python
   @hardware_test(
@@ -154,7 +172,7 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
 **Summary**:  
 - Add pytest markers for your new platform/resources  
 - Implement a marker function (`xxx_marks`)  
- Plug into `hardware_test`  
- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation!
+- Plug into `hardware_marks`  
+- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!

 See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
@@ -2,8 +2,18 @@

 To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes. This document outlines the guidelines for organizing and naming test files.

+## Checklist before submitting your test files
+
+1. The file is saved in an appropriate place and the file name is clear.
+2. The coding style follows the requirements outlined below.
+3. All test functions have appropriate pytest markers.
+4. For tests that need run in CI, please ensure it labeled as ``@pytest.mark.core_model` the test is configured under the `./buildkite/` folder.
+
+
 ## Test Types

+For more details about our [Five Levels Tests design](../ci/CI_5levels.md).
+
 ### Unit Tests and System Tests
 For unit tests and system tests, we strongly recommend placing test files in the same directory structure as the source code being tested, using the naming convention `test_*.py`.

@@ -18,95 +28,142 @@ End-to-end tests verify the complete functionality of a system or component. For

 ## Test Directory Structure

-The ideal directory structure mirrors the source code organization:
+The ideal directory structure mirrors the source code organization. Legend: `✅` = test exists, `⬜` = suggested to add.

 ```
-vllm_omni/                          tests/
-├── config/                    →    ├── config/
-│   └── model.py                    │   └── test_model.py
+vllm_omni/                                    tests/
+├── config/                             →     ├── config/
+│   ├── model.py                              │   └── test_model.py                    ⬜
+│   └── lora.py                               │   └── test_lora.py                      ⬜
 │
-├── core/                      →    ├── core/
-│   └── sched/                      │   └── sched/                    # Maps to core/sched/
-│       ├── omni_ar_scheduler.py    │       ├── test_omni_ar_scheduler.py
-│       ├── omni_generation_scheduler.py │  ├── test_omni_generation_scheduler.py
-│       └── output.py               │       └── test_output.py
+├── core/                               →     ├── core/
+│   └── sched/                                 │   └── sched/
+│       ├── omni_ar_scheduler.py               │       ├── test_omni_ar_scheduler.py    ⬜
+│       ├── omni_generation_scheduler.py       │       ├── test_omni_generation_scheduler.py  ⬜
+│       └── output.py                          │       └── test_output.py               ✅ currently in entrypoints/test_omni_new_request_data.py (tests output.OmniNewRequestData)
 │
-├── diffusion/                 →    ├── diffusion/
-│   ├── diffusion_engine.py         │   ├── test_diffusion_engine.py
-│   ├── omni_diffusion.py           │   ├── test_omni_diffusion.py
-│   ├── attention/                  │   ├── attention/                # Maps to diffusion/attention/
-│   │   └── backends/               │   │   └── test_*.py
-│   ├── models/                     │   ├── models/                   # Maps to diffusion/models/
-│   │   ├── qwen_image/             │   │   ├── qwen_image/
-│   │   │   └── ...                 │   │   │   └── test_*.py
-│   │   └── z_image/                │   │   └── z_image/
-│   │       └── ...                 │   │       └── test_*.py
-│   └── worker/                     │   └── worker/                   # Maps to diffusion/worker/
-│       └── ...                     │       └── test_*.py
+├── diffusion/                          →     ├── diffusion/
+│   ├── diffusion_engine.py                    │   ├── test_diffusion_engine.py          ⬜
+│   ├── attention/                             │   ├── attention/
+│   │   ├── layer.py                            │   │   ├── test_attention_sp.py         ✅
+│   │   └── backends/                           │   │   └── test_flash_attn.py           ✅
+│   ├── distributed/                           │   ├── distributed/
+│   │   └── ...                                 │   │   ├── test_comm.py                 ✅
+│   │                                            │   │   ├── test_cfg_parallel.py        ✅
+│   │                                            │   │   └── test_sp_plan_hooks.py       ✅
+│   ├── lora/                                   │   ├── lora/
+│   │   └── ...                                 │   │   ├── test_base_linear.py          ✅
+│   │                                            │   │   └── test_lora_manager.py        ✅
+│   ├── models/                                 │   ├── models/
+│   │   ├── qwen_image/                         │   │   ├── qwen_image/                 (e2e coverage)
+│   │   ├── z_image/                            │   │   └── z_image/
+│   │   └── ...                                 │   │       └── test_zimage_tp_constraints.py  ✅
+│   └── worker/                                 │   └── worker/
+│       ├── diffusion_worker.py                 │       └── test_diffusion_worker.py   ✅ file at diffusion/test_diffusion_worker.py
+│       └── diffusion_model_runner.py            │
 │
-├── distributed/               →    ├── distributed/
-│   └── ...                         │   └── test_*.py
+├── distributed/                         →     ├── distributed/
+│   └── omni_connectors/                         │   └── omni_connectors/
+│       ├── adapter.py                           │       ├── test_adapter_and_flow.py   ✅
+│       ├── kv_transfer_manager.py               │       ├── test_basic_connectors.py   ✅
+│       ├── connectors/                           │       ├── test_kv_flow.py             ✅
+│       └── utils/                               │       └── test_omni_connector_configs.py  ✅
 │
-├── engine/                    →    ├── engine/
-│   ├── processor.py                │   ├── test_processor.py
-│   └── output_processor.py         │   └── test_output_processor.py
+├── engine/                             →     ├── engine/
+│   ├── input_processor.py                      │   ├── test_input_processor.py         ⬜  (no processor.py in source)
+│   ├── output_processor.py                     │   └── test_output_processor.py         ⬜
+│   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
-├── entrypoints/               →    ├── entrypoints/
-│   ├── omni_llm.py                 │   ├── test_omni_llm.py          # UT: OmniLLM core logic (mocked)
-│   ├── omni_stage.py               │   ├── test_omni_stage.py         # UT: OmniStage logic
-│   ├── omni.py                     │   ├── test_omni.py               # E2E: Omni class (offline inference)
-│   ├── async_omni.py               │   ├── test_async_omni.py         # E2E: AsyncOmni class
-│   ├── cli/                        │   ├── cli/                       # Maps to entrypoints/cli/
-│   │   └── ...                     │   │   └── test_*.py
-│   └── openai/                     │   └── openai/                     # Maps to entrypoints/openai/
-│       ├── api_server.py           │       ├── test_api_server.py     # E2E: API server (online serving)
-│       └── serving_chat.py         │       └── test_serving_chat.py
+├── entrypoints/                        →     ├── entrypoints/
+│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
+│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
+│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
+│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
+│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
+│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
+│   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
+│   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
+│   │   └── ...                                  │   │   └── test_*.py                  ⬜
+│   └── openai/                                  │   └── openai_api/                    # maps to entrypoints/openai/
+│       ├── api_server.py                        │       ├── test_api_server.py         ⬜  (e2e indirect coverage)
+│       ├── serving_chat.py                       │       ├── test_serving_chat_sampling_params.py  ✅
+│       ├── serving_speech.py                     │       ├── test_serving_speech.py     ✅
+│       └── image_api_utils.py                   │       └── test_image_server.py      ✅
 │
-├── inputs/                    →    ├── inputs/
-│   ├── data.py                     │   ├── test_data.py
-│   ├── parse.py                    │   ├── test_parse.py
-│   └── preprocess.py               │   └── test_preprocess.py
+├── inputs/                             →     ├── inputs/
+│   ├── data.py                                 │   ├── test_data.py                   ⬜
+│   ├── parse.py                                │   ├── test_parse.py                 ⬜
+│   └── preprocess.py                            │   └── test_preprocess.py            ✅ currently in entrypoints/test_omni_input_preprocessor.py
 │
-├── model_executor/            →    ├── model_executor/
-│   ├── layers/                     │   ├── layers/
-│   │   └── mrope.py                │   │   └── test_mrope.py
-│   ├── model_loader/               │   ├── model_loader/
-│   │   └── weight_utils.py         │   │   └── test_weight_utils.py
-│   ├── models/                     │   ├── models/
-│   │   ├── qwen2_5_omni/           │   │   ├── qwen2_5_omni/
-│   │   │   ├── qwen2_5_omni_thinker.py │ │   │   ├── test_qwen2_5_omni_thinker.py  # UT
-│   │   │   ├── qwen2_5_omni_talker.py │ │   │   ├── test_qwen2_5_omni_talker.py  # UT
-│   │   │   └── qwen2_5_omni_token2wav.py │ │   │   └── test_qwen2_5_omni_token2wav.py  # UT
-│   │   └── qwen3_omni/             │   │   └── qwen3_omni/
-│   │       └── ...                 │   │       └── test_*.py
-│   ├── stage_configs/              │   └── stage_configs/             # Configuration tests (if needed)
-│   │   └── ...                     │       └── test_*.py
-│   └── stage_input_processors/     │   └── stage_input_processors/
-│       └── ...                     │       └── test_*.py
+├── model_executor/                     →     ├── model_executor/
+│   ├── layers/                                  │   ├── layers/
+│   │   └── mrope.py                             │   │   └── test_mrope.py              ⬜
+│   ├── model_loader/                            │   ├── model_loader/
+│   │   └── weight_utils.py                      │   │   └── test_weight_utils.py      ⬜
+│   ├── models/                                  │   ├── models/
+│   │   ├── qwen2_5_omni/                         │   │   ├── qwen2_5_omni/
+│   │   │   ├── qwen2_5_omni_thinker.py           │   │   │   ├── test_audio_length.py  ✅
+│   │   │   ├── qwen2_5_omni_talker.py            │   │   │   ├── test_qwen2_5_omni_thinker.py  ⬜
+│   │   │   └── qwen2_5_omni_token2wav.py         │   │   │   ├── test_qwen2_5_omni_talker.py  ⬜
+│   │   └── qwen3_omni/                          │   │   │   └── test_qwen2_5_omni_token2wav.py  ⬜
+│   │       └── ...                               │   │   └── qwen3_omni/
+│   ├── stage_configs/                           │   │       └── test_*.py              ⬜
+│   │   └── *.yaml                               │   └── stage_configs/                 (used by e2e, test_*.py can be added)  ⬜
+│   └── stage_input_processors/                  │   └── stage_input_processors/
+│       └── ...                                  │       └── test_*.py                 ⬜
 │
-├── sample/                    →    ├── sample/
-│   └── ...                         │   └── test_*.py
+├── sample/                             →     ├── sample/
+│   └── __init__.py                             │   └── test_*.py                      ⬜
 │
-├── utils/                     →    ├── utils/
-│   └── platform_utils.py           │   └── test_platform_utils.py
+├── utils/                              →     ├── utils/
+│   └── __init__.py                             │   └── test_*.py                       ⬜  (no platform_utils.py currently)
 │
-├── worker/                    →    ├── worker/
-    ├── gpu_ar_worker.py            │   ├── test_gpu_ar_worker.py
-    ├── gpu_generation_worker.py    │   ├── test_gpu_generation_worker.py
-    ├── gpu_model_runner.py         │   ├── test_gpu_model_runner.py
-    └── npu/                        │   └── npu/                       # Maps to worker/npu/
-        └── ...                     │       └── test_*.py
+├── worker/                             →     ├── worker/
+│   ├── gpu_ar_model_runner.py                  │   ├── test_gpu_ar_model_runner.py    ⬜
+│   ├── gpu_ar_worker.py                        │   ├── test_gpu_ar_worker.py           ⬜
+│   ├── gpu_generation_model_runner.py          │   ├── test_gpu_generation_model_runner.py  ✅
+│   ├── gpu_generation_worker.py                │   ├── test_gpu_generation_worker.py  ⬜
+│   ├── gpu_model_runner.py                     │   ├── test_omni_gpu_model_runner.py   ✅
+│   └── mixins.py                               │   └── (npu under platforms/npu/worker/)  # not worker/npu/
 │
-└── e2e/                       →    ├── e2e/                # End-to-end scenarios (no 1:1 source mirror)
-                                    ├── online_serving/       # Full-stack online serving flows
-                                    │   └── (empty for now)
-                                    └── offline_inference/    # Full offline inference flows
-                                        ├── test_qwen2_5_omni.py     # Moved from multi_stages/
-                                        ├── test_qwen3_omni.py       # Moved from multi_stages_h100/
-                                        ├── test_t2i_model.py  # Moved from single_stage/
-                                        └── stage_configs/           # Shared stage configs
-                                            ├── qwen2_5_omni_ci.yaml
-                                            └── qwen3_omni_ci.yaml
+├── platforms/                          →     (no tests/platforms/, e2e and stage_configs provide indirect coverage)
+│   ├── cuda/
+│   ├── npu/worker/                             # NPU worker here, not vllm_omni/worker/npu/
+│   ├── rocm/
+│   └── xpu/worker/
+│
+├── outputs.py                          →     test_outputs.py                         ✅ (at tests root)
+├── (logger, patch, request, version)    →     (no corresponding unit test)
+│
+└── e2e (tests side only)               →     ├── e2e/
+                                               ├── online_serving/                     ✅ non-empty
+                                               │   ├── test_async_omni.py
+                                               │   ├── test_qwen3_omni.py
+                                               │   ├── test_qwen3_omni_expansion.py
+                                               │   ├── test_image_gen_edit.py
+                                               │   ├── test_images_generations_lora.py
+                                               │   └── stage_configs/
+                                               └── offline_inference/                  ✅
+                                                   ├── test_qwen2_5_omni.py
+                                                   ├── test_qwen3_omni.py
+                                                   ├── test_bagel_text2img.py
+                                                   ├── test_t2i_model.py
+                                                   ├── test_t2v_model.py
+                                                   ├── test_ovis_image.py
+                                                   ├── test_zimage_tensor_parallel.py
+                                                   ├── test_cache_dit.py
+                                                   ├── test_teacache.py
+                                                   ├── test_stable_audio_model.py
+                                                   ├── test_diffusion_cpu_offload.py
+                                                   ├── test_diffusion_layerwise_offload.py
+                                                   ├── test_diffusion_lora.py
+                                                   ├── test_sequence_parallel.py
+                                                   ├── test_qwen_image_edit_expansion.py
+                                                   └── stage_configs/
+                                                       ├── qwen2_5_omni_ci.yaml
+                                                       ├── qwen3_omni_ci.yaml
+                                                       ├── bagel_*.yaml
+                                                       └── npu/, rocm/, etc.
 ```


@@ -117,6 +174,8 @@ vllm_omni/                          tests/

 - **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`

+- **Expansion Tests**
+
 ### Best Practices

 1. **Mirror Source Structure**: Test directories should mirror the source code structure
@@ -144,81 +203,217 @@ vllm_omni/                          tests/
 ### Template
 #### E2E - Online serving

+E2E Online tests for Qwen3-Omni model with mix input and audio+text output. Based on `tests/e2e/online_serving/test_qwen3_omni.py`.
+
 ```python
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
-Online E2E smoke test for an omni model (video,text,audio → audio).
+E2E Online tests for Qwen3-Omni model with mix input and audio+text output.
 """
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+import threading
 from pathlib import Path

-import pytest
 import openai
+import pytest

-from tests.utils import hardware_test
+from tests.conftest import (
+    OmniServer,
+    convert_audio_to_text,
+    cosine_similarity_text,
+    dummy_messages_from_mix_data,
+    generate_synthetic_video,
+    merge_base64_and_convert_to_text,
+)
+from vllm_omni.platforms import current_omni_platform

-# Optional: set process start method for workers
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+# Edit: model name and stage config path
+models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
+
+#If you use the default configuration file, you can directly use the following address.
+def get_default_config():
+    return str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
+
+#If you need to modify the configuration file, you can use modify_stage_config.
+def get_chunk_config():
+    path = modify_stage_config(
+        get_default_config(),
+        updates={
+            "async_chunk": True,
+            "stage_args": {
+                0: {
+                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
+                },
+                1: {
+                    "engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk"
+                },
+            },
+        },
+        deletes={"stage_args": {2: ["custom_process_input_func"]}},
+    )
+    return path
+
+stage_configs = [get_default_config(), CHUNK_CONFIG_PATH]

-models = ["{your model name}"] #Edit here to load your model
-stage_configs = [str(Path(__file__).parent / "stage_configs" / {your model yaml})] #Edit here to load your model yaml
 test_params = [(model, stage_config) for model in models for stage_config in stage_configs]

-#OmniServer，Used to start the vllm-omni server
-class OmniServer:
-    xxx
+
+#Please use this method to launch the online instance.
+_omni_server_lock = threading.Lock()
+
+@pytest.fixture(scope="module")
+def omni_server(request):
+    """Start vLLM-Omni server as a subprocess. Use module scope for multi-stage init (10-20+ min)."""
+    with _omni_server_lock:
+        model, stage_config_path = request.param
+        with OmniServer(
+            model,
+            ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"],
+        ) as server:
+            yield server


@pytest.fixture
-def omni_server(request):
-    model, stage_config_path = request.param
-    with OmniServer(model, ["--stage-configs-path", stage_config_path]) as server:
-        yield server
-
-
-#handle request message
-@pytest.fixture(scope="session")
-def base64_encoded_video() -> str:
-    xxx
-
-@pytest.fixture(scope="session")
-def dummy_messages_from_video_data(video_data_url: str, content_text: str) -> str:
-    xxx
-
-@pytest.mark.core_model
-@pytest.mark.omni
-@hardware_test(
-    res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
-    num_cards={"cuda": 2, "rocm": 2, "npu": 4},
-)
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_video_to_audio(
-    client: openai.OpenAI,
-    omni_server,
-    base64_encoded_video: str,
-) -> None:
-    #set message
-    video_data_url = f"data:video/mp4;base64, {base64_encoded_video}"
-    messages = dummy_messages_from_video_data(video_data_url)
-
-    #send request
-    chat_completion = client.chat.completions.create(
-        model=omni_server.model,
-        messages=messages,
+def client(omni_server):
+    """OpenAI client for the running vLLM-Omni server."""
+    return openai.OpenAI(
+        base_url=f"http://{omni_server.host}:{omni_server.port}/v1",
+        api_key="EMPTY",
    )

-    #verify text output
-    text_choice = chat_completion.choices[0]
-    assert text_choice.finish_reason == "length"
+#Please use function definitions above the test function to define the prompts and other parameters you need.
+def get_system_prompt():
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, "
+                    "Alibaba Group, capable of perceiving auditory and visual inputs, "
+                    "as well as generating text and speech."
+                ),
+            }
+        ],
+    }

-    #verify audio output
-    audio_choice = chat_completion.choices[1]
-    audio_message = audio_choice.message
-    if hasattr(audio_message, "audio") and audio_message.audio:
-        assert audio_message.audio.data is not None
-        assert len(audio_message.audio.data) > 0
+...
+
+#Please define test case tags according to the instructions in the marker documentation.
+@pytest.mark.core_model
+@pytest.mark.omni
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None:
+    # PLEASE FOLLOW THESE TEMPLATE INSTRUCTIONS:
+    # ============================================================================
+    # TEMPLATE USAGE GUIDE:
+    # 1. Copy this entire function as a starting point for multi-modal tests
+    # 2. Update the test name to reflect your specific test scenario
+    # 3. Modify input/output modalities as needed (see OPTIONS section below)
+    # 4. Adjust assertions based on your expected outcomes
+    # 5. Add custom validation logic for your specific use case
+    # ============================================================================
+
+    #Please list the relevant test points.
+    """
+    Test multi-modal input processing and text/audio output generation via OpenAI API.
+    Deploy Setting: default yaml
+    Input Modal: text + audio + video + image
+    Output Modal: text + audio
+    Input Setting: stream=True
+    Datasets: single request
+    """
+    # SECTION 1: TEST SETUP AND INITIALIZATION
+    # =========================================
+    # INSTRUCTIONS: Initialize test variables and prepare test environment
+    # MODIFY: Add any additional test setup required for your scenario
+    e2e_list = list()
+    # SECTION 2: TEST DATA GENERATION
+    # ================================
+    # INSTRUCTIONS: Generate or load test data for each input modality
+    # MODIFY: Replace synthetic generators with your actual data sources
+    # VIDEO DATA - Generate synthetic video for testing
+    # FORMAT: data:video/mp4;base64,{base64_encoded_video}
+    # PARAMETERS: width, height, duration_frames
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
+    # IMAGE DATA - Generate synthetic image for testing
+    # FORMAT: data:image/jpeg;base64,{base64_encoded_image}
+    # PARAMETERS: width, height
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+    # AUDIO DATA - Generate synthetic audio for testing
+    # FORMAT: data:audio/wav;base64,{base64_encoded_audio}
+    # PARAMETERS: duration_seconds, channels
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
+
+    # SECTION 3: MESSAGE CONSTRUCTION
+    # ================================
+    # INSTRUCTIONS: Assemble the complete message payload for API request
+    # MODIFY: Add/remove modalities or change prompt structure as needed
+
+    # USAGE: Construct a message containing all input modalities
+    # IMPORTANT: Ensure the message structure matches OpenAI API expectations
+    # CUSTOMIZATION POINTS:
+    #   - system_prompt: Controls the assistant's behavior
+    #   - content_text: The user's text prompt/question
+    #   - *_data_url: URLs for media content (video/image/audio)
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        video_data_url=video_data_url,
+        image_data_url=image_data_url,
+        audio_data_url=audio_data_url,
+        content_text=get_prompt("mix"),
+    )
+
+    # SECTION 4: API REQUEST EXECUTION
+    # =================================
+    # INSTRUCTIONS: Make the API call and measure performance
+    # MODIFY: Add timeout, retry logic, or additional parameters
+    start_time = time.perf_counter()
+    chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, stream=True)
+
+    #Call using your preferred method and obtain the final audio and text outputs.
+    ...
+
+    # SECTION 5: OUTPUT VALIDATION
+    # =============================
+    # INSTRUCTIONS: Verify that outputs meet expected criteria
+    # MODIFY: Adjust validation logic for your specific requirements
+
+    # ASSERTION 1: E2E Validation
+    # PURPOSE: Verify that the E2E latency is less than the baseline.
+    current_e2e = time.perf_counter() - start_time
+    print(f"the request e2e is: {current_e2e}")
+    e2e_list.append(current_e2e)
+
+    print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
+
+
+
+    # ASSERTION 2: Text Output Validation
+    # PURPOSE: Verify that text output was generated with keyword content
+    assert text_content is not None and len(text_content) >= 2, "No text output is generated"
+    assert any(
+        keyword in text_content.lower() for keyword in ["square", "quadrate", "sphere", "globe", "circle", "round"]
+    ), "The output does not contain any of the keywords."
+
+
+    # ASSERTION 3: Cross-Modal Consistency
+    # PURPOSE: Verify text and audio outputs convey the same information
+    # CUSTOMIZATION: Adjust similarity threshold (0.9) based on accuracy requirements
+    assert audio_data is not None, "No audio output is generated"
+    audio_content = merge_base64_and_convert_to_text(audio_data)
+    print(f"text content is: {text_content}")
+    print(f"audio content is: {audio_content}")
+    similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
+    print(f"similarity is: {similarity}")
+    assert similarity > 0.9, "The audio content is not same as the text"
 ```

+
 #### E2E - Offline inference
 ```python
 # SPDX-License-Identifier: Apache-2.0
@@ -271,10 +466,3 @@ def test_video_to_audio(omni_runner: type[OmniRunner], model: str) -> None:
        has_audio = any(o.final_output_type == "audio" for o in outputs)
        assert has_audio
 ```
-
-## Checklist before submitting your test files
-
-1. The file is saved in an appropriate place and the file name is clear.
-2. The coding style follows the requirements outlined above.
-3. **All test functions have appropriate pytest markers**
-4. For tests that need run in CI, please ensure the test is configured under the `./buildkite/` folder.
@@ -0,0 +1,173 @@
+
+# Metrics
+
+You can use these metrics in production to monitor the health and performance of the vLLM-omni system. Typical scenarios include:
+
+- **Performance Monitoring**: Track throughput (e.g., `e2e_avg_tokens_per_s`), latency (e.g., `e2e_total_ms`), and resource utilization to verify that the system meets expected standards.
+
+- **Debugging and Troubleshooting**: Use detailed per-request metrics to diagnose issues, such as high transfer times or unexpected token counts.
+
+## How to Enable and View Metrics
+
+### Start the Service with Metrics Logging
+
+```bash
+vllm serve /workspace/models/Qwen3-Omni-30B-A3B-Instruct --omni --port 8014 --log-stats
+```
+
+### Send a Request
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
+```
+
+### What You Will See
+
+With `--log-stats` enabled, the server will output detailed metrics logs after each request. Example output:
+
+
+#### Overall Summary
+
+| Field                       | Value        |
+|-----------------------------|--------------|
+| e2e_requests                | 1            |
+| e2e_wall_time_ms            | 41,299.190   |
+| e2e_total_tokens            | 5,202        |
+| e2e_avg_time_per_request_ms | 41,299.190   |
+| e2e_avg_tokens_per_s        | 125.959      |
+| e2e_stage_0_wall_time_ms    | 10,192.289   |
+| e2e_stage_1_wall_time_ms    | 30,541.409   |
+| e2e_stage_2_wall_time_ms    |    207.496   |
+
+#### RequestE2EStats
+
+| Field                   | Value      |
+|-------------------------|------------|
+| e2e_total_ms            | 41,299.133 |
+| e2e_total_tokens        | 5,202      |
+| transfers_total_time_ms | 245.895    |
+| transfers_total_kbytes  | 138,089.939|
+
+#### StageRequestStats
+
+| Field                  | 0      | 1      | 2      |
+|------------------------|--------|--------|--------|
+| audio_generated_frames | 0      | 0      | 525,525|
+| batch_id               | 38     | 274    | 0      |
+| batch_size             | 1      | 1      | 1      |
+| num_tokens_in          | 4,860  | 4,826  | 4,384  |
+| num_tokens_out         | 67     | 275    | 0      |
+| postprocess_time_ms    | 256.158| 0.491  | 0.000  |
+| stage_gen_time_ms      | 9,910.007|30,379.198|160.745|
+
+#### TransferEdgeStats
+
+| Field               | 0->1        | 1->2       |
+|---------------------|-------------|------------|
+| size_kbytes         | 109,277.349 | 28,812.591 |
+| tx_time_ms          | 78.701      | 18.790     |
+| rx_decode_time_ms   | 111.865     | 31.706     |
+| in_flight_time_ms   | 2.015       | 2.819      |
+
+
+These logs include:
+
+- **Overall summary**: total requests, wall time, average tokens/sec, etc.
+
+- **E2E table**: per-request latency and token counts.
+
+- **Stage table**: per-stage batch and timing details.
+
+- **Transfer table**: data transfer and timing for each edge.
+
+You can use these logs to monitor system health, debug performance, and analyze request-level metrics as described above.
+
+
+## Metrics Scope: Offline vs Online Inference
+
+For **offline inference** (batch mode), the summary includes both system-level metrics (aggregated across all requests) and per-request metrics. In this case, `e2e_requests` can be greater than 1, reflecting multiple completed requests in a batch.
+
+For **online inference** (serving mode), the summary is always per-request. `e2e_requests` is always 1, and only request-level metrics are reported for each completion.
+
+---
+
+## Parameter Details
+
+### Summary Metrics
+
+| Field                     | Meaning                                                                                       |
+|---------------------------|----------------------------------------------------------------------------------------------|
+| `e2e_requests`            | Number of completed requests.                                                                |
+| `e2e_wall_time_ms`        | Wall-clock time span from run start to last completion, in ms.                               |
+| `e2e_total_tokens`        | Total tokens counted across all completed requests (stage0 input + all stage outputs).       |
+| `e2e_avg_time_per_request_ms` | Average wall time per request: `e2e_wall_time_ms / e2e_requests`.                        |
+| `e2e_avg_tokens_per_s`    | Average token throughput over wall time: `e2e_total_tokens * 1000 / e2e_wall_time_ms`.      |
+| `e2e_stage_{i}_wall_time_ms` | Wall-clock time span for stage i, in ms. Each stage's wall time is reported as a separate field, e.g., `e2e_stage_0_wall_time_ms`, `e2e_stage_1_wall_time_ms`, etc. |
+
+---
+
+### E2E Table (per request)
+
+| Field                     | Meaning                                                               |
+|---------------------------|-----------------------------------------------------------------------|
+| `e2e_total_ms`            | End-to-end latency in ms.                                             |
+| `e2e_total_tokens`        | Total tokens for the request (stage0 input + all stage outputs).      |
+| `transfers_total_time_ms` | Sum of transfer edge `total_time_ms` for this request.                |
+| `transfers_total_kbytes`  | Sum of transfer kbytes for this request.                              |
+
+
+---
+
+### Stage Table (per stage event / request)
+
+| Field                     | Meaning                                                                                         |
+|---------------------------|-------------------------------------------------------------------------------------------------|
+| `batch_id`                | Batch index.                                                                                    |
+| `batch_size`              | Batch size.                                                                                     |
+| `num_tokens_in`           | Input tokens to the stage.                                                                      |
+| `num_tokens_out`          | Output tokens from the stage.                                                                   |
+| `stage_gen_time_ms`       | Stage compute time in ms, excluding postprocessing time (reported separately as `postprocess_time_ms`). |
+| `image_num`               | Number of images generated (for diffusion/image stages).                                        |
+| `resolution`              | Image resolution (for diffusion/image stages).                                                                  |
+| `postprocess_time_ms` | Diffusion/image: post-processing time in ms.                                                    |
+
+---
+
+### Transfer Table (per edge / request)
+
+| Field                | Meaning                                                                   |
+|----------------------|---------------------------------------------------------------------------|
+| `size_kbytes`        | Total kbytes transferred.                                                 |
+| `tx_time_ms`         | Sender transfer time in ms.                                               |
+| `rx_decode_time_ms`  | Receiver decode time in ms.                                               |
+| `in_flight_time_ms`  | In-flight time in ms.                                                     |
+
+
+### Expectation of the Numbers (Verification)
+
+**Formulas:**
+
+- `e2e_total_tokens = Stage0's num_tokens_in + sum(all stages' num_tokens_out)`
+
+- `transfers_total_time_ms = sum(tx_time_ms + rx_decode_time_ms + in_flight_time_ms)` for every edge
+
+**Using the example above:**
+
+**e2e_total_tokens**
+
+- Stage0's `num_tokens_in`: **4,860**
+- Stage0's `num_tokens_out`: **67**
+- Stage1's `num_tokens_out`: **275**
+- Stage2's `num_tokens_out`: **0**
+
+so `e2e_total_tokens = 4,860 + 67 + 275 + 0 = 5,202`, which matches the table value `e2e_total_tokens`.
+
+**transfers_total_time_ms**
+
+For each edge:
+
+- 0->1: tx_time_ms (**78.701**) + rx_decode_time_ms (**111.865**) + in_flight_time_ms (**2.015**) = **192.581**
+
+- 1->2: tx_time_ms (**18.790**) + rx_decode_time_ms (**31.706**) + in_flight_time_ms (**2.819**) = **53.315**
+
+192.581 + 53.315 = **245.896** = transfers_total_time_ms, which matches the calculation (difference is due to rounding)
@@ -460,7 +460,7 @@ def forward(self, ...):
 These keys are then accessible in your stage transition function:
 ```python
 # In stage_input_processors/qwen3_omni.py
-thinker_embeddings = output.multimodal_output["0"]  # Access by key
+thinker_prefill_embeddings = output.multimodal_output["0"]  # Access by key
 thinker_hidden_states = output.multimodal_output["24"]
 ```

@@ -513,11 +513,11 @@ def thinker2talker(
    for thinker_output in thinker_outputs:
        output = thinker_output.outputs[0]
        # Extract thinker embeddings and hidden states
-        thinker_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
+        thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
        thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()

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

@@ -131,9 +131,6 @@ python image_to_video.py \

 2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**:   [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)

-> **Note:**
-As of now, asynchronous (online) profiling is not fully supported in vLLM-Omni. While start_profile() and stop_profile() methods exist, they are only reliable in offline inference scripts (e.g., the provided end2end.py examples). Do not use them in server-mode or streaming scenarios—traces may be incomplete or fail to flush.
-
 ### 4. Analyzing Omni Traces

 Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
@@ -92,6 +92,13 @@ The framework achieves high performance through several optimization techniques:
    * **Quantization:** Supports various quantization implementations including FP8 and AWQ.
    * **FusedOps:** Allows for custom and third-party integration.

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

 vLLM-Omni is designed to be flexible and straightforward for users:
@@ -0,0 +1,216 @@
+# Async Chunk Design
+
+## Table of Contents
+
+1. [Overview](#overview)
+2. [Performance](#performance)
+3. [Architecture](#architecture)
+4. [Configuration](#configuration)
+5. [Related Files](#related-files)
+
+## Overview
+
+The `async_chunk` feature enables asynchronous, chunked processing of data across multiple stages in a multi-stage pipeline (e.g., Qwen3-Omni with Thinker → Talker → Code2Wav stages). Instead of waiting for a complete stage output before forwarding to the next stage, this feature allows stages to process and forward data in chunks as it becomes available, significantly reducing latency and improving throughput.
+
+**Chunk Size Definition**
+
+- **Prefill Phase**: `chunk_size = num_scheduled_tokens` for chunked prefill processing
+- **Decode Phase**: `chunk_size = num_scheduled_tokens = 1 ` for per-token streaming
+
+For qwen3-omni:
+- **Thinker → Talker**: Per decode step (typically chunk_size=1)
+- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFA. Use the per-request `initial_codec_chunk_frames` API field to override.
+- **Code2Wav**: Streaming decode with code2wav chunk_size
+
+With `async_chunk`:
+- Stages can start processing as soon as chunks are available
+- Overlapping execution across stages
+- Reduced latency and improved throughput
+- Better resource utilization
+- Async scheduling: Chunk IO (get/put) overlaps with compute via background threads so the scheduler is not blocked waiting for chunks
+
+## Performance
+1. **Reduced Latency**: Next stage can start processing immediately
+2. **Streaming Support**: Enables streaming for audio generation
+3. **IO-Compute Overlap**: Chunk retrieval happens asynchronously while other requests compute
+4. **Non-blocking Scheduler**: Requests waiting for chunks don't block the entire scheduler
+5. **Code2Wav Batch Inference**: Supports batched processing in code2wav stage
+
+| Input     | Output           | Async_chunk enabled | Code2Wav batch size | Max_Concurrency | Prompts | Mean E2E   | Mean TTFT  | Mean TPOT | Mean TTFP   | Mean RTF | Mean ITL |
+|-----------|------------------|---------------------|---------------------|----------------|---------|------------|------------|-----------|-------------|----------|----------|
+| text 100  | text 100+audio   | False               | 1                   | 1              | 50      | 6581.80    | 43.22      | 8.31      | 6459.34     | 0.24     | 8.22     |
+| text 100  | text 100+audio   | False               | 1                   | 4              | 50      | 7398.63    | 67.57      | 9.14      | 7285.35     | 0.27     | 9.05     |
+| text 100  | text 100+audio   | False               | 1                   | 10             | 50      | 13522.99   | 131.82     | 12.72     | 13410.44    | 0.49     | 12.60    |
+| text 100  | text 100+audio   | False               | 64                  | 1              | 50      | 6505.13    | 43.14      | 8.52      | 6395.40     | 0.24     | 8.44     |
+| text 100  | text 100+audio   | False               | 64                  | 4              | 50      | 7668.15    | 51.15      | 9.36      | 7562.37     | 0.28     | 9.27     |
+| text 100  | text 100+audio   | False               | 64                  | 10             | 50      | 9516.18    | 138.06     | 14.75     | 9409.26     | 0.34     | 14.60    |
+| text 100  | text 100+audio   | True                | 1                   | 1              | 50      | 6179.79    | 44.58      | 8.69      | 522.99      | 0.22     | 8.60     |
+| text 100  | text 100+audio   | True                | 1                   | 4              | 50      | 7692.69    | 103.96     | 10.22     | 785.85      | 0.29     | 10.12    |
+| text 100  | text 100+audio   | True                | 1                   | 10             | 50      | 11152.71   | 685.60     | 17.64     | 1628.88     | 0.41     | 17.62    |
+
+
+Performance data collected on H800 GPUs through comprehensive benchmarking with cudagraph enabled. text input uses random dataset.
+
+Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTFP)—e.g. ~92% at concurrency 1 (6.5s→0.52s)—and improves E2E latency (e.g. ~6% at conc 1, ~17% at conc 10). RTF (Real Time Factor) also improves with async_chunk on (e.g. ~8% at conc 1: 0.24→0.22, ~16% at conc 10: 0.49→0.41). Enabling **Code2Wav batch size 64** (vs 1) improves E2E and TTFP at higher concurrency when async_chunk is off (e.g. ~30% at conc 10: 13.5s→9.5s E2E, 13.4s→9.4s TTFP).
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png">
+    <img alt="TTFP Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png" width=100%>
+  </picture>
+</p>
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png">
+    <img alt="RTF Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png" width=100%>
+  </picture>
+</p>
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png">
+    <img alt="E2E Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png" width=100%>
+  </picture>
+</p>
+
+## Architecture
+### Data Flow
+
+#### Sequential Flow
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png">
+    <img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png" width=100%>
+  </picture>
+</p>
+
+#### Async Chunk Flow
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png">
+    <img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width=100%>
+  </picture>
+</p>
+
+### Async Chunk architecture
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png">
+    <img alt="Async Chunk Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png" width=100%>
+  </picture>
+</p>
+
+
+### Key Components
+
+1. **OmniConnector**: Inter-stage data transport only
+   - Shared memory or other IPC mechanisms
+   - **Transport-only API**: `put(from_stage, to_stage, put_key, data)` and `get(from_stage, to_stage, get_key)` (optionally with timeout)
+   - **No request-specific state**: Connector does not track put_requests, get_requests, request_payload, finished_requests, or other request-bound metadata; it only performs put/get operations
+   - Chunk keys and request/chunk lifecycle are managed by **OmniChunkTransferAdapter**
+
+2. **Transfer Adapter Layer**: Extensible abstraction for managing data transfer via connectors
+   - **OmniTransferAdapterBase**: Base class with background **recv_loop** and **save_loop** threads;
+   - **OmniChunkTransferAdapter**: Chunk-specific implementation that owns the full chunk lifecycle when async_chunk is enabled
+     - **Chunk ID and key construction**: Builds keys like `{req_id}_{stage_id}_{chunk_id}` for put/get
+     - **Async get**: `load_async(request)` enqueues the request; background **recv_loop** polls the connector (non-blocking); when data is available, updates the request and marks it in `_finished_load_reqs`; scheduler calls `get_finished_requests()` to learn which requests have chunks ready
+     - **Async put**: `save_async(pooling_output, request)` invokes `custom_process_next_stage_input_func` in the main thread to build the payload, then enqueues a save task; background **save_loop** performs `connector.put()`; payload processing and chunk accumulation (e.g. code2wav chunk_size) remain in the main thread
+
+3. **Stage Input Processors**: Custom functions that process stage outputs into chunks for different models
+   - Receive **transfer_manager** (OmniChunkTransferAdapter)
+   - Qwen3-omni reference: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
+
+4. **Schedulers**: Modified to handle chunk-based scheduling with async IO-compute overlap
+   - `OmniARScheduler`: For autoregressive stages
+   - `OmniGenerationScheduler`: For generation stages
+   - Both schedulers use **OmniChunkTransferAdapter** and **before/after** hooks around `super().schedule()`:
+     - **Before** `super().schedule()`: `process_pending_chunks(waiting, running)` moves requests waiting for chunks to `WAITING_FOR_CHUNK`, enqueues load tasks for background polling
+     - **After** `super().schedule()`: `restore_queues(waiting, running)` restores requests with ready chunks back to waiting/running, `postprocess_scheduler_output(scheduler_output)` attaches cached additional_information, clears chunk-ready flags
+   - **put_chunk** `save_async(pooler_output, request)`; **get_chunk** / **get_chunk_for_generation** `load_async(request)`
+
+5. **Model Runners**: Handle chunk processing
+   - `OmniGPUModelRunner`: Processes chunks in AR stages
+   - `GPUGenerationModelRunner`: Processes chunks in generation stages
+     - Uses `ubatch_slices` from `get_forward_context()` to track per-request sequence lengths in batched inference
+     - Reuses `ubatch_slices_padded` for code2wav batching to properly split batch outputs
+     - Handles list-type multimodal outputs: iterates through requests and assigns corresponding tensor to each
+     - Improved request state management: removes unscheduled and finished requests from input batch
+
+6. **Model Implementation**: Model-specific chunk handling
+   - `Qwen3OmniMoeForConditionalGeneration`: Main model with async_chunk support
+     - **Code2Wav stage batching**: Uses `ubatch_slices` to construct batched codec codes tensor `[batch_size, 16, max_seq_len]`
+     - **Batch output handling**: `generate_audio()` returns `list[torch.Tensor]`, one audio tensor per request
+     - **Multimodal outputs**: Returns list of audio tensors for batch processing instead of single concatenated tensor
+   - `Qwen3OmniCode2WavDecoder`: Audio generation model
+     - `chunked_decode()` and `chunked_decode_streaming()`: Return `list[torch.Tensor]` (one per request)
+     - Uses `ubatch_slices` to split batched waveform output into per-request audio chunks
+     - Each request gets correctly sized audio based on its code sequence length: `waveform[:, :, :code_seq_len * total_upsample]`
+
+7. **Request status**: `RequestStatus.WAITING_FOR_CHUNK` is added via patch (e.g. in `vllm_omni/patch.py`) so requests waiting for a chunk are not scheduled by the base vLLM scheduler until the chunk is ready.
+
+## Configuration
+
+Enable async_chunk in stage configuration YAML:
+
+```yaml
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    engine_args:
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
+  - stage_id: 1
+    engine_args:
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk
+```
+
+### Stage Configuration
+
+- `async_chunk: bool`: Enable/disable async chunk mode
+- `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
+- `stage_connector_config: dict`: Connector configuration
+- `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
+- `max_batch_size: int`: Maximum batch size for the stage
+
+
+### Connector Configuration
+
+```yaml
+connectors:
+  - from_stage: 0
+    to_stage: 1
+    spec:
+      name: SharedMemoryConnector
+      extra:
+        stage_id: 0
+```
+
+### Code2Wav Batch Configuration
+
+For optimal performance with async_chunk, the code2wav stage should be configured with batching:
+
+```yaml
+stage_args:
+  - stage_id: 2  # code2wav stage
+    runtime:
+      devices: "1"
+      max_batch_size: 64  # Enables batched audio generation
+    engine_args:
+      model_stage: code2wav
+```
+
+## Related Files
+
+- `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`: Chunk processing functions (receive `transfer_manager` as first param)
+- `vllm_omni/distributed/omni_connectors/transfer_adapter/base.py`: OmniTransferAdapterBase (recv_loop, save_loop, load_async, save_async)
+- `vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py`: OmniChunkTransferAdapter (process_pending_chunks, restore_queues, postprocess_scheduler_output)
+- `vllm_omni/distributed/omni_connectors/connectors/shm_connector.py`: SharedMemoryConnector (transport-only put/get)
+- `vllm_omni/core/sched/omni_ar_scheduler.py`: AR scheduler with chunk_transfer_adapter
+- `vllm_omni/core/sched/omni_generation_scheduler.py`: Generation scheduler with same async chunk pattern
+- `vllm_omni/worker/gpu_model_runner.py`: Model runner with chunk handling
+- `vllm_omni/worker/gpu_generation_model_runner.py`: Generation model runner with batch output handling and ubatch_slices support
+- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py`: Model implementation with code2wav batching
+- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py`: Code2wav decoder with batch support
+- `vllm_omni/engine/arg_utils.py`: Configuration definitions (async_chunk, worker_type)
+- `vllm_omni/config/model.py`: Model config with async_chunk field
@@ -0,0 +1,286 @@
+# Cache-DiT
+
+This section describes how to add cache-dit acceleration to a new diffusion pipeline. We use the Qwen-Image pipeline and LongCat-Image pipeline as reference implementations.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Standard Models: Automatic Support](#standard-models-automatic-support)
+- [Custom Architectures: Writing Custom Implementation](#custom-architectures-writing-custom-implementation)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is Cache-DiT?
+
+Cache-DiT is an acceleration library for Diffusion Transformers (DiT) that caches intermediate computation results across denoising steps. The core insight is that adjacent denoising steps often produce similar intermediate features, so we can skip redundant computations by reusing cached results.
+
+The library supports three main caching strategies:
+
+- **DBCache:** Dynamic block-level caching that selectively computes or caches transformer blocks based on residual differences
+- **TaylorSeer:** Calibration-based prediction that estimates block outputs using Taylor expansion
+- **SCM (Step Computation Masking):** Dynamic step skipping based on configurable policies
+
+### Architecture
+
+vLLM-omni integrates cache-dit through the `CacheDiTBackend` class, which provides a unified interface for managing cache-dit acceleration on diffusion models.
+
+| Method/Class | Purpose | Behavior |
+|--------------|---------|----------|
+| [`CacheDiTBackend`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.CacheBackend) | Unified backend interface | Automatically handles enabler selection and cache refresh |
+| [`enable_cache_for_dit()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Apply caching to transformer | Configures DBCache on transformer blocks |
+
+**Key APIs from Cache-DiT:**
+
+[Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)
+
+| API | Description |
+|-----|-------------|
+| `BlockAdapter` | Core abstraction for applying cache-dit to transformers. Specifies transformer module(s), block list(s), and forward signature pattern(s). |
+| `ForwardPattern` | Defines block forward signature patterns: `Pattern_0`, `Pattern_1`, `Pattern_2` |
+| `ParamsModifier` | Per-transformer or per-block-list cache configuration customization |
+| `DBCacheConfig` | Configuration for DBCache parameters (warmup steps, cached steps, thresholds) |
+| `refresh_context()` | Update cache context | Called when `num_inference_steps` changes |
+
+---
+
+## Standard Models: Automatic Support
+
+Most DiT models follow this pattern:
+- Single transformer with one `ModuleList` of blocks
+- Standard forward signature
+- Compatible with cache-dit's automatic detection
+
+**Examples:** Qwen-Image, Z-Image
+
+For standard single-transformer models, **no code changes are needed**. The `CacheDiTBackend` automatically uses `enable_cache_for_dit()`:
+
+```python
+from vllm_omni import Omni
+
+# Works automatically for standard models
+omni = Omni(
+    model="Qwen/Qwen-Image",  # Standard single-transformer model
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+    }
+)
+```
+
+**What happens automatically:**
+
+```python
+def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Default enabler for standard single-transformer DiT models."""
+
+    # Build cache configuration
+    db_cache_config = DBCacheConfig(
+        num_inference_steps=None,  # Will be set during first inference
+        Fn_compute_blocks=cache_config.Fn_compute_blocks,
+        Bn_compute_blocks=cache_config.Bn_compute_blocks,
+        max_warmup_steps=cache_config.max_warmup_steps,
+        max_cached_steps=cache_config.max_cached_steps,
+        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
+        residual_diff_threshold=cache_config.residual_diff_threshold,
+    )
+
+    # Enable cache-dit on transformer
+    cache_dit.enable_cache(
+        pipeline.transformer,
+        cache_config=db_cache_config,
+    )
+
+    # Return refresh function for dynamic num_inference_steps updates
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True):
+        cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+
+    return refresh_cache_context
+```
+
+---
+
+## Custom Architectures: Writing Custom Implementation
+
+Some models require custom handling:
+
+- **Dual-transformer:** Models with separate high-noise and low-noise transformers (e.g., Wan2.2)
+- **Multi-block-list:** Models with multiple block lists in one transformer (e.g., LongCatImage with `transformer_blocks` + `single_transformer_blocks`)
+- **Special forward patterns:** Models with non-standard block execution patterns
+
+### Example 1: Dual-Transformer Model (Wan2.2)
+
+Wan2.2 uses two transformers: one for high-noise steps and one for low-noise steps.
+
+**Key difference:** Use `BlockAdapter` to wrap multiple transformers with separate configurations.
+
+```python
+# Standard: cache_dit.enable_cache(pipeline.transformer, ...)
+# Custom: Use BlockAdapter to handle multiple transformers
+cache_dit.enable_cache(
+    BlockAdapter(
+        transformer=[pipeline.transformer, pipeline.transformer_2],  # Multiple transformers
+        blocks=[pipeline.transformer.blocks, pipeline.transformer_2.blocks],
+        forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
+        params_modifiers=[
+            ParamsModifier(...),  # Config for high-noise transformer
+            ParamsModifier(...),  # Config for low-noise transformer (different params)
+        ],
+    ),
+    cache_config=db_cache_config,
+)
+```
+
+**Key difference:** `refresh_context` must be called on each transformer separately.
+
+```python
+# Standard: cache_dit.refresh_context(pipeline.transformer, num_inference_steps=N)
+# Custom: Refresh each transformer with its own step count
+def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
+    high_steps, low_steps = _split_inference_steps(num_inference_steps)
+    cache_dit.refresh_context(pipeline.transformer, num_inference_steps=high_steps, ...)
+    cache_dit.refresh_context(pipeline.transformer_2, num_inference_steps=low_steps, ...)
+```
+
+### Example 2: Multi-Block-List Model (LongCatImage)
+
+LongCatImage has a single transformer with two block lists: `transformer_blocks` and `single_transformer_blocks`.
+
+**Key difference:** Use `BlockAdapter` to specify multiple block lists within one transformer.
+
+```python
+# Standard: cache_dit.enable_cache(pipeline.transformer, ...)
+#   - Automatically detects single block list
+# Custom: Use BlockAdapter to specify multiple block lists
+cache_dit.enable_cache(
+    BlockAdapter(
+        transformer=pipeline.transformer,  # Single transformer
+        blocks=[
+            pipeline.transformer.transformer_blocks,        # Block list 1
+            pipeline.transformer.single_transformer_blocks, # Block list 2
+        ],
+        forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
+        params_modifiers=[modifier],
+    ),
+    cache_config=db_cache_config,
+)
+```
+
+> **Note:** For single transformer with multiple block lists, `refresh_context` works the same as standard models.
+
+### Registering Custom Implementations
+
+After writing your custom enabler, register it in `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`:
+
+```python
+CUSTOM_DIT_ENABLERS = {
+    "Wan22Pipeline": enable_cache_for_wan22,
+    "LongCatImagePipeline": enable_cache_for_longcat_image,
+    "YourCustomPipeline": enable_cache_for_your_model,  # Add here
+}
+```
+
+---
+
+## Testing
+
+After adding cache-dit support, test with:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# Test your custom model
+omni = Omni(
+    model="your-model-name",
+    cache_backend="cache_dit",
+    cache_config={
+        "Fn_compute_blocks": 1,
+        "Bn_compute_blocks": 0,
+        "max_warmup_steps": 4,
+        "residual_diff_threshold": 0.24,
+    }
+)
+
+images = omni.generate(
+    "a beautiful landscape",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Verify:**
+
+1. Cache is applied (check logs for "Cache-dit enabled successfully on xxx")
+2. Performance improvement (should be around 1.5x-2x faster)
+3. Image quality (compare with `cache_backend=None`)
+
+---
+
+## Troubleshooting
+
+### Issue: Cache not applied
+
+**Symptoms:** No speedup observed, no cache-related log messages.
+
+**Causes & Solutions:**
+
+- **Enabler not registered:**
+
+**Problem:** Pipeline name not in `CUSTOM_DIT_ENABLERS` registry.
+
+**Solution:** Verify `pipeline.__class__.__name__` matches the registry key and add your enabler to `CUSTOM_DIT_ENABLERS`.
+
+### Issue: Quality degradation
+
+**Symptoms:** Generated images have artifacts or lower quality compared to non-cached inference.
+
+**Causes & Solutions:**
+
+- **Cache parameters too aggressive:**
+
+**Solution:**
+```python
+cache_config={
+    "residual_diff_threshold": 0.12,  # Lower from 0.24 (try 0.12-0.18)
+    "max_warmup_steps": 6,            # Increase from 4 (try 6-8)
+    "max_continuous_cached_steps": 2, # Reduce if higher
+}
+```
+
+Check the [user guide for cache_dit](../../user_guide/diffusion/cache_dit_acceleration.md) for more adjustable parameters.
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Pattern | Notes |
+|-------|------|---------|-------|
+| **Standard DiT** | [`cache_dit_backend.py::enable_cache_for_dit`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Default enabler | Single transformer, automatic |
+| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Dual-transformer | Separate high/low noise transformers |
+| **LongCat** | [`cache_dit_backend.py::enable_cache_for_longcat_image`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_longcat_image) | Multi-block-list | Two block lists in one transformer |
+| **BAGEL** | [`cache_dit_backend.py::enable_cache_for_bagel`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_bagel) | Omni model | Complex architecture |
+
+---
+
+## Summary
+
+Adding cache-dit support:
+
+1. ✅ **Check model type** - Standard models work automatically, custom architectures need enablers
+2. ✅ **Write enabler** (if needed) - Use `BlockAdapter` for complex architectures
+3. ✅ **Register enabler** (if needed) - Add to `CUSTOM_DIT_ENABLERS` dictionary
+4. ✅ **Return refresh function** (if needed) - Handle `num_inference_steps` changes
+5. ✅ **Test** - Verify with `cache_backend="cache_dit"`
+
+For most models, the default enabler is sufficient. Only write custom enablers for complex architectures!
@@ -0,0 +1,249 @@
+# CFG-Parallel
+
+This section describes how to add CFG-Parallel (Classifier-Free Guidance Parallel) to a diffusion pipeline. We use the Qwen-Image pipeline as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Customization](#customization)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is CFG-Parallel?
+
+In standard Classifier-Free Guidance, each diffusion step requires two forward passes through the transformer:
+
+1. **Positive/Conditional**: Guided by the text prompt
+2. **Negative/Unconditional**: Typically using empty or negative prompt
+
+CFG-Parallel eliminates this bottleneck by distributing the two forward passes across different GPU ranks, allowing them to execute simultaneously rather than sequentially.
+
+### Architecture
+
+vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic. Pipelines inherit from this mixin and implement a `diffuse()` method that orchestrates the denoising loop.
+
+| Method | Purpose | Automatic Behavior |
+|--------|---------|-------------------|
+| [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with CFG | Detects parallel mode, distributes computation, gathers results |
+| [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler with sync | Rank 0 steps, broadcasts latents to all ranks |
+| [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine positive/negative | Applies CFG formula with optional normalization |
+| [`predict_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Forward pass wrapper | Override for custom transformer calls |
+| [`cfg_normalize_function()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Normalize CFG output | Override for custom normalization |
+
+### How It Works
+
+`predict_noise_maybe_with_cfg()` automatically detects and switches between two execution modes:
+
+- **CFG-Parallel mode** (when `cfg_world_size > 1`):
+  - Rank 0 computes positive prompt prediction
+  - Rank 1 computes negative prompt prediction
+  - Results are gathered via `all_gather()`
+  - Combined on rank 0 using CFG formula
+
+- **Sequential mode** (when `cfg_world_size == 1`):
+  - Single rank computes both positive and negative predictions
+  - Directly combines them with CFG formula
+
+`scheduler_step_maybe_with_cfg()` ensures consistent latent states across all ranks:
+
+- **CFG-Parallel mode**:
+  - Only rank 0 performs the scheduler step (applies noise prediction to update latents)
+  - Updated latents are broadcast to all other ranks via `broadcast()`
+  - All ranks maintain synchronized latent states for the next iteration
+
+- **Sequential mode**:
+  - Single rank directly performs the scheduler step
+  - No synchronization needed
+
+---
+
+## Step-by-Step Implementation
+
+### Step 1: Inherit `CFGParallelMixin`
+
+Allow your pipeline to inherit from `CFGParallelMixin` and implements the `diffuse()` method for your specific model.
+
+**Example (Qwen-Image):**
+
+```python
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+import torch.nn as nn
+class YourModelPipeline(nn.Module, CFGParallelMixin):
+    def diffuse(self, ...) -> torch.Tensor:
+        for i, t in enumerate(timesteps):
+            # Prepare positive_kwargs (conditional) and negative_kwargs (unconditional)
+            positive_kwargs = {...}  # hidden_states, encoder_hidden_states, etc.
+            negative_kwargs = {...} if do_true_cfg else None
+
+            # Key method 1: Predict noise with automatic CFG parallel handling
+            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,
+            )
+
+            # Key method 2: Step scheduler with automatic CFG synchronization
+            latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred, t, latents, do_true_cfg
+            )
+
+        return latents
+```
+
+**Key Points:**
+
+- `positive_kwargs`: transformer arguments for conditional (text-guided) prediction
+- `negative_kwargs`: transformer arguments for unconditional prediction (set to `None` if CFG disabled)
+- For image editing pipelines, add `output_slice=image_seq_len` to extract the generative image portion
+
+### Step 2: Call `diffuse`
+
+Call `self.diffuse` in your pipeline's forward function:
+
+```python
+import torch.nn as nn
+class YourModelPipeline(nn.Module, CFGParallelMixin):
+    def forward(
+        self,
+        prompt: str,
+        negative_prompt: str | None = None,
+        guidance_scale: float = 3.5,
+        num_inference_steps: int = 50,
+        **kwargs,
+    ):
+        # Encode prompts, Initialize latents, Get timesteps
+        ...
+        # Run diffusion loop (calls the mixin's diffuse method)
+        latents = self.diffuse(
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            negative_prompt_embeds=negative_embeds,
+            negative_prompt_embeds_mask=negative_mask,
+            latents=latents,
+            timesteps=timesteps,
+            do_true_cfg=do_true_cfg,
+            true_cfg_scale=guidance_scale,
+            ...
+        )
+```
+
+---
+
+## Customization
+
+### Override `predict_noise()` for Custom Transformer Calls
+
+If your transformer requires custom prediction function, you can rewrite `predict_noise` function. Taking Wan2.2 as an example, which has two transformer models. The actual transformer to be called is determined by `self.transformer`.
+
+```python
+class Wan22Pipeline(nn.Module, CFGParallelMixin):
+    def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
+        if current_model is None:
+            current_model = self.transformer
+        return current_model(**kwargs)[0]
+```
+
+
+### Override `cfg_normalize_function()` for Custom Normalization
+
+Some models have their own normalization function. Taking LongCat Image model as an example:
+
+```python
+class LongCatImagePipeline(nn.Module, CFGParallelMixin):
+    def cfg_normalize_function(self, noise_pred, comb_pred, cfg_renorm_min=0.0):
+        """
+        Normalize the combined noise prediction.
+        """
+        cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+        noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+        scale = (cond_norm / (noise_norm + 1e-8)).clamp(min=cfg_renorm_min, max=1.0)
+        noise_pred = comb_pred * scale
+        return noise_pred
+
+        # The original cfg_normalize_function function in CFGParallelMixin
+        # 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
+```
+
+---
+
+## Testing
+
+After adding CFG-Parallel support, test with:
+
+```bash
+cd examples/offline_inference/text_to_image
+python text_to_image.py \
+    --model Your-org/your-model \
+    --prompt "a cup of coffee on the table" \
+    --negative-prompt "ugly, unclear" \
+    --cfg-scale 4.0 \
+    --num-inference-steps 50 \
+    --output "cfg_enabled.png" \
+    --cfg-parallel-size 2
+```
+
+**Verify:**
+
+1. Check logs for CFG parallel being activated
+2. Record the `e2e_time_ms` in the log and compare with CFG-Parallel disabled
+3. Compare the generated result quality with baseline
+4. Record comparison results in your PR
+
+---
+
+## Troubleshooting
+
+### Issue: CFG parallel not activating
+
+**Symptoms:** Generation still slow, logs don't show CFG parallel being used.
+
+**Causes & Solutions:**
+
+- **CFG is not enabled:**
+
+**Problem:** Guidance scale too low or negative prompt not provided.
+
+**Solution:** Ensure `guidance_scale > 1.0` and negative prompt is provided:
+```python
+images = pipeline(
+    prompt="a cat",
+    negative_prompt="",  # Must provide (even if empty)
+    guidance_scale=3.5,   # Must be > 1.0
+)
+```
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Pattern | Notes |
+|-------|------|---------|-------|
+| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/cfg_parallel.py` | Mixin | Dual-stream transformer |
+| **Qwen-Image-Edit** | `vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py` | Mixin | Image editing with `output_slice` |
+| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` | Mixin | Dual-transformer architecture |
+| **CFGParallelMixin** | `vllm_omni/diffusion/distributed/cfg_parallel.py` | Base implementation | Core mixin class |
+
+---
+
+## Summary
+
+Adding CFG-Parallel support:
+
+1. ✅ **Create mixin** - Inherit from `CFGParallelMixin` and implement `diffuse()` method
+2. ✅ **(Optional) Customize** - Override `predict_noise()` or `cfg_normalize_function()` for custom behavior
+3. ✅ **Test** - Verify with `--cfg-parallel-size 2` and compare performance
@@ -6,7 +6,8 @@ This guide explains how to configure and use distributed connectors
 Backend-specific setup lives in separate docs:

 - [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
- [MooncakeConnector](omni_connectors/mooncake_connector.md)
+- [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
+- [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
 - [YuanrongConnector](omni_connectors/yuanrong_connector.md)

 ## Overview
@@ -19,7 +20,8 @@ Current connectors operate in D2H2D (device to host to device) mode.
 | 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 (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
+| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
 | Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |

 ## Core API
@@ -0,0 +1,146 @@
+# HSDP
+
+This section describes how to add HSDP (Hybrid Sharded Data Parallel) support to a diffusion transformer model. We use the Wan2.2 transformer as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Reference Implementations](#reference-implementations)
+
+---
+
+## Overview
+
+### What is HSDP?
+
+HSDP (Hybrid Sharded Data Parallel) is a memory optimization technique that **shards model weights** across multiple GPUs using PyTorch's FSDP2. Unlike Tensor Parallelism which splits computation, HSDP:
+
+- Shards weights across GPUs to reduce per-GPU memory usage
+- Gathers weights on-demand during forward passes
+- Can work standalone or combined with other parallelism (e.g., Sequence Parallel)
+
+This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
+
+**Important constraints:**
+- HSDP cannot be used with Tensor Parallelism
+- For standalone HSDP (no other parallelism), `hsdp_shard_size` must be specified explicitly
+
+### Architecture
+
+HSDP implementation relies on:
+
+1. **`_hsdp_shard_conditions`**: Model attribute specifying which modules to shard
+2. **`apply_hsdp_to_model`**: Function that applies FSDP2 sharding
+3. **`HSDPInferenceConfig`**: Runtime configuration for HSDP
+
+---
+
+## Step-by-Step Implementation
+
+### Step 1: Identify Modules to Shard
+
+Determine which modules in your transformer should be sharded. Typically, these are:
+
+- Transformer blocks (e.g., `blocks.0`, `blocks.1`, ...)
+- Large submodules with significant weight memory
+
+**Key questions:**
+- Which modules have the largest weights?
+- Which modules are repeated (like transformer blocks)?
+
+### Step 2: Define Shard Conditions
+
+Add `_hsdp_shard_conditions` to your model class. This is a list of functions that return `True` for modules that should be sharded.
+
+**Example (Transformer Blocks):**
+
+```python
+class MyTransformerModel(nn.Module):
+
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        """Match transformer blocks for HSDP sharding.
+
+        Args:
+            name: Module name from named_modules() (e.g., "blocks.0", "blocks.0.attn")
+            module: The module instance
+
+        Returns:
+            True if this module should be sharded
+        """
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+```
+
+**Multiple Conditions Example:**
+
+```python
+class MyModel(nn.Module):
+
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "blocks" in name and name.split(".")[-1].isdigit()
+
+    @staticmethod
+    def _is_moe_expert(name: str, module) -> bool:
+        # Also shard MoE expert layers
+        return "experts" in name and name.split(".")[-1].isdigit()
+
+    # Module is sharded if ANY condition returns True
+    _hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]
+```
+
+---
+
+## Testing
+
+After adding HSDP support, test with:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+parallel_config = DiffusionParallelConfig(
+    use_hsdp=True,
+    hsdp_shard_size=8,  # Shard across 8 GPUs
+)
+omni = Omni(model="your-model-name", parallel_config=parallel_config)
+
+output = omni.generate(
+    "a cup of coffee on the table",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Or via command line:**
+
+```bash
+vllm serve Your-org/your-model --omni --port 8091 --use-hsdp
+```
+
+**Verify:**
+
+1. Check logs for "HSDP Inference: replicate_size=..., shard_size=..."
+2. Check logs for "Sharded N modules + root"
+3. Verify memory usage is reduced proportionally
+4. Compare generated output quality with HSDP disabled
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Notes |
+|-------|------|-------|
+| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Reference implementation |
+| **HSDP Core** | `vllm_omni/diffusion/distributed/hsdp.py` | `apply_hsdp_to_model`, `shard_model` |
+| **HSDP Tests** | `tests/diffusion/distributed/test_hsdp.py` | Unit tests |
+
+---
@@ -1,74 +0,0 @@
-# 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,312 @@
+# MooncakeStoreConnector
+
+## 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: MooncakeStoreConnector
+      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).
+
+---
+
+## Design
+
+### 1. Overview
+
+`MooncakeStoreConnector` is the storage-oriented remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake distributed store APIs and provides a uniform `put()` / `get()` abstraction for transferring stage payloads across nodes.
+
+Compared with `SharedMemoryConnector`, this connector is intended for multi-node deployments. Compared with `MooncakeTransferEngineConnector`, it is simpler and more object-oriented: data is written into a distributed store and fetched back by key, rather than transferred through an explicit remote-memory write protocol.
+
+Its primary role is to provide a general-purpose remote connector for stage payloads when:
+
+- the pipeline spans multiple hosts
+- a shared-memory connector is not possible
+- a simpler remote backend is preferred over the RDMA-oriented transfer engine
+
+### 2. Relationship with the OmniConnector System
+
+`MooncakeStoreConnector` implements `OmniConnectorBase`, so it integrates with the same connector lifecycle as the other backends:
+
+- `OmniConnectorFactory` instantiates it from `ConnectorSpec`
+- `load_omni_transfer_config()` maps edge config to the connector spec
+- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
+
+This means the rest of the pipeline does not need store-specific logic. It only relies on the generic connector contract.
+
+### 3. Design Goals
+
+The connector is designed around the following goals:
+
+- **Cross-node object transfer** using a single key-based abstraction
+- **Minimal connector-specific control plane**, since the store itself is the shared medium
+- **Compatibility with arbitrary Python payloads** through the Omni serializer
+- **Operational simplicity** compared with the transfer-engine-based remote connector
+
+The connector is intentionally not optimized for zero-copy tensor movement or direct remote-memory access.
+
+### 4. Core Design
+
+#### 4.1 Object-Oriented Transfer Model
+
+`MooncakeStoreConnector` treats the transport backend as a distributed object store:
+
+1. Serialize the Python object into bytes.
+2. Generate a unique connector key.
+3. Store the bytes in Mooncake.
+4. Fetch the bytes from Mooncake on the receiver side.
+5. Deserialize the bytes back into the original object.
+
+This design makes the connector conceptually close to a distributed key-value transport.
+
+#### 4.2 Key Construction
+
+The connector uses `OmniConnectorBase._make_key()` to derive the internal store key:
+
+```text
+{request_key}@{from_stage}_{to_stage}
+```
+
+This adds stage routing information to the logical request key and avoids collisions between different pipeline edges.
+
+#### 4.3 No Extra Metadata Path
+
+Unlike `SharedMemoryConnector` and `MooncakeTransferEngineConnector`, this connector returns:
+
+```python
+metadata = None
+```
+
+from `put()`.
+
+This is a meaningful design difference:
+
+- the store itself is the shared rendezvous point
+- consumers only need the same key
+- no transport handle, address, or side-channel metadata needs to be propagated
+
+As a result, the control plane is simpler for this connector than for the SHM and RDMA variants.
+
+### 5. Initialization
+
+#### 5.1 Required Mooncake Components
+
+The connector requires the Mooncake Python bindings to expose:
+
+- `MooncakeDistributedStore`
+- `ReplicateConfig`
+
+If these symbols are unavailable, construction fails immediately with `ImportError`. This keeps startup failures explicit and avoids silent fallback behavior.
+
+#### 5.2 Store Setup
+
+During `_init_store()`, the connector:
+
+1. creates `MooncakeDistributedStore()`
+2. calls `store.setup(...)`
+3. validates the return code
+4. creates a `ReplicateConfig`
+5. enables `with_soft_pin = True`
+
+The setup step is completed during connector construction, not lazily at first use.
+
+### 6. Put / Get Flow
+
+#### 6.1 Producer Flow: `put()`
+
+The producer-side flow is:
+
+1. Validate that the store has been initialized.
+2. Serialize the payload into bytes.
+3. Build the internal key with stage routing information.
+4. Call `store.put(key, serialized_data, self.pin)`.
+5. Update metrics and return success.
+
+The returned tuple is:
+
+```python
+(True, len(serialized_data), None)
+```
+
+#### 6.2 Consumer Flow: `get()`
+
+The consumer-side flow is:
+
+1. Validate that the store has been initialized.
+2. Build the same internal key.
+3. Poll `store.get(key)` for up to 20 retries.
+4. Sleep 50 ms between retries.
+5. If data is found, deserialize it and return `(data, payload_size)`.
+6. If all retries are exhausted, record a timeout and return `None`.
+
+This design gives the connector a bounded waiting model rather than an indefinite blocking get.
+
+### 7. Integration with Stage Communication
+
+All callers use the connector through the same `put()` / `get()` contract:
+
+- the sender calls `put()` to serialize and store the payload
+- the receiver calls `get()` to retrieve and deserialize it
+- no connector-specific metadata is required, since the store key is the rendezvous point
+
+Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, which makes the connector functional but not the highest-performance option for large payloads such as KV cache blocks.
+
+### 8. Data Flow in the Pipeline
+
+The end-to-end transfer model is:
+
+```mermaid
+sequenceDiagram
+    participant SenderStage
+    participant MooncakeStoreConnector
+    participant MooncakeStore
+    participant ReceiverStage
+
+    SenderStage->>MooncakeStoreConnector: put(from_stage, to_stage, put_key, data)
+    MooncakeStoreConnector->>MooncakeStoreConnector: serialize object
+    MooncakeStoreConnector->>MooncakeStore: put(internal_key, bytes)
+    MooncakeStoreConnector-->>SenderStage: (success, size, None)
+
+    ReceiverStage->>MooncakeStoreConnector: get(from_stage, to_stage, get_key)
+    MooncakeStoreConnector->>MooncakeStore: get(internal_key)
+    MooncakeStoreConnector->>MooncakeStoreConnector: deserialize bytes
+    MooncakeStoreConnector-->>ReceiverStage: (data, size)
+```
+
+This is a store-mediated remote transfer model rather than a direct peer-to-peer transport model.
+
+### 9. Strengths and Trade-offs
+
+#### Strengths
+
+- Simple conceptual model: store by key, fetch by key.
+- No connector-specific metadata handoff is required.
+- Works naturally for remote multi-node stage transfer.
+- Easy to integrate into the existing connector abstraction.
+
+#### Trade-offs
+
+- Full serialization and deserialization are always required.
+- Large payloads are more expensive than in direct-memory transports.
+- Runtime behavior depends on external Mooncake services being available.
+- Cleanup semantics are weaker than request-scoped local buffer management.
+
+### 10. Important Implementation Characteristics
+
+#### 10.1 Cleanup Is a No-op
+
+`cleanup()` only logs a debug message and does not actively delete remote data. The current design assumes Mooncake-side lifecycle management rather than explicit per-request removal.
+
+This keeps the connector implementation simple but means request-level reclamation is not modeled inside the connector itself.
+
+#### 10.2 Close Releases the Store Handle
+
+Unlike `SharedMemoryConnector`, where `close()` is a no-op, `MooncakeStoreConnector.close()` performs a meaningful shutdown:
+
+1. Calls `self.store.close()` to release the Mooncake store handle.
+2. Sets `self.store = None` to prevent further operations.
+
+This ensures the connector releases its connection to the external Mooncake service on shutdown. Errors during close are logged but do not propagate.
+
+#### 10.3 Health Output Uses a Shared Schema
+
+`health()` reports:
+
+- `host`
+- `metadata_server`
+- `master`
+- metrics
+
+and also includes placeholder fields such as `protocol`, `pool_device`, and `pool_size` to keep a more uniform shape across connector health outputs. This reflects a system-level consistency choice rather than a store-specific design need.
+
+#### 10.4 Get Is Retry-Based, Not Event-Driven
+
+The receiver side polls the store with a short retry loop. This is simple and robust, but it means the connector is latency-sensitive to:
+
+- store visibility delay
+- network jitter
+- payload size
+
+If the payload is not visible within the retry window, the connector reports a timeout.
+
+### 11. Summary
+
+`MooncakeStoreConnector` is the remote, store-based connector in the OmniConnector stack. Its design is straightforward:
+
+- serialize the payload
+- store it in Mooncake under a stage-qualified key
+- fetch it by the same key on the receiver side
+- deserialize it back into the original object
+
+This connector fills an important gap in the system:
+
+1. It enables cross-node transfer without requiring shared memory.
+2. It keeps the stage communication model uniform.
+3. It provides a simpler operational alternative to the transfer-engine-based remote connector.
+
+Its simplicity is also its main trade-off: it favors a clean object-transfer model over the fast-path and peer-to-peer optimizations implemented by `MooncakeTransferEngineConnector`.
@@ -0,0 +1,793 @@
+# MooncakeTransferEngineConnector
+
+## When to Use
+
+Best for high-performance multi-node data transfer between stages using Mooncake
+Transfer Engine. Supports both RDMA and TCP protocols with a managed memory pool,
+zero-copy deserialization, and optional GPUDirect RDMA. Applicable to any
+inter-stage data (KV caches, request payloads, etc.), not limited to KV cache transfer.
+
+Compared to `MooncakeStoreConnector` (TCP key-value store), this connector
+provides **~60x faster** data transfer via RDMA direct memory access.
+
+## Installation
+
+```bash
+pip install mooncake-transfer-engine
+```
+
+Ensure RDMA drivers are installed on all nodes (e.g., Mellanox OFED for
+InfiniBand/RoCE NICs).
+
+## Configuration
+
+Define the connector in runtime:
+
+```yaml
+runtime:
+  connectors:
+    rdma_connector:
+      name: MooncakeTransferEngineConnector
+      extra:
+        host: "auto"                  # Auto-detect local RDMA IP
+        zmq_port: 50051               # ZMQ base port (see "Port Offset Scheme" below)
+        protocol: "rdma"              # "rdma" or "tcp"
+        device_name: ""               # RDMA device (e.g., "mlx5_0"), empty for auto-detect
+        memory_pool_size: 2147483648  # 2GB memory pool
+        memory_pool_device: "cpu"     # "cpu" for pinned memory, "cuda" for GPUDirect RDMA
+```
+
+Wire stages to the connector:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: rdma_connector
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: rdma_connector
+```
+
+## Parameters
+
+### Required
+
+| Parameter | Description |
+|---|---|
+| `role` | **Internal, do not set manually.** Auto-injected by the orchestration layer (`"sender"` for `output_connectors`, `"receiver"` for `input_connectors`). Defaults to `"sender"` if omitted. |
+| `host` | Local IP address for RDMA. `"auto"` detects from network interfaces. |
+| `protocol` | Transport protocol: `"rdma"` (InfiniBand/RoCE) or `"tcp"`. |
+
+### Memory Pool
+
+| Parameter | Default | Description |
+|---|---|---|
+| `memory_pool_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. |
+| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity). |
+
+### Networking
+
+| Parameter | Default | Description |
+|---|---|---|
+| `zmq_port` | 50051 | ZMQ **base** port. The orchestration layer computes the actual port as `base + purpose_offset + stage_offset` (see table below). Users only set this base value. |
+| `sender_host` | `None` | **Internal.** Receiver-side only — dynamically resolved via `update_sender_info()`. Not needed in YAML. |
+| `sender_zmq_port` | `None` | **Internal.** Receiver-side only — defaults to the sender's adjusted port. Not needed in YAML. |
+| `device_name` | `""` | RDMA device name (e.g., `"mlx5_0"`). Empty for auto-detect. Can also be set via `RDMA_DEVICE_NAME` env var. |
+
+#### ZMQ Port Offset Scheme
+
+To avoid port conflicts when multiple edges, purposes, DP replicas, or TP ranks share the same node, the actual ZMQ port is computed as:
+
+```
+side_channel_port = zmq_port + purpose_offset + stage_offset + dp_index * tp_size
+sender_listen     = side_channel_port + tp_rank
+receiver_connect  = remote_side_channel_port + tp_rank
+```
+
+| Component | Value | Description |
+|---|---|---|
+| `zmq_port` | 50051 (default) | Base port from YAML config |
+| `purpose_offset` | `request_forwarding` = 0, `kv_transfer` = 100 | Separates control-plane vs KV-cache connections |
+| `stage_offset` | `int(from_stage)` (0, 1, 2...) | Separates edges from different source stages |
+| `dp_index * tp_size` | e.g., DP1 × TP2 = 2 | Each DP replica reserves a port range of size `tp_size` (following vLLM convention: `VLLM_MOONCAKE_BOOTSTRAP_PORT + dp_index * tp_size`) |
+| `tp_rank` | 0, 1, 2... | Each TP rank within a DP replica uses its own port |
+| orchestrator | +200 | Extra offset when caller is the orchestrator (avoids collision with stage workers on the same node) |
+
+**Example** (base=50051, stage 0→1, DP=2, TP=2, kv_transfer):
+
+| Caller | DP | TP rank | Port |
+|---|---|---|---|
+| Stage worker | DP0 | rank 0 | `50051 + 100 + 0 + 0×2 + 0 = 50151` |
+| Stage worker | DP0 | rank 1 | `50051 + 100 + 0 + 0×2 + 1 = 50152` |
+| Stage worker | DP1 | rank 0 | `50051 + 100 + 0 + 1×2 + 0 = 50153` |
+| Stage worker | DP1 | rank 1 | `50051 + 100 + 0 + 1×2 + 1 = 50154` |
+| Orchestrator | — | — | `50051 + 200 + 0 = 50251` |
+
+## Memory Pool Modes
+
+| Mode | Config | Data Flow | Best For |
+|---|---|---|---|
+| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
+| GPUDirect | `memory_pool_device: "cuda"` | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
+
+> **Note**: GPUDirect RDMA requires the NIC and GPU to share a direct PCIe
+> switch (PIX topology). On systems where they are connected via PXB or NODE,
+> CPU pinned memory is faster due to GPU BAR1 bandwidth limitations.
+
+## Environment Variables
+
+| Variable | Description |
+|---|---|
+| `RDMA_DEVICE_NAME` | Override RDMA device name (e.g., `mlx5_0`). |
+| `MC_IB_PCI_RELAXED_ORDERING` | Set to `1` to enable PCIe relaxed ordering for GPUDirect. |
+
+## Docker / Container Setup
+
+RDMA requires host-level device access:
+
+```bash
+docker run -it \
+    --cap-add=SYS_PTRACE \
+    --cap-add=IPC_LOCK \
+    --security-opt seccomp=unconfined \
+    --network=host \
+    --device=/dev/infiniband \
+    -v /sys/class/infiniband:/sys/class/infiniband:ro \
+    your-image:tag
+```
+
+## Performance
+
+Benchmark results on H800 GPUs with mlx5_0 RDMA NIC (~186 MB KV cache):
+
+| Metric | MooncakeStoreConnector | MooncakeTransferEngineConnector (CPU) |
+|---|---|---|
+| KV transfer wall time | ~810 ms | **~14 ms** |
+| RDMA throughput | N/A (TCP) | ~22 GB/s |
+| Speedup | 1x | **58x** |
+
+## Troubleshooting
+
+### Quick Diagnostics
+
+```bash
+# 1. Check RDMA devices and link status
+ibdev2netdev
+# Expected: "mlx5_X port 1 ==> <iface> (Up)"
+# RoCE devices map to Ethernet interfaces (e.g., enp75s0f0)
+# IB devices map to ib0, ib1, etc.
+
+# 2. Check InfiniBand device details
+ibstat
+
+# 3. Verify /dev/infiniband is accessible (required in containers)
+ls /dev/infiniband/
+
+# 4. Check Mooncake installation
+python -c "from mooncake.engine import TransferEngine; print('OK')"
+
+# 5. Check environment variables
+echo "RDMA_DEVICE_NAME=${RDMA_DEVICE_NAME:-<not set>}"
+echo "MC_IB_PCI_RELAXED_ORDERING=${MC_IB_PCI_RELAXED_ORDERING:-<not set>}"
+```
+
+### Common Issues
+
+| Symptom | Cause | Fix |
+|---|---|---|
+| `Failed to modify QP to RTR` | Cross-NIC QP handshake failure (multi-NIC DGX) | Set `device_name` to a single RoCE NIC (e.g., `mlx5_2`) or set `RDMA_DEVICE_NAME` env var |
+| `transport retry counter exceeded` | RDMA path between incompatible NICs | Same as above — restrict to one NIC |
+| `zmq.error.Again: Resource temporarily unavailable` | ZMQ recv timeout (transfer took too long) | Check NIC selection; increase data may need longer timeout |
+| `Mooncake Engine initialization failed` | Missing RDMA drivers or `/dev/infiniband` | Install Mellanox OFED; in Docker add `--device=/dev/infiniband` |
+| `MemoryError` in allocator | Memory pool too small for payload | Increase `memory_pool_size` |
+| GPU transfer slower than CPU | GPU BAR1 bandwidth limitation (PXB/NODE topology) | Use `memory_pool_device: "cpu"` instead of `"cuda"` |
+
+### Multi-NIC Environments (DGX)
+
+On DGX machines with 12+ RDMA NICs, only RoCE NICs (with a bound network
+interface) reliably support loopback. IB-only NICs may fail cross-NIC QP
+handshakes. To identify RoCE NICs:
+
+```bash
+ibdev2netdev | grep -v "ib[0-9]"
+# RoCE devices show Ethernet interface names like enp75s0f0
+```
+
+Then configure the connector:
+```yaml
+device_name: "mlx5_2"  # or set RDMA_DEVICE_NAME=mlx5_2
+```
+
+See the RDMA Test README in tests/distributed/omni_connectors/README.md
+for test-specific setup instructions.
+
+For more details on the underlying engine, refer to the
+[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
+
+---
+
+## Design
+
+### 1. Overview
+
+`MooncakeTransferEngineConnector` is the high-performance remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake `TransferEngine` and combines:
+
+- a **direct data plane** for remote memory writes
+- a **ZMQ side channel** for metadata lookup, handshake, and completion signaling
+- a **managed local memory pool** for both send and receive buffers
+
+Unlike `MooncakeStoreConnector`, which treats the backend as a distributed store, `MooncakeTransferEngineConnector` is designed as a peer-to-peer transport. Its goal is to move large stage payloads efficiently while still fitting the common `put()` / `get()` API defined by `OmniConnectorBase`.
+
+It is the most performance-oriented connector in the current OmniConnector family and is intended for large remote payloads such as:
+
+- KV cache transfer
+- stage hidden-state payloads
+- streaming chunk payloads
+- other binary-heavy inter-stage artifacts
+
+### 2. Relationship with the OmniConnector System
+
+`MooncakeTransferEngineConnector` implements the same connector contract as the other backends:
+
+- `put(from_stage, to_stage, put_key, data)`
+- `get(from_stage, to_stage, get_key, metadata=None)`
+- `cleanup(request_id, ...)`
+- `health()`
+- `close()`
+
+It is integrated into the system through the standard connector plumbing:
+
+- `OmniConnectorFactory` constructs the connector from `ConnectorSpec`
+- `load_omni_transfer_config()` resolves the edge-level connector configuration
+- `get_connectors_config_for_stage()` and `resolve_omni_kv_config_for_stage()` inject the connector role
+- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
+
+The key system-level distinction is that this connector is **role-aware**:
+
+- sender instances expose data and listen for pull requests
+- receiver instances allocate buffers and actively pull data from the sender
+
+### 3. Design Goals
+
+The connector is designed around four primary goals:
+
+1. **High-throughput remote transfer**
+   Avoid store-mediated round trips and write directly into the receiver memory region.
+
+2. **Fast path for raw payloads**
+   Support `torch.Tensor`, `bytes`, and `ManagedBuffer` without forcing all traffic through full object serialization.
+
+3. **Unified connector abstraction**
+   Preserve the same `put()` / `get()` interface used by the rest of the OmniConnector stack.
+
+4. **Safe lifecycle management**
+   Manage allocation, reuse, cleanup, and failure recovery for a registered memory pool.
+
+### 4. Architecture Overview
+
+At a high level, the connector is composed of four main subsystems:
+
+```mermaid
+classDiagram
+    class OmniConnectorBase {
+        <<abstract>>
+        +put(from_stage, to_stage, put_key, data)
+        +get(from_stage, to_stage, get_key, metadata)
+        +cleanup(request_id)
+        +health()
+        +close()
+    }
+
+    class MooncakeTransferEngineConnector {
+        +supports_raw_data: bool
+        -engine: TransferEngine
+        -allocator: BufferAllocator
+        -pool: torch.Tensor
+        -zmq_ctx: zmq.Context
+        -_local_buffers: dict
+        -_sender_executor: ThreadPoolExecutor
+        -_listener_thread: threading.Thread
+        +put(...)
+        +get(...)
+        +update_sender_info(sender_host, sender_zmq_port)
+        +get_connection_info()
+        +cleanup(request_id, from_stage, to_stage)
+        +close()
+    }
+
+    class BufferAllocator {
+        -total_size: int
+        -alignment: int
+        -free_blocks: list
+        +alloc(size) int
+        +free(offset, size)
+    }
+
+    class ManagedBuffer {
+        -allocator: BufferAllocator
+        -offset: int
+        -size: int
+        -pool_tensor: torch.Tensor
+        +tensor
+        +as_tensor(dtype, shape) torch.Tensor
+        +to_bytes() bytes
+        +release()
+    }
+
+    class TransferEngine {
+        +initialize(host, handshake, protocol, device_name)
+        +register_memory(base_ptr, size)
+        +batch_transfer_sync_write(remote_session, src_addrs, dst_addrs, lengths)
+        +unregister_memory(base_ptr)
+        +get_rpc_port() int
+    }
+
+    class QueryRequest {
+        +request_id: str
+    }
+
+    class QueryResponse {
+        +request_id: str
+        +data_size: int
+        +is_fast_path: bool
+    }
+
+    class MooncakeAgentMetadata {
+        +remote_hostname: str
+        +remote_port: int
+        +request_id: str
+        +dst_addrs: list[int]
+        +lengths: list[int]
+    }
+
+    OmniConnectorBase <|-- MooncakeTransferEngineConnector
+    MooncakeTransferEngineConnector *-- BufferAllocator
+    MooncakeTransferEngineConnector *-- TransferEngine
+    ManagedBuffer --> BufferAllocator : releases to
+    ManagedBuffer --> "1" torch.Tensor : views
+    MooncakeTransferEngineConnector ..> ManagedBuffer : returns / retains
+    MooncakeTransferEngineConnector ..> QueryRequest : decodes
+    MooncakeTransferEngineConnector ..> QueryResponse : encodes
+    MooncakeTransferEngineConnector ..> MooncakeAgentMetadata : exchanges
+```
+
+#### 4.1 Transfer Engine
+
+Mooncake `TransferEngine` is responsible for the actual data-plane transfer. It registers local memory and performs synchronous remote writes through:
+
+```python
+batch_transfer_sync_write(...)
+```
+
+#### 4.2 Managed Memory Pool
+
+Each connector instance owns a large pre-registered memory pool:
+
+- CPU pinned memory when `memory_pool_device == "cpu"`
+- GPU memory when `memory_pool_device == "cuda"`
+
+This avoids repeated memory registration and allows each transfer to allocate subranges from one long-lived pool.
+
+#### 4.3 Buffer Manager
+
+Two helper classes control local memory ownership:
+
+- `BufferAllocator`
+  Manages aligned subrange allocation and free-list merging.
+
+- `ManagedBuffer`
+  Represents one live slice of the pool and exposes:
+  - `.tensor`
+  - `.as_tensor(dtype, shape)`
+  - `.to_bytes()`
+  - `.release()`
+
+#### 4.4 ZMQ Side Channel
+
+ZMQ is used for transport coordination, not for the data payload itself. It handles:
+
+- metadata query from receiver to sender
+- pull request submission
+- completion or error signaling
+- internal notification from worker threads back to the listener thread
+
+This split makes the control plane lightweight while keeping the bulk payload on the transfer engine data plane.
+
+### 5. Role Model
+
+#### 5.1 Sender Role
+
+A sender connector:
+
+- accepts `put()` calls
+- stores live transfer-ready buffers in `_local_buffers`
+- starts a ZMQ listener thread
+- responds to metadata queries and pull requests from receivers
+
+#### 5.2 Receiver Role
+
+A receiver connector:
+
+- does not bind the sender-side ZMQ listener
+- accepts `get()` calls
+- allocates receive buffers from its own pool
+- requests metadata or transfer service from the sender
+
+The role is not inferred dynamically. It is injected by the stage configuration layer:
+
+- incoming edge for a stage -> `role="receiver"`
+- outgoing edge for a stage -> `role="sender"`
+
+This is important because incorrect role assignment would break initialization semantics.
+
+#### 5.3 Host Auto-Detection
+
+The `host` configuration field supports the special value `"auto"`. When set, the connector auto-detects the local IP address that would be used for external communication (via a UDP socket probe to `8.8.8.8`). If that fails, it falls back to hostname resolution, and ultimately to `127.0.0.1`.
+
+This is useful in environments where the operator does not want to hard-code IP addresses in the connector config.
+
+#### 5.4 RDMA Device Filtering
+
+The `device_name` configuration field allows the operator to specify which RDMA NICs to use (comma-separated, e.g. `"mlx5_0,mlx5_1"`). If not set in config, the connector also checks the `RDMA_DEVICE_NAME` environment variable.
+
+This is important in environments with mixed InfiniBand/RoCE NICs, where not all devices are suitable for the transfer engine.
+
+### 6. Local Memory Management
+
+#### 6.1 Memory Pool Registration
+
+During initialization, the connector:
+
+1. allocates a large pool tensor
+2. records its base pointer
+3. registers that memory with Mooncake
+4. creates a `BufferAllocator` for subrange management
+
+This means every later transfer only allocates offsets inside the pre-registered pool rather than registering memory per request.
+
+#### 6.2 BufferAllocator
+
+`BufferAllocator` maintains a sorted free list of `(offset, size)` blocks and enforces alignment. Its responsibilities include:
+
+- aligned allocation
+- freeing previously allocated blocks
+- adjacent block merging
+- double-free detection
+- overlap detection to catch corruption
+
+This is a critical piece of the connector because both sender and receiver depend on long-lived pool reuse.
+
+#### 6.3 ManagedBuffer
+
+`ManagedBuffer` is the main fast-path data wrapper. It can:
+
+- expose the pool slice as a zero-copy 1D `uint8` tensor
+- reinterpret that slice as a typed tensor
+- copy out the contents as Python `bytes`
+- release the slice back to the allocator
+
+The connector uses `ManagedBuffer` in two different ways:
+
+- as a send-side holder to keep the pool slice alive
+- as a receive-side return type when `is_fast_path=True`
+
+### 7. Put Flow
+
+#### 7.1 High-Level Behavior
+
+`put()` is only valid in sender mode. Its job is to expose a payload for later remote pull by the receiver.
+
+The high-level flow is:
+
+1. validate connector state and role
+2. convert the input into a pool-backed transferable representation
+3. store the transfer metadata in `_local_buffers`
+4. return lightweight metadata describing how the receiver can fetch the data
+
+#### 7.2 Payload Type Handling
+
+`put()` supports three payload classes:
+
+**A. `ManagedBuffer`**
+
+If the buffer belongs to the same pool, the connector can use it directly without copying. This is the most efficient path.
+
+If the buffer comes from a different pool, the connector falls back to a copy path.
+
+**B. `torch.Tensor` or `bytes`**
+
+These payloads are copied into the local pool and marked as fast-path data:
+
+- no Omni object serialization is required
+- receiver can get a `ManagedBuffer` back
+
+**C. Generic Python object**
+
+Any other payload is serialized via `OmniSerializer.serialize(...)` and then copied into the pool.
+
+In this case:
+
+- `is_fast_path=False`
+- the receiver will deserialize back into a Python object
+
+#### 7.3 Sender Metadata
+
+The sender returns:
+
+```python
+{
+    "source_host": self.host,
+    "source_port": self.zmq_port,
+    "data_size": size,
+    "is_fast_path": is_fast_path,
+}
+```
+
+This metadata is intentionally lightweight. It tells the receiver:
+
+- where the sender-side control plane lives
+- how large the remote transfer will be
+- whether the payload should be returned as a `ManagedBuffer` or a deserialized object
+
+#### 7.4 Sender Buffer Table
+
+The sender stores each live payload in `_local_buffers` under the stage-qualified key. Each entry contains:
+
+- source addresses
+- lengths
+- the holder object
+- ownership information (`should_release`)
+- `is_fast_path`
+- creation time
+
+This table is the sender-side truth source for both metadata queries and pull requests.
+
+### 8. Get Flow
+
+#### 8.1 High-Level Behavior
+
+`get()` runs on the receiver side and performs four steps:
+
+1. resolve metadata
+2. allocate a destination buffer in the local pool
+3. request the sender to write into that destination buffer
+4. return either a `ManagedBuffer` or a deserialized object
+
+#### 8.2 Metadata Resolution Paths
+
+The `metadata` parameter in `get()` is optional. The connector supports two resolution modes depending on whether the caller supplies it.
+
+**With metadata**
+
+When the caller passes metadata, the connector uses it directly. The metadata carries:
+
+- `source_host` / `source_port` — sender ZMQ endpoint
+- `data_size` — payload byte count
+- `is_fast_path` — whether the receiver gets a `ManagedBuffer` or a deserialized object
+
+This mode is suitable when the control plane already forwards the sender's `put()` output to the receiver.
+
+**Without metadata**
+
+When `get(metadata=None)` is called, the connector queries the sender over ZMQ to discover the same fields (`data_size`, `is_fast_path`). The caller must first call:
+
+```python
+update_sender_info(sender_host, sender_zmq_port)
+```
+
+so that the connector knows where to send the query.
+
+This mode is suitable for polling-based flows (e.g. KV transfer, async chunk transfer) where the receiver does not have metadata from the control plane.
+
+#### 8.3 Destination Allocation
+
+Once metadata is resolved, the receiver:
+
+1. allocates a subrange from its own local pool
+2. wraps it in a `ManagedBuffer`
+3. builds a `MooncakeAgentMetadata` request containing:
+   - receiver hostname
+   - receiver RPC port
+   - request ID
+   - destination addresses
+   - transfer lengths
+
+This tells the sender exactly where to write the incoming data.
+
+#### 8.4 Transfer Completion
+
+The receiver then sends the pull request over ZMQ and waits for:
+
+- `TRANS_DONE`
+- or `TRANS_ERROR`
+
+If the transfer succeeds:
+
+- for `is_fast_path=True`, the receiver returns `(ManagedBuffer, size)`
+- for `is_fast_path=False`, the receiver copies to bytes, deserializes, releases the buffer, and returns `(object, size)`
+
+### 9. Sender-Side Listener Design
+
+#### 9.1 Listener Thread
+
+In sender mode, the connector starts `_zmq_listener_loop()` after initialization. The listener:
+
+- binds `tcp://{host}:{zmq_port}`
+- receives incoming requests
+- uses a poller for socket events and internal notifications
+- periodically reclaims stale buffers
+
+If the bind fails, initialization fails immediately. The code does not silently downgrade the connector role.
+
+#### 9.2 Worker Thread Pool
+
+The listener hands work to `_sender_executor` so that the listener thread does not block on transfer work.
+
+There are two request types:
+
+- metadata query -> `_handle_query_request(...)`
+- transfer request -> `_handle_pull_request(...)`
+
+#### 9.3 Query Handling
+
+For metadata queries, the sender looks up the request ID in `_local_buffers` and returns:
+
+- data size
+- fast-path flag
+
+This supports consumers that only know the sender endpoint but not the original sender metadata.
+
+#### 9.4 Pull Handling
+
+For a transfer request, the sender:
+
+1. locates the source addresses in `_local_buffers`
+2. constructs the remote session identifier
+3. calls `batch_transfer_sync_write(...)`
+4. replies `TRANS_DONE` or `TRANS_ERROR`
+
+On success, the sender immediately calls `cleanup(meta.request_id)` and frees the producer-side buffer if it owns it.
+
+This choice is important: it makes the connector effectively a **single-consumer transfer model** for each successful put/get pair.
+
+### 10. Fast Path Semantics
+
+This connector explicitly advertises:
+
+```python
+supports_raw_data = True
+```
+
+That means it can move raw payloads without forcing everything through the Omni object serializer.
+
+#### Fast Path
+
+For `torch.Tensor`, `bytes`, or pool-local `ManagedBuffer`:
+
+- sender returns `is_fast_path=True`
+- receiver returns a `ManagedBuffer`
+- caller is responsible for calling `release()`
+
+This avoids an unnecessary copy on the receiver side.
+
+#### Serialized Object Path
+
+For arbitrary Python objects:
+
+- sender serializes the payload
+- receiver converts the receive buffer to bytes
+- receiver deserializes the object
+- receive buffer is released internally
+
+This preserves a uniform object-oriented API while still allowing optimized raw-data transport when possible.
+
+### 11. Failure Handling and Cleanup
+
+#### 11.1 Timeouts and Socket Recovery
+
+The receiver caches ZMQ REQ sockets per thread, but invalidates them after failures. This avoids reusing sockets that may be stuck in a bad state after timeout or receive errors.
+
+Timeout is scaled based on payload size:
+
+- a base timeout
+- plus additional time for large payloads
+
+This is intended to reduce false timeouts for large remote writes.
+
+#### 11.2 Stale Buffer Reclamation
+
+The sender periodically reclaims old entries from `_local_buffers` using a TTL policy. This protects the memory pool from permanent leaks if a receiver crashes or never consumes a prepared payload.
+
+This is a practical recovery mechanism, although the code notes that TTL cleanup can still race with very long-running in-flight transfers.
+
+#### 11.3 Connector Shutdown
+
+`close()` is a full resource teardown routine. It:
+
+- stops the listener thread
+- shuts down the worker executor
+- releases all pending buffers
+- closes cached sockets
+- unregisters memory from the engine when supported
+- terminates the ZMQ context
+- drops the pool reference
+
+This makes `MooncakeTransferEngineConnector` the most lifecycle-aware connector in the current connector family.
+
+### 12. Current Implementation Constraints
+
+The current code documents several important topology constraints.
+
+#### 12.1 One Sender to One Receiver per Successful Transfer
+
+After a successful transfer, the sender-side buffer is cleaned up immediately. This means the same prepared payload is not retained for multiple independent receivers.
+
+#### 12.2 One Receiver to One Active Sender Endpoint
+
+The receiver only stores one `(sender_host, sender_zmq_port)` pair through `update_sender_info(...)`. So the metadata-query mode is currently single-sender at a time.
+
+#### 12.3 Explicit Buffer Ownership Matters
+
+When the connector allocates a pool slice internally, it is responsible for releasing it. When a caller passes an externally owned `ManagedBuffer`, the connector keeps it alive for transfer but does not assume ownership of its eventual release.
+
+These constraints are consistent with the current implementation and should be treated as design assumptions rather than incidental behavior.
+
+### 13. Data Flow in the Pipeline
+
+The end-to-end sender/receiver interaction is:
+
+```mermaid
+sequenceDiagram
+    participant SenderStage
+    participant SenderConnector
+    participant ReceiverConnector
+    participant ReceiverStage
+
+    SenderStage->>SenderConnector: put(from_stage, to_stage, put_key, data)
+    SenderConnector->>SenderConnector: place payload in local memory pool
+    SenderConnector-->>SenderStage: metadata(source_host, source_port, data_size, is_fast_path)
+
+    ReceiverStage->>ReceiverConnector: get(..., metadata)
+    ReceiverConnector->>ReceiverConnector: allocate destination buffer
+    ReceiverConnector->>SenderConnector: ZMQ pull request with dst addr
+    SenderConnector->>ReceiverConnector: TransferEngine remote write
+    SenderConnector-->>ReceiverConnector: TRANS_DONE
+    ReceiverConnector-->>ReceiverStage: ManagedBuffer or deserialized object
+```
+
+For metadata-less polling, the flow simply adds a metadata query step before the pull request.
+
+### 14. Strengths and Trade-offs
+
+#### Strengths
+
+- Best remote-transfer design in the current connector stack for large payloads.
+- Supports raw-data fast path.
+- Keeps stage communication under the same connector abstraction.
+- Includes real lifecycle and memory-pool management.
+- Works for both stage payload transfer and KV transfer scenarios.
+
+#### Trade-offs
+
+- More complex than the store-based connector.
+- Correctness depends on role injection and endpoint coordination.
+- Caller must release fast-path receive buffers.
+- Current implementation is optimized for single-consumer transfer semantics.
+
+### 15. Summary
+
+`MooncakeTransferEngineConnector` is the high-performance peer-to-peer transport in the OmniConnector system. Its design combines:
+
+- a registered memory pool
+- a safe subrange allocator
+- a ZMQ control plane
+- a Mooncake transfer-engine data plane
+
+This allows the connector to support both:
+
+1. a **fast path** for raw tensors and bytes
+2. a **generic object path** for arbitrary Python payloads
+
+Within vLLM-Omni, it is the connector that most directly targets performance-sensitive remote transfer, especially for large payloads and KV cache movement. Its additional complexity is deliberate: it is the connector that turns the generic OmniConnector abstraction into a transport capable of efficient remote memory movement rather than simple object storage.
@@ -7,9 +7,10 @@ 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.
+All payloads are serialized and stored in shared memory (`/dev/shm`); the SHM
+segment name is returned in metadata. The configuration exposes a
+`shm_threshold_bytes` field for a future inline-vs-SHM split, but the current
+implementation always uses shared memory regardless of payload size.

 ## Configuration

@@ -25,3 +26,234 @@ runtime:
 ## Notes

 - Auto-mode uses SharedMemoryConnector if no connector is declared for an edge.
+
+---
+
+## Design
+
+### 1. Overview
+
+`SharedMemoryConnector` is the default same-node connector in `vllm_omni/distributed/omni_connectors`. It is designed for stage-to-stage transfer when producer and consumer processes run on the same host and can share `/dev/shm`.
+
+The connector provides a unified `put()` / `get()` API for arbitrary Python objects while keeping the control plane lightweight:
+
+- The payload is serialized by the connector.
+- The serialized bytes are placed in shared memory.
+- The queue/control plane only carries a small metadata handle.
+
+This makes `SharedMemoryConnector` the simplest connector in the OmniConnector family and the default fallback when an edge does not explicitly configure another backend.
+
+### 2. Relationship with the OmniConnector System
+
+`SharedMemoryConnector` implements `OmniConnectorBase`, so it follows the same lifecycle and API contract as the other connectors:
+
+- `put(from_stage, to_stage, put_key, data)`
+- `get(from_stage, to_stage, get_key, metadata=None)`
+- `cleanup(request_id)`
+- `health()`
+- `close()`
+
+Within the larger system:
+
+- `load_omni_transfer_config()` automatically fills missing edges with `SharedMemoryConnector`.
+- Callers interact with the connector exclusively through the `put()` / `get()` / `cleanup()` contract — the connector does not require caller-specific logic.
+
+Compared with the remote Mooncake-based connectors, `SharedMemoryConnector` is intentionally minimal and local-only.
+
+### 3. Design Goals
+
+The connector is built around the following goals:
+
+- **Low-friction local transfer** for single-node multi-process pipelines.
+- **Unified object semantics** for arbitrary Python payloads.
+- **Small control-plane overhead** by passing only metadata through queues.
+- **Zero external dependencies** beyond Python shared memory and the existing stage utilities.
+
+It is not intended to provide cross-node transfer, RDMA, or raw tensor zero-copy semantics across processes.
+
+### 4. Core Design
+
+#### 4.1 Serialization Model
+
+`SharedMemoryConnector` always starts from a Python object and serializes it through the shared Omni serializer:
+
+```python
+payload = self.serialize_obj(data)
+```
+
+This keeps the connector behavior consistent with the rest of the connector stack:
+
+- producer code does not need connector-specific serialization logic
+- consumer code always receives the original object after deserialization
+- the connector can reuse the same serializer used by other backends
+
+#### 4.2 Shared Memory as the Data Plane
+
+The actual data plane is a shared-memory segment created by:
+
+- `shm_write_bytes(...)`
+- `shm_read_bytes(...)`
+
+The connector stores a small metadata object such as:
+
+```python
+{
+    "shm": {"name": ..., "size": ...},
+    "size": ...
+}
+```
+
+This metadata is passed over the control plane and allows the downstream stage to locate the shared-memory segment.
+
+#### 4.3 Locking Model
+
+To avoid races between the producer and consumer, the connector uses a lock file per request:
+
+```text
+/dev/shm/shm_{put_key}_lockfile.lock
+```
+
+Locking is done with `fcntl.flock`:
+
+- producer uses `LOCK_EX`
+- consumer uses `LOCK_EX`
+
+Both sides acquire an exclusive lock. This ensures that the shared-memory segment is not read while it is still being written and makes the handoff safer in a multi-process environment.
+
+### 5. Put / Get Flow
+
+#### 5.1 Producer Flow: `put()`
+
+The producer-side flow is:
+
+1. Serialize the input object to bytes.
+2. Compute the payload size.
+3. Acquire the per-request lock file.
+4. Write the bytes into shared memory.
+5. Return lightweight metadata to the caller.
+
+The returned tuple is:
+
+```python
+(success, serialized_size, metadata)
+```
+
+where `metadata` contains the shared-memory handle needed by the consumer.
+
+#### 5.2 Consumer Flow: `get(metadata=...)`
+
+The primary consumer path is metadata-driven:
+
+1. Extract the shared-memory handle from `metadata`.
+2. Acquire the exclusive lock.
+3. Read the raw bytes from shared memory.
+4. Deserialize the bytes back into the original Python object.
+5. Remove the lock file if it still exists.
+
+This is the path used by the current stage-to-stage connector flow.
+
+#### 5.3 Compatibility Flow: `get(metadata=None)`
+
+The connector also keeps a compatibility path for callers that only know the key:
+
+1. Attempt to open the shared-memory segment by name via `SharedMemory(name=get_key)`.
+2. If the segment exists and has non-zero size, acquire the exclusive lock and read the bytes.
+3. Deserialize the bytes and return the object.
+
+If the segment does not exist or any exception occurs, the call returns `None` immediately. There is no retry loop in this path -- it is a single-attempt open.
+
+This path is mainly for older code paths and is not the preferred mode for the current connector pipeline.
+
+### 6. Key Implementation Characteristics
+
+#### 6.1 Threshold Exists, but the Current Code Always Uses SHM
+
+The class keeps a `shm_threshold_bytes` field and still exposes metrics for inline writes. However, the current implementation uses:
+
+```python
+if True:
+    ...
+```
+
+inside `put()`, which means the current code path always writes to shared memory.
+
+So the design still suggests a future split between:
+
+- small payloads inline
+- large payloads in shared memory
+
+but the current behavior is effectively:
+
+- all payloads go through shared memory
+
+This should be documented because it affects real runtime behavior.
+
+#### 6.2 Cleanup Is Currently Passive
+
+`cleanup()` is currently a no-op. The intended assumption is:
+
+- the consumer reads the segment
+- the underlying shared-memory helpers unlink it
+
+If the consumer never executes `get()`, the shared-memory segment may remain allocated. This means the connector relies on the normal success path for resource reclamation.
+
+#### 6.3 Close Is Currently Minimal
+
+`close()` is also a no-op. There is no connector-owned background thread, socket, or memory pool to tear down, so the lifecycle is simple. The trade-off is that `close()` does not scan or recover leaked shared-memory resources.
+
+### 7. Data Flow in the Pipeline
+
+The typical flow with `SharedMemoryConnector` is:
+
+```mermaid
+sequenceDiagram
+    participant SenderStage
+    participant SharedMemoryConnector
+    participant QueueOrControlPlane
+    participant ReceiverStage
+
+    SenderStage->>SharedMemoryConnector: put(from_stage, to_stage, put_key, data)
+    SharedMemoryConnector->>SharedMemoryConnector: serialize object
+    SharedMemoryConnector->>SharedMemoryConnector: write bytes to /dev/shm
+    SharedMemoryConnector-->>SenderStage: metadata {shm, size}
+    SenderStage->>QueueOrControlPlane: forward connector metadata
+    QueueOrControlPlane->>ReceiverStage: task + connector metadata
+    ReceiverStage->>SharedMemoryConnector: get(from_stage, to_stage, get_key, metadata)
+    SharedMemoryConnector->>SharedMemoryConnector: read bytes from /dev/shm
+    SharedMemoryConnector->>SharedMemoryConnector: deserialize object
+    SharedMemoryConnector-->>ReceiverStage: (data, size)
+```
+
+This is a classic split-control-plane / data-plane design, but constrained to a single host.
+
+### 8. Strengths and Trade-offs
+
+#### Strengths
+
+- Very simple deployment model.
+- No external service dependency.
+- Fits naturally into the existing queue-driven orchestration flow.
+- Good default for local multi-process pipelines.
+
+#### Trade-offs
+
+- Same-node only.
+- Full object serialization and deserialization are still required.
+- Resource cleanup depends on the normal consumer path.
+- Shared memory capacity is limited by host configuration.
+
+### 9. Summary
+
+`SharedMemoryConnector` is the baseline local transport for the OmniConnector system. Its design is intentionally straightforward:
+
+- serialize object
+- place bytes in shared memory
+- pass metadata through the control plane
+- deserialize on the receiving side
+
+It plays two important roles in vLLM-Omni:
+
+1. It is the simplest production-ready connector for same-node stage pipelines.
+2. It serves as the automatic fallback connector when no explicit edge transport is configured.
+
+Although the current implementation is deliberately minimal, it provides the foundation for reliable local connector semantics and keeps the stage communication model uniform across the system.
@@ -94,8 +94,265 @@ stage_args:
 Parameters:

 - host: datasystem worker host.
- port: datasystem worker port.
+- port: datasystem worker port (default: `35001` if omitted; the example above uses `31501` to match the worker startup command).
 - 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).
+
+---
+
+## Design
+
+### 1. Overview
+
+`YuanrongConnector` is the Datasystem-based remote connector in `vllm_omni/distributed/omni_connectors`. It uses Yuanrong Datasystem's distributed key-value client as the transport backend and exposes the same `put()` / `get()` interface as the other OmniConnectors.
+
+Like `MooncakeStoreConnector`, it is a store-oriented remote connector rather than a direct peer-to-peer transport. Its role is to let stage payloads move across nodes through a deterministic key-based storage abstraction while keeping the rest of the pipeline on the common connector API.
+
+It is intended for deployments that already use Yuanrong Datasystem and want a remote connector that integrates with the existing OmniConnector configuration and orchestration model.
+
+### 2. Relationship with the OmniConnector System
+
+`YuanrongConnector` implements `OmniConnectorBase`, so it participates in the same connector lifecycle as the other backends:
+
+- `OmniConnectorFactory` constructs it from a `ConnectorSpec`
+- stage edge configuration is resolved by `load_omni_transfer_config()`
+- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
+
+This means the connector is not exposed directly to stage logic. Stages only interact with the generic connector contract, and the backend choice remains a configuration concern.
+
+### 3. Design Goals
+
+The connector is built around the following goals:
+
+1. **Cross-node payload transfer through Datasystem**
+   Reuse Yuanrong Datasystem as the remote exchange medium for stage data.
+
+2. **Uniform object transfer semantics**
+   Allow arbitrary Python objects to be transmitted through the shared Omni serializer.
+
+3. **Minimal connector-specific control plane**
+   Use deterministic keys so that consumers can retrieve data without an extra transport metadata handoff.
+
+4. **Operational reuse of existing infrastructure**
+   Fit into environments that already deploy Yuanrong Datasystem workers and etcd.
+
+The connector is not designed for direct remote-memory writes or tensor-specific fast-path transfer.
+
+### 4. Core Design
+
+#### 4.1 Store-Oriented Transfer Model
+
+`YuanrongConnector` treats the transport backend as a distributed object store:
+
+1. serialize the Python payload
+2. build a deterministic connector key
+3. write the serialized bytes into Datasystem
+4. read the bytes back on the receiver side
+5. deserialize them into the original object
+
+This is the same broad architectural class as `MooncakeStoreConnector`, but implemented on top of Yuanrong Datasystem APIs instead of Mooncake store APIs.
+
+#### 4.2 Deterministic Keying
+
+Unlike the default `_make_key()` in `OmniConnectorBase`, `YuanrongConnector` defines its own key format:
+
+```text
+{request_id}:{from_stage}_{to_stage}
+```
+
+This has two design implications:
+
+- the request identifier remains the primary lookup handle
+- stage routing information is embedded in the key so that the same logical request ID can safely appear on different edges
+
+The explicit override also makes the key format easier to align with Datasystem-side debugging and operational inspection.
+
+#### 4.3 No Extra Metadata Hand-off
+
+`put()` returns:
+
+```python
+(success, serialized_size, None)
+```
+
+and does not generate connector-specific metadata.
+
+This design works because the receiver can reconstruct the exact same key from:
+
+- `get_key`
+- `from_stage`
+- `to_stage`
+
+As a result, the connector does not require a separate side-channel metadata handoff.
+
+### 5. Initialization
+
+#### 5.1 Datasystem Client Dependency
+
+The connector requires the Datasystem Python bindings to expose:
+
+- `KVClient`
+- `SetParam`
+- `WriteMode`
+
+If any of these symbols are unavailable, connector construction fails immediately with `ImportError`. This keeps configuration errors explicit and avoids a partially initialized runtime.
+
+#### 5.2 Client Setup
+
+During `_init_client()`, the connector:
+
+1. reads `host` and `port`
+2. creates `KVClient(host, port)`
+3. calls `client.init()`
+
+At construction time it also creates a `SetParam` and fixes:
+
+```python
+self.set_param.write_mode = WriteMode.NONE_L2_CACHE_EVICT
+```
+
+This means the connector has a stable write policy for all writes and does not currently expose write-mode selection as a higher-level connector option.
+
+### 6. Put / Get Flow
+
+#### 6.1 Producer Flow: `put()`
+
+The producer-side flow is:
+
+1. verify that the Datasystem client has been initialized
+2. serialize the input object with the shared Omni serializer
+3. build the Datasystem key using the connector-specific `_make_key()`
+4. call `client.set(key, serialized_data, self.set_param.write_mode)`
+5. update metrics and return success
+
+The returned metadata is always `None`, because the Datasystem key itself is the lookup contract between producer and consumer.
+
+#### 6.2 Consumer Flow: `get()`
+
+The consumer-side flow is:
+
+1. verify that the Datasystem client has been initialized
+2. rebuild the same key with `from_stage`, `to_stage`, and `get_key`
+3. call:
+
+```python
+client.get([key], False, self.get_sub_timeout_ms)
+```
+
+4. take the first returned element if present
+5. deserialize it and return `(data, payload_size)`
+
+If the returned list is empty or contains no data for the key, `get()` returns `None`.
+
+### 7. Timeout and Retrieval Semantics
+
+The connector uses `get_sub_timeout_ms` as its read timeout. Unlike `MooncakeStoreConnector`, which performs an explicit retry loop in Python, `YuanrongConnector` delegates waiting behavior more directly to the Datasystem client call.
+
+This leads to a slightly different retrieval model:
+
+- `MooncakeStoreConnector`: retry-oriented polling in connector code
+- `YuanrongConnector`: single client call with backend-managed timeout behavior
+
+From the connector API perspective the result is the same, but operationally the waiting behavior is more dependent on Datasystem client semantics.
+
+### 8. Integration with Stage Communication
+
+All callers use the connector through the same `put()` / `get()` contract:
+
+- the sender calls `put()` to serialize and store the payload
+- the receiver calls `get()` to retrieve and deserialize it
+- no connector-specific metadata is required, since the Datasystem key is the rendezvous point
+
+Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, and there is no raw tensor fast path, which makes the connector functional but not optimized for the largest payloads.
+
+### 9. Data Flow in the Pipeline
+
+The end-to-end transfer model is:
+
+```mermaid
+sequenceDiagram
+    participant SenderStage
+    participant YuanrongConnector
+    participant Datasystem
+    participant ReceiverStage
+
+    SenderStage->>YuanrongConnector: put(from_stage, to_stage, put_key, data)
+    YuanrongConnector->>YuanrongConnector: serialize object
+    YuanrongConnector->>Datasystem: set(key, bytes)
+    YuanrongConnector-->>SenderStage: (success, size, None)
+
+    ReceiverStage->>YuanrongConnector: get(from_stage, to_stage, get_key)
+    YuanrongConnector->>Datasystem: get([key], timeout)
+    YuanrongConnector->>YuanrongConnector: deserialize bytes
+    YuanrongConnector-->>ReceiverStage: (data, size)
+```
+
+This is a store-mediated remote connector design with deterministic key lookup and no explicit side-channel metadata exchange.
+
+### 10. Strengths and Trade-offs
+
+#### Strengths
+
+- Reuses existing Yuanrong Datasystem infrastructure.
+- Keeps the connector contract simple and uniform.
+- Requires no connector-specific metadata handoff.
+- Suitable for remote stage transfer in Datasystem-based deployments.
+
+#### Trade-offs
+
+- Always pays full serialization and deserialization cost.
+- Does not support raw bytes / tensor fast-path semantics.
+- Depends on external Datasystem worker availability.
+- Retrieval and timeout behavior are partly delegated to the backend client.
+
+### 11. Important Implementation Characteristics
+
+#### 11.1 Cleanup Is a No-op
+
+`cleanup()` only logs a debug message and does not explicitly remove data from Datasystem. The current design assumes backend-side lifecycle or garbage collection rather than request-scoped delete semantics inside the connector.
+
+This mirrors the same design trade-off seen in other store-based connectors: simplicity over explicit per-request data reclamation.
+
+#### 11.2 Close Only Releases the Local Client Handle
+
+`close()` does not perform a remote shutdown. It simply clears the local `client` reference and marks the connector as closed from the process perspective:
+
+```python
+self.client = None
+```
+
+This is appropriate for a client-based store connector, but it also means that backend resource lifecycle remains outside the connector's control.
+
+#### 11.3 Error and Timeout Metrics Are Coarse-Grained
+
+The connector tracks:
+
+- `puts`
+- `gets`
+- `bytes_transferred`
+- `errors`
+- `timeouts`
+
+In the current code, failed `put()` increments `errors`, while `get()` exceptions increment `timeouts`. This is operationally useful, but it does not distinguish between:
+
+- backend timeout
+- not-found result
+- transport failure
+- deserialization failure
+
+So the metrics should be read as high-level health indicators, not detailed root-cause diagnostics.
+
+### 12. Summary
+
+`YuanrongConnector` is the Datasystem-backed remote connector in the OmniConnector family. Its design is straightforward:
+
+- serialize payloads
+- store them under a deterministic stage-qualified key
+- retrieve them by the same key
+- deserialize them on the receiving side
+
+Within vLLM-Omni, it provides a clean integration point for Yuanrong Datasystem-based deployments while preserving the same connector abstraction used by the rest of the pipeline.
+
+Its design priorities are simplicity, infrastructure reuse, and API consistency rather than direct-memory transport optimization.
@@ -45,7 +45,7 @@ ray start --address=<HEAD_NODE_IP>:6399

 When running on Ray, the system automatically adapts its communication strategy:

-*   **Cross-Node**: Recommended to use `MooncakeConnector` (requires separate configuration).
+*   **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback).
 *   **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
 *   **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.

@@ -0,0 +1,490 @@
+# Sequence Parallel
+
+This section describes how to add Sequence Parallel (SP) to a diffusion transformer model. We use the Qwen-Image transformer and Wan2.2 transformer as reference implementations.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Approach 1: Non-Intrusive `_sp_plan` (Recommended)](#approach-1-non-intrusive-_sp_plan-recommended)
+- [Approach 2: Intrusive Modification (For Complex Cases)](#approach-2-intrusive-modification-for-complex-cases)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+
+### What is Sequence Parallel?
+
+**Terminology Note:** Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py). We use "Sequence Parallelism" to align with vLLM-Omni's terminology.
+
+Diffusion transformers process long sequences of image patches or video frames. For high-resolution generation, these sequences can become very large. Enabling SP allows each GPU to process only a portion of the sequence, with attention mechanisms (Ulysses/Ring) handling cross-GPU communication transparently.
+
+### Architecture
+
+The major APIs for Sequence Parallel:
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,   # For sharding (splitting) tensors
+    SequenceParallelOutput,  # For gathering tensors
+)
+from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
+```
+
+| Method/Class | Purpose | Behavior |
+|--------------|---------|----------|
+| `SequenceParallelInput` | Declare input sharding in `_sp_plan` | Auto-shards tensors at module input |
+| `SequenceParallelOutput` | Declare output gathering in `_sp_plan` | Auto-gathers tensors at module output |
+| `sp_shard()` | Manual tensor sharding | Splits tensor across SP workers |
+| `sp_gather()` | Manual tensor gathering | Gathers sharded tensors from all workers |
+
+---
+
+## Approach 1: Non-Intrusive `_sp_plan` (Recommended)
+
+The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.
+
+**When to use:**
+- Standard transformer architectures
+- Tensor operations happen at `nn.Module` boundaries
+- Predictable sharding/gathering patterns
+
+This is the ideal approach for integrating sequence parallelism into new models, as it is easier to maintain and ensure compatibility with other types of acceleration.
+
+**How it works:**
+1. Declare `_sp_plan` dict in your transformer class
+2. Framework automatically applies hooks when `sequence_parallel_size > 1`
+3. Hooks shard/gather tensors at specified module boundaries
+4. Attention layers handle cross-GPU communication internally
+
+```python
+class StandardTransformer(nn.Module):
+    _sp_plan = {
+        # Shard hidden_states at first transformer block input
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        # Gather at final output projection
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+`StandardTransformer` has a transformer blocks list `self.blocks = nn.ModuleList([...])`, and a projection output layer `self.proj_out`. The `_sp_plan` above defines that when SP is enabled, sharding the input tensor to the first transformer block, and gathering the sharded tensor at the final output projection layer.
+
+**Requirements:**
+- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries**
+- Inline Python operations (e.g., `torch.cat`, `pad_sequence`) **cannot be hooked**
+
+**Solution for inline operations:** Extract into a submodule (see Step 2 below).
+
+### Step 1: Understand Module Boundaries
+
+Identify where tensors need to be sharded or gathered in your model's forward pass:
+
+```python
+class MyTransformer(nn.Module):
+    def __init__(self):
+        self.patch_embed = PatchEmbed()      # ← Boundary 1
+        self.pos_embed = RoPE()              # ← Boundary 2
+        self.blocks = nn.ModuleList([...])   # ← Boundary 3
+        self.norm_out = LayerNorm()
+        self.proj_out = Linear()             # ← Boundary 4
+
+    def forward(self, x):
+        x = self.patch_embed(x)              # ← Shard before this?
+        pos = self.pos_embed(x)              # ← Shard RoPE outputs?
+        for block in self.blocks:
+            x = block(x, pos)                # ← Blocks process sharded x
+        x = self.norm_out(x)
+        output = self.proj_out(x)            # ← Gather after this?
+        return output
+```
+
+### Step 2: Handle Inline Operations
+
+If your `forward()` contains inline tensor operations, **extract them into submodules**.
+
+**Example: Z-Image concatenates image + text features inline**
+
+```python
+# ❌ BAD: Inline operation - hooks cannot intercept
+class ZImageTransformer(nn.Module):
+    def forward(self, x, cap_feats):
+        # This concatenation happens inline - _sp_plan can't shard it!
+        unified = torch.cat([x, cap_feats], dim=1)
+
+        for layer in self.layers:
+            unified = layer(unified)
+
+        return unified
+
+# ✅ GOOD: Extract into submodule
+class UnifiedPrepare(nn.Module):
+    """Submodule to concatenate image and text features."""
+    def forward(self, x, cap_feats):
+        return torch.cat([x, cap_feats], dim=1)
+
+class ZImageTransformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.unified_prepare = UnifiedPrepare()  # Now a module!
+        self.layers = nn.ModuleList([...])
+
+    def forward(self, x, cap_feats):
+        # Now _sp_plan can shard the output of unified_prepare!
+        unified = self.unified_prepare(x, cap_feats)
+
+        for layer in self.layers:
+            unified = layer(unified)
+
+        return unified
+```
+
+**Other common cases:**
+- `pad_sequence()` → `PadSequenceModule`
+- `torch.cat()` → `ConcatModule`
+- `tensor.reshape()` → `ReshapeModule`
+- Complex preprocessing → `PreprocessModule`
+
+### Step 3: Write `_sp_plan` for Your Model
+
+Create a class-level `_sp_plan` dictionary specifying where to shard/gather tensors.
+
+Typically, there are two patterns for diffusion models:
+
+**Pattern 1: Shard at first block, gather at output projection**
+
+Most common pattern for standard transformers:
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,   # For sharding (splitting) tensors
+    SequenceParallelOutput,  # For gathering tensors
+)
+class StandardTransformer(nn.Module):
+    _sp_plan = {
+        # Shard hidden_states at first transformer block input
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        # Gather at final output projection
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+**Pattern 2: Shard RoPE embeddings separately**
+
+When RoPE is computed in a separate module:
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,   # For sharding (splitting) tensors
+    SequenceParallelOutput,  # For gathering tensors
+)
+class TransformerWithRoPE(nn.Module):
+    _sp_plan = {
+        # Shard RoPE module OUTPUTS (returns tuple of cos, sin)
+        "rope": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # cos
+            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # sin
+        },
+        # Shard transformer block INPUT
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        # Gather at output
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+**Pattern 3: Shard RoPE for Dual Stream Attention**
+In some cases, different streams in attention may need to handle sequence parallelism differently. For example, we may want to shard the image embeddings, while replicating the text embeddings to correctly configure joint attention.
+
+```python
+class DualStreamTransformer(nn.Module):
+    """
+    Dual-stream model where we need to replicate the text components, but shard
+    the image components to correctly handle sequence parallelism.
+    """
+    _sp_plan = {
+        # In this case, the rope_preparer returns a tuple of len 4, where the
+        # first 2 items correspond to the text, and the second 2 correspond to
+        # visual inputs, so we only shard the second.
+        "rope_preparer": {
+            # Outputs 0, 1 (text) - NOT sharded (replicated)
+            # Outputs 2, 3 (image) - sharded
+            2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),  # img_cos
+            3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),  # img_sin
+        },
+        # Shard transformer block INPUT
+        "transformer_blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        # Gather at output
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+NOTE: be careful to test adequately when refactoring classes that take this style of plan, as changing the order of the return values will break sequence parallelism.
+
+### API Reference
+
+**SequenceParallelInput Parameters:**
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `split_dim` | int | Dimension to split (usually `1` for sequence) |
+| `expected_dims` | int \| None | Expected tensor rank for validation (optional) |
+| `split_output` | bool | `False`: shard **input** params; `True`: shard **output** tensors |
+| `auto_pad` | bool | Auto-pad if sequence not divisible by world_size (default: `False`) |
+
+**SequenceParallelOutput Parameters:**
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `gather_dim` | int | Dimension to gather (usually `1` for sequence) |
+| `expected_dims` | int \| None | Expected tensor rank for validation (optional) |
+
+**Module Naming Conventions:**
+
+| Key | Meaning | Python equivalent |
+|-----|---------|-------------------|
+| `""` | Root model | `model` |
+| `"blocks.0"` | First element of ModuleList | `model.blocks[0]` |
+| `"blocks.*"` | All elements of ModuleList | `for b in model.blocks` |
+| `"rope"` | Named submodule | `model.rope` |
+| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` |
+
+**Dictionary Value Types:**
+
+| Key type | `split_output` | Description |
+|----------|----------------|-------------|
+| `"param_name"` (str) | `False` | Shard **input parameter** by name |
+| `0`, `1`, ... (int) | `True` | Shard **output tuple** by index |
+
+---
+
+## Approach 2: Intrusive Modification (For Complex Cases)
+
+For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
+
+
+**When to use:**
+- Dynamic/conditional sharding logic
+- Complex tensor manipulations that can't be encapsulated
+- Temporary workaround during development
+
+```python
+from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
+
+def forward(self, hidden_states, ...):
+    if self.parallel_config.sequence_parallel_size > 1:
+        hidden_states = sp_shard(hidden_states, dim=1)
+
+    # ... computation ...
+
+    if self.parallel_config.sequence_parallel_size > 1:
+        output = sp_gather(output, dim=1)
+
+    return output
+```
+
+---
+
+## Testing
+
+After implementing Sequence Parallel support, thoroughly test your implementation to ensure correctness and performance across different configurations.
+
+**Test Different `sp_size`:**
+
+Test your model with various sequence parallel world sizes to verify correctness and identify optimal configurations:
+
+```bash
+cd examples/offline_inference/text_to_image
+python text_to_image.py \
+    --model Your-org/your-model \
+    --prompt "a cup of coffee on the table" \
+    --num-inference-steps 50 \
+    --ulysses-degree 2 \
+    --ring-degree 2 \
+    --output sp_test_image_ulysses=2_ring=2.png
+```
+
+**Verify:**
+
+1. **Correctness:** Output should be identical across all `sp_size` values
+2. **Speed:** Throughput should remain stable or improve (especially for large sequences)
+3. **Logs:** Check for any shape mismatch or communication errors
+
+**Test with Tensor Parallel:**
+
+Sequence Parallel can be combined with other parallelism strategies:
+
+```bash
+cd examples/offline_inference/text_to_image
+python text_to_image.py \
+    --model Your-org/your-model \
+    --prompt "a cup of coffee on the table" \
+    --num-inference-steps 50 \
+    --ulysses-degree 2 \
+    --tensor-parallel-size 2 \
+    --output sp_test_image_ulysses=2_tp=2.png
+```
+
+---
+
+## Troubleshooting
+
+### Issue: Shape mismatch errors
+
+**Symptoms:** `RuntimeError: shape mismatch` during forward pass.
+
+**Causes & Solutions:**
+
+- **RoPE dimension mismatch:**
+
+**Problem:** RoPE embeddings not sharded, but hidden_states is sharded.
+
+**Solution:** Shard RoPE outputs in `_sp_plan`:
+```python
+_sp_plan = {
+    "rope": {
+        0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+        1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+    },
+    ...
+}
+```
+
+- **Sequence Length not divisible by sp_size:**
+
+**Problem:** `SequenceParallelInput(auto_pad=False)` - auto_pad should be True to enable automatic sequence padding.
+
+**Solution:** In `SequenceParallelInput`, set `auto_pad=True` and add attention mask support.
+
+> **Experimental Feature:** `auto_pad=True` is an experimental feature and may be changed in the future. We plan to improve this solution to involve minimal changes to model files. More details are [here](https://github.com/vllm-project/vllm-omni/issues/1324).
+
+**Constraints of auto_pad:**
+
+| Constraint | Description |
+|------------|-------------|
+| **Attention Backend Compatibility** | The attention backends must support `attention_mask`. Currently only `TORCH_SDPA` and `FLASH_ATTN` (default for diffusion models) are supported. |
+| **Ring Attention Limitation** | Ring attention does not support `attention_mask`. Therefore, when using `auto_pad=True`, the combination of Ulysses + Ring attention is not feasible. |
+
+1. Enable `auto_pad=True` for all sequence-dimension inputs in `_sp_plan`:
+```python
+_sp_plan = {
+    "rope": {
+        0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
+        1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
+    },
+    "blocks.0": {
+        "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True)
+    },
+    ...
+}
+```
+
+2. Create attention mask dynamically when padding is applied:
+```python
+from vllm_omni.diffusion.forward_context import get_forward_context
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+# In model forward(), before transformer blocks:
+hidden_states_mask = None
+ctx = get_forward_context()
+if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
+    batch_size = hidden_states.shape[0]
+    padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
+    hidden_states_mask = torch.ones(batch_size, padded_seq_len, dtype=torch.bool, device=hidden_states.device)
+    hidden_states_mask[:, ctx.sp_original_seq_len:] = False
+
+# Pass mask to attention layers
+attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask) if hidden_states_mask is not None else None
+output = self.attn(query, key, value, attn_metadata)
+```
+
+**Important Quality Considerations:**
+
+While `auto_pad` enables generation for irregular resolutions, be aware of potential quality impacts:
+
+| Aspect | Impact |
+|--------|--------|
+| **Training Distribution** | Models perform best on aspect ratios seen during training (e.g., 1:1, 16:9, 4:3). Unusual ratios like 700x400 (1.75:1) may produce lower quality results. |
+| **Padding Overhead** | Padded positions consume compute even when masked. For best efficiency, prefer resolutions divisible by `sp_size`. |
+
+**Recommendations for users:**
+- Use standard aspect ratios when possible (e.g., 768x432 for 16:9 instead of 700x400)
+- Ensure post-patch dimensions are divisible by `sp_size` for optimal quality
+- Test generation quality when using unusual resolutions
+
+### Issue: Inline operations not sharded
+
+**Symptoms:** Some tensors remain full-sized, not sharded.
+
+**Causes & Solutions:**
+
+- **Operations happen inline in `forward()`, not at module boundaries:**
+
+**Problem:**
+```python
+def forward(self, x, cap):
+    unified = torch.cat([x, cap], dim=1)  # ← Inline operation!
+    # _sp_plan can't hook this
+```
+
+**Solution:** Extract into submodule:
+```python
+class ConcatModule(nn.Module):
+    def forward(self, x, cap):
+        return torch.cat([x, cap], dim=1)
+
+class MyModel(nn.Module):
+    _sp_plan = {
+        "concat": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
+        },
+        ...
+    }
+    def __init__(self):
+        self.concat = ConcatModule()  # Now hookable!
+
+    def forward(self, x, cap):
+        unified = self.concat(x, cap)  # ← Can be sharded via _sp_plan
+```
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Pattern | Notes |
+|-------|------|---------|-------|
+| **LongCat** | `vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py` | Dual-stream | Text components replicated, image components sharded |
+| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Dual-stream + preprocessing | auto_pad, separate RoPE |
+| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Dual-Transformer + RoPE | Video transformer |
+| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Unified sequence | Concatenated input |
+| **SP Plan Types** | `vllm_omni/diffusion/distributed/sp_plan.py` | Type definitions | SequenceParallelInput/Output |
+| **Hook Implementation** | `vllm_omni/diffusion/hooks/sequence_parallel.py` | Hook mechanics | How hooks work |
+| **Tests** | `tests/diffusion/distributed/test_sp_plan_hooks.py` | Test examples | Validation patterns |
+
+---
+
+## Summary
+
+Adding Sequence Parallel support to a transformer:
+
+1. ✅ **Choose approach** - Use `_sp_plan` for standard cases, intrusive modification for complex cases
+2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered? And which module boundaries need to be moved to facilitate this?
+3. ✅ **Extract inline operations** - Move `torch.cat`, `pad_sequence`, etc. to submodules
+4. ✅ **Define `_sp_plan`** - Declare shard/gather points as class attribute
+5. ✅ **Use `auto_pad` for variable lengths** - Support non-uniform sequences
+6. ✅ **Test** - Verify with different `ulysses_degree` and `ring_degree` combinations
@@ -0,0 +1,459 @@
+# TeaCache
+
+This section describes how to add TeaCache to a diffusion transformer model. We use the Qwen-Image transformer as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Customization](#customization)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is TeaCache?
+
+TeaCache speeds up diffusion inference by caching transformer block computations when consecutive timesteps are similar. It provides **1.5x-2.0x speedup** with minimal quality loss.
+
+The core insight is that the modulated input (after normalization and timestep conditioning) changes gradually across timesteps. By measuring the L1 distance between consecutive modulated inputs and comparing it to a threshold, TeaCache decides whether to execute the full transformer blocks or reuse the cached residual from the previous step.
+
+vLLM-omni provides a **hook-based** TeaCache system that requires **zero changes to model code**. The hook completely intercepts the transformer's forward pass and implements adaptive caching transparently. This design allows easy integration with any transformer model by simply writing an extractor function.
+
+### Architecture
+
+The TeaCache system consists of three main components:
+
+| Component | Purpose | Location |
+|-----------|---------|----------|
+| [`CacheContext`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.CacheContext) | Dataclass containing model-specific information for caching | `vllm_omni/diffusion/cache/teacache/context.py` |
+| [`EXTRACTOR_REGISTRY`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/teacache/extractors/#vllm_omni.diffusion.cache.teacache.extractors.EXTRACTOR_REGISTRY) | Maps transformer class names to extractor functions | `vllm_omni/diffusion/cache/teacache/extractors.py` |
+| [`TeaCacheConfig`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.TeaCacheConfig) | Configuration including thresholds and polynomial coefficients | `vllm_omni/diffusion/cache/teacache/config.py` |
+
+The hook handles all caching logic automatically, including:
+
+- CFG-aware state management (separate states for positive/negative branches)
+- CFG-parallel compatibility
+- L1 distance computation with polynomial rescaling
+- Residual caching and reuse
+
+
+---
+
+## Step-by-Step Implementation
+
+To add TeaCache support for a new model, you need to:
+
+1. Write an **extractor function** that returns a `CacheContext` object
+2. Register the extractor in the `EXTRACTOR_REGISTRY`
+3. Add model-specific polynomial coefficients to `TeaCacheConfig`
+
+### Step 1: Model-Specific Preprocessing
+
+Extract and process model inputs. This typically involves:
+- Embedding image/latent inputs
+- Processing text encoder outputs (if dual-stream)
+- Creating timestep embeddings
+- Computing positional embeddings
+
+**Example (Qwen-Image):**
+
+```python
+def extract_qwen_context(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    encoder_hidden_states_mask: torch.Tensor,
+    timestep: torch.Tensor,
+    img_shapes: torch.Tensor,
+    txt_seq_lens: torch.Tensor,
+    guidance: torch.Tensor | None = None,
+    **kwargs: Any,
+) -> CacheContext:
+    # Validate model structure
+    if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
+        raise ValueError("Module must have transformer_blocks")
+
+    # Preprocessing: embed inputs
+    hidden_states = module.img_in(hidden_states)
+    timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype)
+    encoder_hidden_states = module.txt_norm(encoder_hidden_states)
+    encoder_hidden_states = module.txt_in(encoder_hidden_states)
+
+    # Create timestep embedding
+    if guidance is not None:
+        guidance = guidance.to(hidden_states.dtype) * 1000
+    temb = (
+        module.time_text_embed(timestep, hidden_states)
+        if guidance is None
+        else module.time_text_embed(timestep, guidance, hidden_states)
+    )
+
+    # Compute position embeddings
+    image_rotary_emb = module.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
+```
+
+### Step 2: Extract Modulated Input
+
+The modulated input is used for cache decisions. Extract it from the **first transformer block** after normalization and modulation.
+
+**Example (Qwen-Image):**
+
+```python
+    # Extract modulated input from first transformer block
+    block = module.transformer_blocks[0]
+    img_mod_params = block.img_mod(temb)
+    img_mod1, _ = img_mod_params.chunk(2, dim=-1)
+    img_modulated, _ = block.img_norm1(hidden_states, img_mod1)
+```
+
+**Key Points:**
+
+- Use the **first block** to extract modulated input early
+- Apply the same normalization and modulation as the actual forward pass
+- The tensor should represent the processed features that will change across timesteps
+
+### Step 3: Define Transformer Execution
+
+Create a callable that executes all transformer blocks. This encapsulates the main computation loop.
+
+**Example (Qwen-Image dual-stream):**
+
+```python
+    def run_transformer_blocks():
+        """Execute all Qwen transformer blocks."""
+        h = hidden_states
+        e = encoder_hidden_states
+
+        for block in module.transformer_blocks:
+            e, h = block(
+                hidden_states=h,
+                encoder_hidden_states=e,
+                encoder_hidden_states_mask=encoder_hidden_states_mask,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+            )
+        return (h, e)  # Return both image and text hidden states
+```
+
+**Example (Single-stream model like Flux):**
+
+```python
+    def run_transformer_blocks():
+        """Execute all Flux transformer blocks."""
+        h = hidden_states
+
+        for block in module.transformer_blocks:
+            h = block(h, temb=temb)
+        return (h,)  # Return only image hidden states
+```
+
+**Key Points:**
+
+- Return format:
+- For single-stream models: return `(hidden_states,)`
+- For dual-stream models: return `(hidden_states, encoder_hidden_states)`
+
+### Step 4: Define Postprocessing
+
+Create a callable that applies final transformations to produce the model output.
+
+**Example (Qwen-Image):**
+
+```python
+    return_dict = kwargs.get("return_dict", True)
+
+    def postprocess(h):
+        """Apply Qwen-specific output postprocessing."""
+        h = module.norm_out(h, temb)
+        output = module.proj_out(h)
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+```
+
+### Step 5: Return CacheContext
+
+Package all information into a `CacheContext` object.
+
+```python
+    return CacheContext(
+        modulated_input=img_modulated,
+        hidden_states=hidden_states,
+        encoder_hidden_states=encoder_hidden_states,  # or None for single-stream
+        temb=temb,
+        run_transformer_blocks=run_transformer_blocks,
+        postprocess=postprocess,
+    )
+```
+
+**CacheContext Fields:**
+
+| Field | Type | Purpose |
+|-------|------|---------|
+| `modulated_input` | `torch.Tensor` | Tensor used for cache decision (similarity comparison) |
+| `hidden_states` | `torch.Tensor` | Current hidden states (will be modified by caching) |
+| `encoder_hidden_states` | `torch.Tensor | None` | Encoder states for dual-stream models, `None` for single-stream |
+| `temb` | `torch.Tensor` | Timestep embedding tensor |
+| `run_transformer_blocks` | `Callable[[], tuple]` | Executes transformer blocks, returns `(hidden_states, [encoder_hidden_states])` |
+| `postprocess` | `Callable[[torch.Tensor], Any]` | Applies final transformations to produce model output |
+| `extra_states` | `dict | None` | Optional dict for additional model-specific state |
+
+### Step 6: Register the Extractor
+
+Add your extractor to the `EXTRACTOR_REGISTRY` in `vllm_omni/diffusion/cache/teacache/extractors.py`:
+
+```python
+EXTRACTOR_REGISTRY: dict[str, Callable] = {
+    "QwenImageTransformer2DModel": extract_qwen_context,
+    "Bagel": extract_bagel_context,
+    "ZImageTransformer2DModel": extract_zimage_context,
+    "YourModelTransformer2DModel": extract_your_model_context,  # Add here
+}
+```
+
+**Key:** Use the transformer class name (`module.__class__.__name__`)
+
+### Step 7: Add Model Coefficients
+
+Add polynomial rescaling coefficients to `vllm_omni/diffusion/cache/teacache/config.py`:
+
+```python
+_MODEL_COEFFICIENTS = {
+    "QwenImageTransformer2DModel": [
+        -4.50000000e02,
+        2.80000000e02,
+        -4.50000000e01,
+        3.20000000e00,
+        -2.00000000e-02,
+    ],
+    "YourModelTransformer2DModel": [  # Add your model's coefficients
+        # 5 polynomial coefficients (can reuse similar model's coefficients initially)
+    ],
+}
+```
+
+
+**Initial approach:** Start with coefficients from a similar model architecture, then tune empirically following [Customization](#customization) section.
+
+---
+
+## Customization
+
+### Coefficient Estimation
+
+While you can start with coefficients from a similar model architecture, estimating custom coefficients for your specific model typically improves TeaCache performance.
+
+**Why Estimate Coefficients?**
+
+The polynomial coefficients rescale L1 distances between consecutive modulated inputs to better predict when cached residuals can be reused. Model-specific coefficients account for:
+
+- Architecture differences (layer count, hidden size, attention patterns)
+- Training data characteristics
+- Noise prediction behavior across timesteps
+
+| Approach | Performance | Effort |
+|----------|-------------|--------|
+| Using defaults from similar model | Within 5-10% of optimal | Low |
+| Estimating custom coefficients | Best performance | Medium |
+
+#### Implement Data Collection Adapter
+
+Add an adapter in `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py`:
+
+```python
+class YourModelAdapter:
+    """Adapter for coefficient estimation on your model."""
+
+    @staticmethod
+    def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any:
+        """Load your diffusion pipeline."""
+        from your_model_package import YourModelPipeline
+
+        pipeline = YourModelPipeline.from_pretrained(
+            model_path,
+            torch_dtype=dtype,
+        )
+        pipeline = pipeline.to(device)
+        return pipeline
+
+    @staticmethod
+    def get_transformer(pipeline: Any) -> tuple[Any, str]:
+        """Extract transformer from pipeline."""
+        return pipeline.transformer, "YourTransformer2DModel"
+
+    @staticmethod
+    def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
+        """Install data collection hook on transformer."""
+        from vllm_omni.diffusion.hooks import HookRegistry
+
+        registry = HookRegistry.get_or_create(transformer)
+        registry.register_hook(hook._HOOK_NAME, hook)
+
+
+# Register your adapter
+_MODEL_ADAPTERS["YourModel"] = YourModelAdapter
+```
+
+#### Collect Data and Estimate
+
+```python
+from vllm_omni.diffusion.cache.teacache.coefficient_estimator import (
+    TeaCacheCoefficientEstimator,
+)
+from datasets import load_dataset
+from tqdm import tqdm
+
+# Initialize estimator
+estimator = TeaCacheCoefficientEstimator(
+    model_path="/path/to/your/model",
+    model_type="YourModel",
+)
+
+# Load diverse prompts (paper recommends ~70 prompts)
+dataset = load_dataset("nateraw/parti-prompts", split="train")
+prompts = dataset["Prompt"][:70]
+
+# Collect data
+for prompt in tqdm(prompts, desc="Collecting data"):
+    estimator.collect_from_prompt(prompt=prompt, num_inference_steps=50)
+
+# Estimate coefficients
+coeffs = estimator.estimate(poly_order=4)
+print(f"Estimated coefficients: {coeffs.tolist()}")
+```
+
+**Data Statistics Guide:**
+
+| Metric | Good Range | Warning Signs |
+|--------|------------|---------------|
+| **Count** | 2000-5000+ | < 1000: too few prompts |
+| **Input Diffs (x)** | 0.01-0.10 | Very small (<0.001): model may not modulate properly |
+| **Output Diffs (y)** | Should correlate with x | No correlation: check extractor |
+| **Coefficient magnitude** | -1e6 to 1e6 | > 1e8: numerical instability |
+
+---
+
+## Testing
+
+After adding TeaCache support, test with:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="your-model-name",
+    cache_backend="tea_cache",
+    cache_config={
+        "rel_l1_thresh": 0.2,
+        "coefficients": [1.33e6, -1.69e5, 7.95e3, -1.64e2, 1.26],  # Your coefficients
+    }
+)
+
+images = omni.generate(
+    "a beautiful landscape",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Verify:**
+
+1. **Check logs** - Look for TeaCache initialization messages
+2. **Compare performance** - Measure speedup vs baseline (expect 1.5x-2.0x)
+3. **Verify output quality** - Visually compare cached vs uncached outputs (should be nearly identical)
+
+See more detailed examples in [user guide for teacache](../../user_guide/diffusion/teacache.md).
+
+---
+
+## Troubleshooting
+
+### Issue: "Unknown model type"
+
+**Symptoms:** Error message indicating the model type is not recognized when enabling TeaCache.
+
+**Causes & Solutions:**
+
+- **Extractor not registered:**
+
+**Problem:** The transformer class name doesn't exist in `EXTRACTOR_REGISTRY`.
+
+**Solution:** Check the class name and add to registry:
+```python
+# Check transformer class name
+print(pipeline.transformer.__class__.__name__)
+
+# Add to EXTRACTOR_REGISTRY
+EXTRACTOR_REGISTRY["YourTransformer2DModel"] = extract_your_context
+```
+
+- **Transformer class name mismatch:**
+
+**Solution:** Ensure the registry key matches exactly with `module.__class__.__name__`.
+
+### Issue: "Cannot find coefficients"
+
+**Symptoms:** Error when initializing TeaCache about missing model coefficients.
+
+**Causes & Solutions:**
+
+- **Missing coefficients in config:**
+
+**Solution:** Add coefficients to `_MODEL_COEFFICIENTS` in `config.py`, or pass custom coefficients:
+```python
+omni = Omni(
+    model="your-model",
+    cache_backend="tea_cache",
+    cache_config={"coefficients": [1.0, -0.5, 0.1, -0.01, 0.001]}
+)
+```
+
+### Issue: Quality Degradation
+
+**Symptoms:** Output images look noticeably different or have artifacts compared to baseline.
+
+**Causes & Solutions:**
+
+- **Threshold too high:**
+
+**Problem:** `rel_l1_thresh` is too aggressive, causing cache reuse when outputs differ significantly.
+
+**Solution:** Lower the threshold:
+```python
+cache_config={"rel_l1_thresh": 0.1}  # Try 0.1-0.2
+```
+
+- **Coefficients not tuned:**
+
+**Solution:** Estimate model-specific coefficients using the coefficient estimation process described above.
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Pattern | Notes |
+|-------|------|---------|-------|
+| **Qwen-Image** | `vllm_omni/diffusion/cache/teacache/extractors.py` | Dual-stream | `extract_qwen_context` |
+| **Bagel** | `vllm_omni/diffusion/cache/teacache/extractors.py` | Omni model | `extract_bagel_context` |
+| **TeaCache Core** | `vllm_omni/diffusion/cache/teacache/` | Base implementation | Hook and config |
+| **Coefficient Estimator** | `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py` | Estimation tool | Adapter pattern |
+
+---
+
+## Summary
+
+Adding TeaCache support:
+
+1. ✅ **Write extractor** - Create function returning `CacheContext` with model-specific preprocessing
+2. ✅ **Register extractor** - Add to `EXTRACTOR_REGISTRY` with transformer class name
+3. ✅ **Add coefficients** - Add polynomial coefficients to `_MODEL_COEFFICIENTS`
+4. ✅ **Test** - Verify with `cache_backend="tea_cache"`
@@ -0,0 +1,279 @@
+# Tensor Parallel
+
+This section describes how to add Tensor Parallel (TP) to a diffusion transformer model. We use the Z-Image transformer as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Troubleshooting](#troubleshooting)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is Tensor Parallel?
+
+Tensor Parallel (TP) is a model parallelism technique that **shards model weights** across multiple GPUs. Each GPU holds only a portion of the model's parameters and computes only part of each layer's output.
+
+Diffusion transformers contain large attention and MLP layers. We can use Tensor Parallel to shard the model dimension across multiple GPUs, allowing larger models to fit in memory while achieving near-linear speedup.
+
+### Architecture
+
+The Tensor Parallel implementation relies vLLM's Parallel Layers:
+
+[vLLM Parallel Layers API Reference](https://docs.vllm.ai/en/latest/contributing/model/basic/?h=column#3-optional-implement-tensor-parallelism-and-quantization-support)
+
+**Parallel Layer Types:**
+
+| Layer Type | Purpose | Weight Partitioning |
+|------------|---------|---------------------|
+| `ColumnParallelLinear` | First FFN layer, separated QKV | Columns (output dimension) |
+| `RowParallelLinear` | Second FFN layer, attention output | Rows (input dimension) |
+| `QKVParallelLinear` | Multi-head/grouped-query attention QKV | Handles head replication automatically |
+| `ReplicatedLinear` | Layers that shouldn't be sharded | No partitioning (replicated) |
+
+---
+
+## Step-by-Step Implementation
+
+
+### Step 1: Identify Linear Layers
+
+Find all `nn.Linear` layers in your transformer that need to be sharded.
+
+**Key questions:**
+- Which layers should be column parallel (weight split by columns)?
+- Which layers should be row parallel (weight split by rows)?
+
+### Step 2: Replace Linear Layers with Parallel Equivalents
+
+Replace `nn.Linear` with parallel layers from `vllm.model_executor.layers.linear`.
+
+**Example (MLP Block - Up-Down Pattern):**
+
+```python
+class FeedForward(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int):
+        super().__init__()
+        # Column parallel: weight split by columns [hidden_dim/N, dim]
+        self.w1 = ColumnParallelLinear(
+            dim,
+            hidden_dim,
+            bias=False,
+            return_bias=False,
+        )
+        self.act = nn.GELU()
+
+        self.w2 = RowParallelLinear(
+            hidden_dim,
+            dim,
+            bias=False,
+            input_is_parallel=True,  # Input already sharded from w1
+            return_bias=False,
+        )
+
+    def forward(self, x):
+        # x: [batch, seq, dim] (replicated on all GPUs)
+        # w1 outputs sharded [batch, seq, hidden_dim/N]
+        x = self.w1(x)
+        # act operates on sharded tensors (no communication)
+        x = self.act(x)
+        # w2 outputs full dim [batch, seq, dim] via all-reduce
+        x = self.w2(x)
+        return x
+```
+
+**Example (Attention - QKV-Out Pattern):**
+
+```python
+from vllm_omni.diffusion.attention.layer import Attention
+class YourModelAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, num_kv_heads: int):
+        super().__init__()
+        self.head_dim = dim // num_heads
+
+        # Column parallel: QKV weight split by columns
+        # Each GPU gets num_heads/N heads
+        self.to_qkv = QKVParallelLinear(
+            hidden_size=dim,
+            head_size=self.head_dim,
+            total_num_heads=num_heads,
+            total_num_kv_heads=num_kv_heads,
+            bias=False,
+            return_bias=False,
+        )
+
+        # Row parallel: output weight split by rows
+        self.to_out = RowParallelLinear(
+            dim,
+            dim,
+            bias=False,
+            input_is_parallel=True,  # Input sharded from attention
+            return_bias=False,
+        )
+
+        self.attn = Attention(
+            num_heads=self.to_qkv.num_heads, # Each GPU gets num_heads/N heads
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            num_kv_heads=self.to_qkv.num_kv_heads,
+        )
+
+    def forward(self, x):
+        # x: [batch, seq, dim] (replicated)
+        # to_qkv outputs sharded [batch, seq, (q+k+v) * head_dim/N]
+        qkv = self.to_qkv(x)
+        # Split into Q, K, V (each sharded on heads)
+        q, k, v = qkv.split([...], dim=-1)
+        # Attention computed independently on each GPU
+        out = self.attn(q, k, v)
+        # to_out all-reduces to full dim
+        out = self.to_out(out)
+        return out
+```
+
+**Key Points:**
+
+- `ColumnParallelLinear` → `RowParallelLinear` is the standard pairing
+- Set `input_is_parallel=True` on `RowParallelLinear` when input comes from `ColumnParallelLinear`
+- Use `QKVParallelLinear` for attention projections (handles head replication automatically)
+
+### Step 3: Validate TP Constraints
+
+For correct TP operation, these dimensions **must be divisible** by `tensor_parallel_size`:
+
+| Dimension | Reason | Example Error |
+|-----------|--------|---------------|
+| `num_heads` | Heads sharded by QKVParallelLinear | `num_heads=30, tp=4` ❌ (30 % 4 ≠ 0) |
+| `num_kv_heads` | KV heads sharded by QKVParallelLinear | `num_kv_heads=30, tp=4` ❌ (30 % 4 ≠ 0) |
+
+---
+
+## Testing
+
+After adding Tensor Parallel support, test with:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+parallel_config = DiffusionParallelConfig(tensor_parallel_size=2)
+omni = Omni(model="your-model-name", parallel_config=parallel_config)
+
+output = omni.generate(
+    "a cup of coffee on the table",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+**Or via command line:**
+
+```bash
+cd examples/offline_inference/text_to_image
+python text_to_image.py \
+    --model Your-org/your-model \
+    --prompt "a cup of coffee on the table" \
+    --negative-prompt "ugly, unclear" \
+    --cfg-scale 4.0 \
+    --num-inference-steps 50 \
+    --output "tp_enabled.png" \
+    --tensor-parallel-size 2
+```
+
+**Verify:**
+
+1. Check the `e2e_time_ms` in the log for speedup
+2. Compare generated image quality with TP disabled
+3. Verify memory usage is reduced proportionally
+4. Record comparison results in your PR
+
+---
+
+## Troubleshooting
+
+### Issue: TP not activating
+
+**Symptoms:** Model runs on single GPU, no memory savings or speedup.
+
+**Causes & Solutions:**
+
+- **Still using `nn.Linear`:**
+
+**Problem:** Linear layers not replaced with parallel equivalents.
+
+**Solution:** Replace with parallel layers:
+```python
+# ❌ BAD
+self.proj = nn.Linear(dim, dim)
+
+# ✅ GOOD
+self.proj = RowParallelLinear(dim, dim, input_is_parallel=True)
+```
+
+### Issue: Dimension mismatch errors
+
+**Symptoms:** `RuntimeError: shape mismatch` during forward pass.
+
+**Causes & Solutions:**
+
+- **Missing `input_is_parallel=True`:**
+
+**Problem:** RowParallelLinear expects sharded input but receives full tensor.
+
+**Solution:** Set `input_is_parallel=True` when input comes from ColumnParallelLinear:
+```python
+# ✅ GOOD: Correct pairing
+self.w1 = ColumnParallelLinear(dim, hidden_dim, return_bias=False,)
+self.w2 = RowParallelLinear(
+    hidden_dim,
+    dim,
+    input_is_parallel=True,  # Input sharded from w1
+    return_bias=False,
+)
+```
+
+- **Incorrect split dimensions:**
+
+**Problem:** QKV split sizes don't match sharded dimensions.
+
+**Solution:** Use `self.to_qkv.num_heads` (local heads per GPU):
+```python
+# ❌ BAD: Uses total heads
+q_size = self.total_num_heads * self.head_dim
+
+# ✅ GOOD: Uses local heads
+q_size = self.to_qkv.num_heads * self.head_dim
+```
+
+---
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Pattern | Notes |
+|-------|------|---------|-------|
+| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Standard TP | Full implementation with validation |
+| **FLUX** | `vllm_omni/diffusion/models/flux/flux_transformer.py` | Dual-stream | Image + text streams |
+| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Standard TP | With RoPE |
+| **TP Tests** | `tests/e2e/offline_inference/test_zimage_parallelism.py` | E2E testing | TP correctness and performance |
+| **Constraint Tests** | `tests/diffusion/models/z_image/test_zimage_tp_constraints.py` | Unit testing | Validation logic |
+
+---
+
+## Summary
+
+Adding Tensor Parallel support to a transformer:
+
+1. ✅ **Identify linear layers** - Which layers should be sharded?
+2. ✅ **Replace with parallel layers** - Use QKVParallelLinear, ColumnParallelLinear, RowParallelLinear
+3. ✅ **Validate TP constraints** - Ensure dimensions divisible by TP size
+4. ✅ **Test** - Verify with `tensor_parallel_size=N`, check memory, speed, and quality
@@ -0,0 +1,289 @@
+# VAE Patch Parallelism
+
+This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
+We use **Qwen-Image** as the reference implementation.
+
+---
+
+## Table of Contents
+
+- [Overview](#overview)
+- [Step-by-Step Implementation](#step-by-step-implementation)
+- [Testing](#testing)
+- [Reference Implementations](#reference-implementations)
+- [Summary](#summary)
+
+---
+
+## Overview
+
+### What is Vae Patch parallel?
+
+**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
+
+ Split into multiple spatial tiles
+
+ Distributed across multiple ranks
+
+ Decoded in parallel
+
+ Merged to reconstruct the final output
+
+This approach:
+
+ Distributes computation across multiple devices
+
+ Reduces peak memory usage per device
+
+ Accelerates decoding latency
+
+### Architecture
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
+
+The executor is model-agnostic and accepts three function parameters:
+
+ split – Partition the latent into tiles
+
+ exec – Decode a single tile
+
+ merge – Combine decoded tiles into the final output
+
+#### Execution Flow
+
+ Call split(z) to generate a list of TileTask and a GridSpec
+
+ Dispatch tasks across ranks using workload-based balancing
+
+ Each rank executes exec(task) on its assigned tiles
+
+ Gather decoded tile results to rank 0
+
+ Rank 0 performs merge(...)
+
+ (Optional) Broadcast final result to all ranks
+
+This design separates:
+
+ Distributed execution logic
+
+ Model-specific tiling and merging logic
+
+#### Why split / exec / merge is necessary?
+
+The latent tensor cannot be arbitrarily partitioned.
+
+During decoding:
+
+ Each output pixel may depend on neighboring pixels
+
+ The receptive field is model-dependent
+
+Therefore:
+
+ Tiles must include overlap
+
+ Merge must perform blending to avoid seams
+
+## Step-by-Step Implementation
+
+### Step 1: Implement DistributedAutoencoderKLQwenImage
+`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
+
+
+```
+class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+```
+**Key points**:
+ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
+ Call init_distributed() after loading weights
+
+### Step 2: Implement split/exec/merge
+Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`:
+```
+class GridSpec:
+    split_dims: tuple[int, ...]  # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W))
+    grid_shape: tuple[int, ...]  # Tile grid layout (num_rows, num_cols)
+    tile_spec: dict = field(default_factory=dict) # Metadata required for merging
+    output_dtype: torch.dtype | None = None # Final output dtype
+```
+```
+class TileTask:
+    tile_id: int # task id
+    grid_coord: tuple[int, ...]  # Tile position in grid
+    tensor: torch.Tensor | list[torch.Tensor]  # The tile tensor
+    workload: int | float = 1 # Used for load balancing (e.g., tile area)
+```
+And tiled base split/exec/merge as follow:
+```
+def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
+    # mostly copy from AutoencoderKL
+    _, _, num_frames, height, width = z.shape
+    sample_height = height * self.spatial_compression_ratio
+    sample_width = width * self.spatial_compression_ratio
+
+    tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
+    tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
+    tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
+    tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
+
+    blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+    blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+    # Split z into overlapping tiles and decode them separately.
+    # The tiles have an overlap to avoid seams between tiles.
+    tiletask_list = []
+    for i in range(0, height, tile_latent_stride_height):
+        for j in range(0, width, tile_latent_stride_width):
+            time_list = []
+            for k in range(num_frames):
+                self._conv_idx = [0]
+                tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
+                time_list.append(tile)
+            tiletask_list.append(
+                TileTask(
+                    len(tiletask_list),
+                    (i // tile_latent_stride_height, j // tile_latent_stride_width),
+                    time_list,
+                    workload=time_list[0].shape[3] * time_list[0].shape[4],
+                )
+            )
+    tile_spec = {
+        "sample_height": sample_height,
+        "sample_width": sample_width,
+        "blend_height": blend_height,
+        "blend_width": blend_width,
+    }
+    grid_spec = GridSpec(
+        split_dims=(3, 4),
+        grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
+        tile_spec=tile_spec,
+        output_dtype=self.dtype,
+    )
+    return tiletask_list, grid_spec
+
+def tile_exec(self, task: TileTask) -> torch.Tensor:
+    """Decode a single latent tile into RGB space."""
+    self.clear_cache()
+    time = []
+    for k in range(len(task.tensor)):
+        self._conv_idx = [0]
+        tile = self.post_quant_conv(task.tensor[k])
+        decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
+        time.append(decoded)
+    result = torch.cat(time, dim=2)
+    return result
+
+def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
+    """Merge decoded tiles into a full image."""
+    grid_h, grid_w = grid_spec.grid_shape
+    result_rows = []
+    self.clear_cache()
+
+    result_rows = []
+    for i in range(grid_h):
+        result_row = []
+        for j in range(grid_w):
+            tile = coord_tensor_map[(i, j)]
+            if i > 0:
+                tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
+            if j > 0:
+                tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
+            result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
+        result_rows.append(torch.cat(result_row, dim=-1))
+    dec = torch.cat(result_rows, dim=3)[
+        :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
+    ]
+    return dec
+```
+
+### Step 3: Override tiled_decode
+We need to override tiled_decode, the main logic is:
+ check distributed is enabled
+ select split/exec/merge
+ Invoke self.distributed_decoder.execute to decode
+```
+def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
+    if not self.is_distributed_enabled():
+        return super().tiled_decode(z, return_dict=return_dict)
+
+    logger.info("Decode run with distributed executor")
+    result = self.distributed_decoder.execute(
+        z,
+        DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
+        broadcast_result=True,
+    )
+    if not return_dict:
+        return (result,)
+
+    return DecoderOutput(sample=result)
+```
+`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0.
+
+### Step 4: Modify Pipeline
+Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage
+```
+class YourModelPipeline(nn.Module):
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        ...
+-       self.vae = AutoencoderKL.from_pretrained(
+-           model, subfolder="vae", local_files_only=local_files_only).to(self.device)
+       self.vae = DistributedAutoencoderKL.from_pretrained(
+           model, subfolder="vae", local_files_only=local_files_only
+       ).to(self.device)
+```
+
+## Testing
+Verify numerical consistency between:
+ vae_patch_parallel_size = 1
+
+ vae_patch_parallel_size = N
+
+Example:
+torch.allclose(output_1, output_n, atol=1e-5)
+
+Testing requirements:
+ Fix random seed
+ Use identical tiling strategy
+
+```python
+m = Omni(
+        model=model_name,
+        vae_use_tiling=True,
+        parallel_config=DiffusionParallelConfig(
+            tensor_parallel_size=2,
+            vae_patch_parallel_size=1, # or 2
+        ),
+    )
+```
+When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead.
+
+## Reference Implementations
+
+Complete examples in the codebase:
+
+| Model | Path | Notes |
+|-------|------|-------|
+| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
+| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
+| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
+
+---
+
+## Summary
+
+Adding Vae Patch Parallel support to diffusion model:
+
+1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
+2. **Change vae model in pipeline to Distributed Vae**
+3. **Test** - Verify with `tensor_parallel_size=N` quality
@@ -47,7 +47,8 @@ If no, check your shell running the ComfyUI process. There may be some error mes
 This extension offers the following nodes based on the output modalities:

 - **Generate Image** for text-to-image and image-to-image tasks
- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
+- **Generate Video** for text-to-video and image-to-video tasks
+- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
 - **TTS** and **TTS Voice Clone** for TTS tasks

 This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
@@ -59,7 +60,7 @@ To build a simple workflow yourself,

 - Drag a generation node onto the canvas.
 - Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
 - If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
    - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
    - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
@@ -0,0 +1,151 @@
+# Custom Pipeline Extension Guide
+
+Transformer already support Custom Pipeline via
+https://github.com/huggingface/diffusers/blob/main/docs/source/en/using-diffusers/custom_pipeline_overview.md
+
+This guide demonstrates how to use the newly added features for extending vLLM-Omni's diffusion pipeline with custom functionality.
+
+## Overview
+
+Three main features enable custom pipeline extension:
+
+1. **`WorkerWrapperBase`**: A wrapper class that enables dynamic worker extension with custom functionality
+2. **`load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines
+3. **`CustomPipelineWorkerExtension`**: An extension class that enables pipeline re-initialization with custom implementations
+
+## Features
+
+### WorkerWrapperBase
+
+`WorkerWrapperBase` is a wrapper class that creates `DiffusionWorker` instances with optional extension support. It enables dynamic inheritance, allowing you to add custom methods and functionality to workers without modifying the base worker class.
+
+**Key capabilities:**
+- Dynamic worker class extension via `worker_extension_cls`
+- Support for custom pipeline initialization via `custom_pipeline_args`
+- Method delegation to underlying worker
+- Attribute access forwarding
+
+**Location:** `vllm_omni/diffusion/worker/diffusion_worker.py`
+
+### load_format Parameter
+
+The `load_format` parameter controls how diffusion models are loaded. It supports the following values:
+
+- **`"default"`**: Standard model loading using the model registry (default behavior)
+- **`"custom_pipeline"`**: Load a custom pipeline class specified by `custom_pipeline_name`
+- **`"dummy"`**: Skip model loading (useful for testing or when pipeline will be initialized separately)
+
+**Location:** `vllm_omni/diffusion/model_loader/diffusers_loader.py`
+
+### CustomPipelineWorkerExtension
+
+`CustomPipelineWorkerExtension` is a mixin class that extends `DiffusionWorker` with the ability to re-initialize the pipeline with a custom implementation.
+
+**Key method:**
+- `re_init_pipeline(custom_pipeline_args)`: Re-initializes the pipeline with custom arguments, properly cleaning up the old pipeline first
+
+**Location:** `vllm_omni/diffusion/worker/diffusion_worker.py`
+
+## Usage Example
+
+### Step 1: Create a Custom Pipeline
+
+Create a custom pipeline class that extends an existing pipeline. In this example, we extend `QwenImageEditPipeline` to add trajectory tracking:
+
+```python
+# custom_pipeline.py
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline
+import torch
+
+class CustomPipeline(QwenImageEditPipeline):
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+
+    def forward(self, req, prompt=None, negative_prompt=None, **kwargs):
+        # Call parent's forward to get normal output
+        output = super().forward(req=req, prompt=prompt, negative_prompt=negative_prompt, **kwargs)
+
+        # Add custom trajectory data
+        actual_num_steps = req.sampling_params.num_inference_steps or kwargs.get('num_inference_steps', 50)
+        output.trajectory_timesteps = torch.linspace(1000, 0, actual_num_steps, dtype=torch.float32)
+        output.trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32)
+
+        return output
+```
+
+### Step 2: Use the Custom Pipeline with Omni
+
+Initialize the `Omni` engine with custom pipeline configuration:
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# Initialize with custom pipeline
+omni = Omni(
+    model="Qwen/Qwen-Image-Edit",
+    diffusion_load_format="dummy",  # Skip initial loading
+    custom_pipeline_args={
+        "pipeline_class": "custom_pipeline.CustomPipeline"
+    },
+)
+
+# Generate with the custom pipeline
+outputs = omni.generate(
+    ...
+)
+
+# Access custom trajectory data
+output = outputs[0].request_output[0]
+print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
+print(f"Trajectory latents shape: {output.latents.shape}")
+```
+
+### Step 3: Run the Example
+
+The example provided in this directory demonstrates the complete workflow:
+
+```bash
+cd examples/offline_inference/custom_pipeline/image_to_image
+
+# Run with custom pipeline
+python image_edit.py \
+    --model Qwen/Qwen-Image-Edit-2511 \
+    --image cherry_blossom.jpg \
+    --prompt "Let this mascot dance under the moon, surrounded by floating stars" \
+    --output output_image_edit.png \
+    --num-inference-steps 10
+```
+
+## Advanced Usage
+
+### Custom Worker Extension
+
+You can create custom worker extensions to add new methods beyond pipeline re-initialization:
+
+```python
+from typing import Any
+from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
+
+class MyCustomExtension(CustomPipelineWorkerExtension):
+    def custom_method(self):
+        """Your custom worker method."""
+        return "custom_result"
+
+    def another_method(self, data: Any):
+        """Another custom method."""
+        # Access worker internals via self
+        return self.model_runner.some_operation(data)
+
+omni = Omni(
+    model="Qwen/Qwen-Image-Edit",
+    diffusion_load_format="dummy",
+    custom_pipeline_args={
+        "pipeline_class": "custom_pipeline.CustomPipeline"
+    },
+    worker_extension_cls=MyCustomExtension,
+    # Note: worker_extension_cls is an internal parameter
+    # CustomPipelineWorkerExtension will automatically init pipeline when custom_pipeline_args is provided
+)
+```
@@ -5,4 +5,5 @@ vLLM-Omni supports the following hardware platforms:
 - [GPU](gpu.md)
    - [NVIDIA CUDA](gpu.md)
    - [AMD ROCm](gpu.md)
+    - [Intel XPU](gpu.md)
 - [NPU](npu.md)
@@ -18,6 +18,10 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr

    --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:requirements"

+=== "Intel XPU"
+
+    --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:requirements"
+
 ## Set up using Python

 ### Create a new Python environment
@@ -26,17 +30,20 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr

 ### Pre-built wheels

-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).

 === "NVIDIA CUDA"

    --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-wheels"

-
 === "AMD ROCm"

    --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:pre-built-wheels"

+=== "Intel XPU"
+
+    --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-wheels"
+
 [](){ #build-from-source }

 ### Build wheel from source
@@ -49,6 +56,10 @@ Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.1

    --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-wheel-from-source"

+=== "Intel XPU"
+
+    --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-wheel-from-source"
+
 ## Set up using Docker

 ### Pre-built images
@@ -61,8 +72,16 @@ Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.1

    --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:pre-built-images"

+=== "Intel XPU"
+
+    --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-images"
+
 ### Build your own docker image

 === "AMD ROCm"

    --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-docker"
+
+=== "Intel XPU"
+
+    --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-docker"
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**

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

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

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

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

-The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
+The release 0.17.0 of vLLM is based on PyTorch 2.10.0 which requires CUDA 12.9 environment.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -56,11 +56,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.15.0
+git checkout v0.17.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.15.0/vllm-0.15.0-cp38-abi3-manylinux_2_31_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -91,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
    --env "HF_TOKEN=$HF_TOKEN" \
    -p 8091:8091 \
    --ipc=host \
-    vllm/vllm-omni:v0.14.0 \
+    vllm/vllm-omni:v0.16.0 \
    --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
				`@@ -1 +0,0 @@`
				{"id":"1c99f525-0a37-45ba-a28a-7df7c3af66b4","revision":0,"last_node_id":12,"last_link_id":14,"nodes":[{"id":1,"type":"VLLMOmniComprehension","pos":[1191.2177053682556,144.66829928181377],"size":[400,268],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":1},{"localized_name":"video","name":"video","shape":7,"type":"VIDEO","link":2},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":3},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":14},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"output_text","name":"output_text","type":"BOOLEAN","widget":{"name":"output_text"},"link":null},{"localized_name":"output_audio","name":"output_audio","type":"BOOLEAN","widget":{"name":"output_audio"},"link":null},{"localized_name":"use_audio_in_video","name":"use_audio_in_video","type":"BOOLEAN","widget":{"name":"use_audio_in_video"},"link":null}],"outputs":[{"localized_name":"text_response","name":"text_response","type":"STRING","links":[8]},{"localized_name":"audio_response","name":"audio_response","type":"AUDIO","links":[9]}],"properties":{"Node name for S&R":"VLLMOmniComprehension"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen2.5-Omni-7B","",true,true,true]},{"id":3,"type":"LoadVideo","pos":[729.5984141255855,-198.631920454299],"size":[282.798828125,233.0743408203125],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[2]}],"properties":{"Node name for S&R":"LoadVideo"},"widgets_values":["draw.mp4","image"]},{"id":4,"type":"LoadAudio","pos":[729.8037086965753,99.86963519703949],"size":[282.798828125,136],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":5,"type":"VLLMOmniARSampling","pos":[510.3517536642828,658.073751009259],"size":[270,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[12]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":7,"type":"VLLMOmniARSampling","pos":[503.33235181647115,419.34158016181806],"size":[270,178],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[5,13]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":8,"type":"VLLMOmniSamplingParamsList","pos":[820.6056617389042,426.38372037182273],"size":[263.066015625,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"param1","name":"param1","type":"SAMPLING_PARAMS","link":5},{"localized_name":"param2","name":"param2","shape":7,"type":"SAMPLING_PARAMS","link":12},{"localized_name":"param3","name":"param3","shape":7,"type":"SAMPLING_PARAMS","link":13}],"outputs":[{"localized_name":"param list","name":"param list"Line truncated