Initial plan

Signed-off-by: Taichang Zhou <tzhouam@connect.ust.hk>
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: ZJY0516 <zhu.jiangyun@foxmail.com>
Co-authored-by: youkaichao <youkaichao@gmail.com>

.buildkite/pipeline.yml

@@ -9,30 +9,18 @@ steps:
     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: "Upload Nightly Pipeline"
+    depends_on: image-build
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yaml
+    agents:
+      queue: "cpu_queue_premerge"
 
   - label: "Simple Unit Test"
     depends_on: image-build
     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
+      - "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:
@@ -49,24 +37,7 @@ steps:
     timeout_in_minutes: 20
     depends_on: image-build
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_t2i_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 Images API LoRA E2E"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+      - 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" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -118,7 +89,7 @@ steps:
     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
+      - 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:
@@ -135,25 +106,7 @@ steps:
     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
+      - 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:
@@ -186,12 +139,13 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
 
-  - label: "Benchmark Test"
+  - label: "Benchmark&Engine 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
+      - pytest -s -v tests/engine/test_async_omni_engine_abort.py
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -225,58 +179,77 @@ steps:
     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
+      - 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"
+      - 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: "Omni Model Test with H100"
+    timeout_in_minutes: 20
+    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 -m "core_model" --run-level "core_model"
+    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: "Qwen3-TTS E2E Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - export VLLM_LOGGING_LEVEL=DEBUG
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
+    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)"
     timeout_in_minutes: 20

.buildkite/test-amd.yaml

@@ -44,7 +44,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 +56,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 +66,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
@@ -101,7 +101,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
+    - pytest -s -v tests/engine/test_async_omni_engine_abort.py
 
 - label: "Diffusion Image Edit Test"
   timeout_in_minutes: 15

.buildkite/test-merge.yml

@@ -0,0 +1,329 @@
+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: "Simple Unit Test"
+    depends_on: image-build
+    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: image-build
+    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: image-build
+    commands:
+      - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model CPU offloading Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - 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&Engine 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
+      - pytest -s -v tests/engine/test_async_omni_engine_abort.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 -m "advanced_model" --run-level "advanced_model"
+    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

.buildkite/test-nightly.yaml

@@ -1,20 +1,69 @@
 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
+    timeout_in_minutes: 60
     depends_on: image-build
+    if: build.env("NIGHTLY") == "1"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
       - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py
+      - pytest -s -v tests/examples/online_serving/test_qwen3_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: "Omni Model Test"
+    timeout_in_minutes: 60
+    depends_on: image-build
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - export VLLM_LOGGING_LEVEL=DEBUG
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/examples/online_serving/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 Perf Test"
+    timeout_in_minutes: 120
+    depends_on: image-build
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/perf/scripts/run_benchmark.py
     agents:
       queue: "mithril-h100-pool"
     plugins:

.github/PULL_REQUEST_TEMPLATE.md

@@ -12,9 +12,9 @@ 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.
+- [ ] The test plan. Please providing 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 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. **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>
 

.gitignore

@@ -242,3 +242,6 @@ configs/development.yaml
 Dockerfile.dev
 discussion
 tmp_test
+
+# output files
+*.wav

benchmarks/build_dataset/download_process_data_seedtts.md

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

benchmarks/build_dataset/extract_tts_prompts.py

@@ -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()

benchmarks/qwen3-omni/README.md

@@ -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.

benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh

@@ -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)"

benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py

@@ -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

benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh

@@ -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"

docker/Dockerfile.ci

@@ -8,13 +8,38 @@ COPY . .
 
 # Install system dependencies
 RUN apt-get update && \
-    apt-get install -y ffmpeg && \
+    apt-get install -y ffmpeg sox libsox-fmt-all && \
     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]"
+# Install vLLM wheel for the specific commit
+# URL format: https://wheels.vllm.ai/<commit_hash>/
+# Current vLLM wheel commit: 2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
+# Note: --prerelease=allow is required because nightly/commit wheels use
+# pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default,
+# causing it to fall back to the older stable release from PyPI.
+RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
 
+
+RUN uv pip uninstall --system nvidia-cublas-cu12
+# Install vllm-omni into the same uv-managed Python environment used by the base image.
+# Use bash -c so that $(python3 -c ...) is expanded inside the container.
+RUN uv pip install --system --no-cache-dir ".[dev]"
+
+# pplx_kernels is optional; the base image may already ship a version.
+# Reinstall attempt may fail if no compatible wheel exists for the current
+# PyTorch/CUDA combination – allow the build to continue regardless.
+RUN pip install --upgrade --force-reinstall "pplx_kernels" || true
+
+RUN /bin/bash -c 'FLASHINFER_CUDA_TAG="$(python3 -c "import torch; print((torch.version.cuda or \"12.4\").replace(\".\", \"\"))")" && \
+    uv pip install --system --upgrade --force-reinstall \
+      "flashinfer-python==0.6.3" \
+      "flashinfer-cubin==0.6.3" \
+      "flashinfer-jit-cache==0.6.3" \
+      --extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"'
+# Pin nvidia-cublas-cu12 after vllm-omni install so dependency resolution does not overwrite it
+RUN uv pip install --system --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
+RUN uv pip install --system --upgrade --force-reinstall "numpy==2.2.6"
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 
 ENTRYPOINT []

docs/.nav.yml

@@ -8,26 +8,27 @@ 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
       - 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
+      - LoRA Inference(Diffusion): user_guide/examples/offline_inference/lora_inference.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
-      - Qwen3-TTS Offline Inference: user_guide/examples/offline_inference/qwen3_tts.md
+      - 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
     - Online Serving:
       - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
-      - Online LoRA Inference (Diffusion): user_guide/examples/online_serving/lora_inference.md
+      - LoRA Inference(Diffusion): user_guide/examples/online_serving/lora_inference.md
       - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
-      - Qwen3-TTS Online Serving: user_guide/examples/online_serving/qwen3_tts.md
+      - Qwen3-TTS: user_guide/examples/online_serving/qwen3_tts.md
       - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
   - General:
     - usage/*
@@ -42,6 +43,9 @@ nav:
       - 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
       - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
       - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
     - ComfyUI: features/comfyui.md
@@ -54,6 +58,12 @@ nav:
     - contributing/model/README.md
     - contributing/model/adding_omni_model.md
     - contributing/model/adding_diffusion_model.md
+    - Advanced Features:
+      - contributing/features/cfg_parallel.md
+      - contributing/features/sequence_parallel.md
+      - contributing/features/tensor_parallel.md
+      - contributing/features/cache_dit.md
+      - contributing/features/teacache.md
   - CI: contributing/ci
   - Design Documents:
     - design/index.md

docs/api/README.md

@@ -7,17 +7,12 @@ 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.chat_utils.extract_audio_from_video_async][]
 - [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][]
@@ -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
@@ -107,10 +101,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][]

docs/contributing/ci/CI_5levels.md

@@ -0,0 +1,678 @@
+# 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}_expansion.py<br>
+        /tests/e2e/offline_inference/test_{model_name}_expansion.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_request(request_config, request_num=1)
+    ```
+
+    **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_request method 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.

docs/contributing/ci/failures.md

@@ -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.

docs/contributing/ci/tests_markers.md

@@ -5,33 +5,34 @@ 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 *                              |
+| `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, and they will be added by the `@hardware_test` decorator.
 
 ### Example usage for markers
 
@@ -71,10 +72,7 @@ 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
@@ -94,7 +92,6 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
     ```
 - `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"`

docs/contributing/ci/tests_style.md

@@ -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,141 @@ 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
+                                                   └── stage_configs/
+                                                       ├── qwen2_5_omni_ci.yaml
+                                                       ├── qwen3_omni_ci.yaml
+                                                       ├── bagel_*.yaml
+                                                       └── npu/, rocm/, etc.
 ```
 
 
@@ -117,6 +173,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 +202,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 +465,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.

docs/contributing/features/cache_dit.md

@@ -0,0 +1,286 @@
+# Support 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!

docs/contributing/features/cfg_parallel.md

@@ -0,0 +1,249 @@
+# Support 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

docs/contributing/features/sequence_parallel.md

@@ -0,0 +1,405 @@
+# Support 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
+
+**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),
+    }
+```
+
+### 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`:
+```python
+"blocks.0": {
+    "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True)
+}
+```
+
+### 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 |
+|-------|------|---------|-------|
+| **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?
+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

docs/contributing/features/teacache.md

@@ -0,0 +1,459 @@
+# Support 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"`

docs/contributing/features/tensor_parallel.md

@@ -0,0 +1,279 @@
+# Support 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

docs/contributing/metrics.md

@@ -0,0 +1,165 @@
+
+# Metrics vLLM-Omni:
+
+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
+
+### 1. Start the Service with Metrics Logging
+
+```bash
+vllm serve /workspace/models/Qwen3-Omni-30B-A3B-Instruct --omni --port 8014 --log-stats
+```
+
+### 2. Send a Request
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
+```
+
+### 3. 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
+
+| 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
+```
+This matches the table value: `e2e_total_tokens = 5,202`.
+
+### 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**
+
+Sum: 192.581 + 53.315 = **245.896**
+
+The table shows `transfers_total_time_ms = 245.895`, which matches the calculation (difference is due to rounding).

docs/contributing/profiling.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
 

docs/features/custom_pipeline.md

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

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

@@ -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==0.16.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -34,10 +34,10 @@ uv pip install vllm-omni
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.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.16.0 release version of the library
 
 ```bash
-uv pip install vllm==0.15.0 --torch-backend=auto
+uv pip install vllm==0.16.0 --torch-backend=auto
 ```
 
 The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
@@ -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.16.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.16.0/vllm-0.16.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.15.0 --torch-backend=auto
+uv pip install vllm==0.16.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.15.0 --extra-index-url https://wheels.vllm.ai/rocm/0.15.0/rocm700
+uv pip install vllm==0.16.0 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni

docs/models/supported_models.md

@@ -20,6 +20,7 @@ th {
 | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
+| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |

docs/serving/speech_api.md

@@ -0,0 +1,267 @@
+# Speech API
+
+vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation using Qwen3-TTS models.
+
+Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
+
+## Quick Start
+
+### Start the Server
+
+```bash
+# CustomVoice model (predefined speakers)
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
+```
+
+### Generate Speech
+
+**Using curl:**
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "vivian",
+        "language": "English"
+    }' --output output.wav
+```
+
+**Using Python:**
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/speech",
+    json={
+        "input": "Hello, how are you?",
+        "voice": "vivian",
+        "language": "English",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+**Using OpenAI SDK:**
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
+
+response = client.audio.speech.create(
+    model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+    voice="vivian",
+    input="Hello, how are you?",
+)
+
+response.stream_to_file("output.wav")
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/speech
+Content-Type: application/json
+```
+
+### Request Parameters
+
+#### OpenAI Standard Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | The text to synthesize into speech |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
+
+#### vLLM-Omni Extension Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `task_type` | string | "CustomVoice" | TTS task type: CustomVoice, VoiceDesign, or Base |
+| `language` | string | "Auto" | Language (see supported languages below) |
+| `instructions` | string | "" | Voice style/emotion instructions |
+| `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
+
+**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
+
+#### Voice Clone Parameters (Base task)
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `ref_audio` | string | null | Reference audio (URL or base64 data URL) |
+| `ref_text` | string | null | Transcript of reference audio |
+| `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
+
+### Response Format
+
+Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
+
+### Voices Endpoint
+
+```
+GET /v1/audio/voices
+```
+
+Lists available voices for the loaded model.
+
+```json
+{
+    "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
+}
+```
+
+## Examples
+
+### CustomVoice with Style Instruction
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "I am so excited!",
+        "voice": "vivian",
+        "instructions": "Speak with great enthusiasm"
+    }' --output excited.wav
+```
+
+### VoiceDesign (Natural Language Voice Description)
+
+```bash
+# Start server with VoiceDesign model first
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
+```
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello world",
+        "task_type": "VoiceDesign",
+        "instructions": "A warm, friendly female voice with a gentle tone"
+    }' --output designed.wav
+```
+
+### Base (Voice Cloning)
+
+```bash
+# Start server with Base model first
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
+```
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, this is a cloned voice",
+        "task_type": "Base",
+        "ref_audio": "https://example.com/reference.wav",
+        "ref_text": "Original transcript of the reference audio"
+    }' --output cloned.wav
+```
+
+## Supported Models
+
+| Model | Task Type | Description |
+|-------|-----------|-------------|
+| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
+
+## Error Responses
+
+### 400 Bad Request
+
+Invalid parameters:
+
+```json
+{
+    "error": {
+        "message": "Input text cannot be empty",
+        "type": "BadRequestError",
+        "param": null,
+        "code": 400
+    }
+}
+```
+
+### 404 Not Found
+
+Model not found:
+
+```json
+{
+    "error": {
+        "message": "The model `xxx` does not exist.",
+        "type": "NotFoundError",
+        "param": "model",
+        "code": 404
+    }
+}
+```
+
+## Troubleshooting
+
+### "TTS model did not produce audio output"
+
+Ensure you're using the correct model variant for your task type:
+- CustomVoice task → CustomVoice model
+- VoiceDesign task → VoiceDesign model
+- Base task → Base model
+
+### Server Not Running
+
+```bash
+# Check if server is responding
+curl http://localhost:8091/v1/audio/voices
+```
+
+### Out of Memory
+
+If you encounter OOM errors:
+1. Use smaller model variant: `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`
+2. Reduce `--gpu-memory-utilization`
+
+### Unsupported Speaker
+
+Use `/v1/audio/voices` to list available voices for the loaded model.
+
+## Development
+
+Enable debug logging:
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager \
+    --uvicorn-log-level debug
+```

docs/user_guide/diffusion/cache_dit_acceleration.md

@@ -188,8 +188,8 @@ cd examples/offline_inference/text_to_image
 python text_to_image.py \
     --model Qwen/Qwen-Image \
     --prompt "a cup of coffee on the table" \
-    --cache_backend cache_dit \
-    --num_inference_steps 50
+    --cache-backend cache_dit \
+    --num-inference-steps 50
 ```
 
 

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -1,101 +1,143 @@
-# CPU Offloading for Diffusion Model
+# CPU Offloading for Diffusion Models
 
 ## Overview
 
-vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models, allowing you to run larger models on GPUs with limited VRAM:
+vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models:
 
-1. **Model-level (Component) Offloading**: Swaps entire model components (DiT transformer, VAE, encoders) between GPU and CPU.
-2. **Layerwise (Blockwise) Offloading**: Keeps only a single or a few transformer blocks on GPU at a time, with compute - memory copy overlap.
+1. **Model-level (Sequential) Offloading**: Mutual exclusion between DiT model and encoder - only one is on GPU at a time.
+2. **Layerwise (Blockwise) Offloading**: Keeps only one transformer block on GPU at a time with compute-memory overlap.
 
-Both approaches use pinned memory for faster CPU-GPU transfers. For now, the two offloading strategies could not be used at the same time.
+Both strategies use pinned memory for faster CPU-GPU transfers. The strategies are **mutually exclusive** for now - if both are enabled, layerwise takes priority.
 
 
-## Model-level CPU Offloading
+## Model-level (Sequential) Offloading
 
-### Implementation
+### How It Works
 
-CPU offload lets the diffusion worker move large model components between GPU and CPU memory on demand. It keeps the DiT transformer resident on GPU only while it is actively running, and swaps it out when encoders modules need the device. This reduces peak VRAM usage so bigger checkpoints run on smaller GPUs, or multiple requests can share the same GPU.
+Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks:
 
-**Execution Flow**:
-1. Text encoders run on GPU while the DiT transformer is offloaded to CPU.
-2. Before denoising, weights are prefetched back to GPU, honoring pinned-memory copies for speed.
-3. After the diffusion step, the transformer returns to CPU and the process repeats as needed.
+- **When encoders run**: DiT transformer is offloaded to CPU
+- **When DiT runs**: Encoders are offloaded to CPU
+- **VAE**: Stays resident on GPU
 
-Transfers use pinned host buffers, and the worker coordinates swaps via mutex-style hooks so components never compete for memory.
+Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed.
 
-### Configuration
-You can enable CPU offload in two ways:
-
-1. **Python API**: set `enable_cpu_offload=True`.
+### Usage
 
+**Python API:**
 ```python
 from vllm_omni import Omni
 
-if __name__ == "__main__":
-
-    m = Omni(model="Qwen/Qwen-Image",enable_cpu_offload=True)
+m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
 ```
 
-2. **CLI**: pass `--enable-cpu-offload` to the diffusion service entrypoint.
+**CLI:**
+```bash
+vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload
+```
 
 ### Limitations
-- Cold start latency increases for over one minute for some models(e.g., Qwen-Image)
+- Cold start latency increases
+- Adds overhead from CPU-GPU transfers between encoder and denoising phases
+- Support single GPU only for now
 
 
 ## Layerwise (Blockwise) Offloading
 
-### Implementation
-Layerwise offload operates at transformer block granularity, keeping a single transformer block, or a specified number of blocks, on GPU while others stay in CPU memory.
+### How It Works
 
-Unlike full model-wise CPU offload which swaps entire components like DiT and encoders, layerwise offloading applies a sliding window way of loading and offloading weights between gpu and cpu: while block `i` computes, block `i+1` get prefetched asynchronously via pinned memory. In this way, only partial blocks(s) reside on GPU at any moment during inference, so that greatly decrease the memory occupancy.
+Layerwise offloading keeps only one transformer block on GPU at a time.
 
-**Execution Flow**:
+As each block completes, the next block is prefetched to GPU while the current block is freed. The pre and forward hooks utilized by layerwise offloading apply a separate CUDA stream (`copy_stream`) to overlap weight transfer with computation, and retain flattened tensors in pinned CPU memory for block parameters re-materialization. Encoders, VAE, and non-block DiT modules (embeddings, norms) always stay on GPU.
 
-1. During model initialization, all components are loaded to CPU first. Then components other than DiT model(s) in the pipeline, such as VAE and encoders, are moved to GPU. The weights of target transformer blocks are collected as contiguous tensors per layer on CPU with pinned memory; and non-block modules (embeddings, norms, etc) in the DiT model are moved to and stay on GPU.
-2. The first block(s) are transferred to GPU during initialization of `LayerwiseOffloader`, before the first denoising step of the very first request.
-3. As each block executes, the next block prefetches on a separate CUDA stream for compute - memory copy overlap. After execution, the current block is immediately freed from GPU memory.
-4. When the last block completes, the first block prefetches for the next denoising step.
+**Execution Flow:**
 
+| Block | Pre-forward Hook | Forward | Post-forward Hook |
+|-------|------------------|---------|-------------------|
+| block-0 | Prefetch block-1 (async) | Compute block-0 | Free block-0 |
+| block-1 | Prefetch block-2 (async) | Compute block-1 | Free block-1 |
+| ... | ... | ... | ... |
+| block-(n-1) | **Prefetch block-0** (async) | Compute block-(n-1) | Free block-(n-1) |
 
-Example of hook executions of a DiT model with n layers, by default keep a single layer on GPU:
-| Layer (block) idx | forward pre-hook               | forward          | forward post-hook         |
-|-------------------|--------------------------------|------------------|---------------------------|
-| layer-0           | prefetch layer 1 (copy stream) | compute layer 0  | free layer-0 gpu weights  |
-| layer-1           | prefetch layer 2 (copy stream) | compute layer 1  | free layer-1 gpu weights  |
-| layer-2           | prefetch layer 3 (copy stream) | compute layer 2  | free layer-2 gpu weights  |
-| ...               | ...                            | ...              | ...                       |
-| layer-(n-1)       | **prefetch layer 0 (copy stream)** | compute layer (n-1) | free layer (n-1) gpu weights  |
+Each transformer block has a `LayerwiseOffloadHook` that prefetches the next block before forward and frees the current block after forward.
 
+Layerwise offloading is primarily recommended for large **video generation models** where the compute cost per block is high enough to effectively overlap with memory prefetch operations. For example, Wan2.2 T2V and I2V pipelines.
 
-### Configuration
-
-1. **Python API**: set `enable_layerwise_offload=True` and optionally `layerwise_num_gpu_layers`.
+### Usage
 
+**Python API:**
 ```python
 from vllm_omni import Omni
 
-if __name__ == "__main__":
-    m = Omni(
-        model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-        enable_layerwise_offload=True,
-        ...
-    )
+# Text-to-video
+m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_layerwise_offload=True)
+
+# Or image-to-video
+m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True)
 ```
 
-2. **CLI**: pass `--enable-layerwise-offload` and `--layerwise-num-gpu-layers` to the diffusion service entrypoint.
+**CLI:**
+```bash
+# Text-to-video
+vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-layerwise-offload
 
-### Supported Models
+# Or image-to-video
+vllm-omni serve diffusion Wan-AI/Wan2.2-I2V-A14B-Diffusers --enable-layerwise-offload
+```
 
-| Architecture | Models | Example HF Models | DiT Model Cls | Blocks Attr Name |
-|--------------|--------|-------------------|----------|----------|
-| `QwenImagePipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | "transformer_blocks" |
-| `Wan22Pipeline` | Wan2.2 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | "blocks" |
+### To Support a Model
 
-NOTE: Models must define `_layerwise_offload_blocks_attr` class attribute so that the layerwise offloader finds the target transformer blocks.
+Models must define the blocks attribute name for layerwise offloading:
+
+```python
+class WanTransformer3DModel(nn.Module):
+    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
+
+    def __init__(self):
+        self.blocks = nn.ModuleList([...])  # Transformer blocks
+```
 
 ### Limitations
 - Cold start latency increases because of
-    1) components are loaded to CPU first at the very first during initialization,  
+    1) components are loaded to CPU first at the very first during initialization,
     2) weight consolidation and pinning
-- Performance depends on CPU <-> GPU interconnection (e.g., PCIe bandwidth).
+- Performance depends on compute cost and H2D bandwidth as well
 - Support single GPU only for now
+
+
+### Implementation Notes
+
+**Module Discovery**
+
+The offloader automatically discovers pipeline components:
+
+- **DiT modules**: `transformer`, `transformer_2`, `dit`
+- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
+- **VAE**: `vae`
+
+**Hook System**
+
+Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelHook`) to register pre/post forward callbacks on modules, enabling automatic swapping without modifying model code.
+
+**Backend Architecture**
+
+```
+OffloadBackend (base class)
+├── ModelLevelOffloadBackend → uses SequentialOffloadHook
+└── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
+```
+
+Factory function `get_offload_backend()` selects the appropriate backend based on configuration.
+
+
+## Supported Models
+
+| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attr (Layerwise specific) |
+|--------------|----------------|-----------|---------------------|-------------------|-------------|
+| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
+| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
+| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` |
+
+**Notes:**
+- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
+- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attr` pointing to transformer blocks

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -14,23 +14,25 @@ The following parallelism methods are currently supported in vLLM-Omni:
 
 4. [Tensor Parallelism](#tensor-parallelism): Tensor parallelism shards model weights across devices. This can reduce per-GPU memory usage. Note that for diffusion models we currently shard the majority of layers within the DiT.
 
+5. [VAE Patch Parallelism](#vae-patch-parallelism): VAE patch parallelism shards VAE decode spatially across ranks. This can reduce the peak memory of VAE decode and (depending on resolution and communication overhead) speed up VAE decode.
+
 The following table shows which models are currently supported by parallelism method:
 
 ### ImageGen
 
-| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
-|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|
-| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |
-| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |
-| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |
-| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |
-| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |
-| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ❌        |
-| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |
-| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |
+| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel |
+|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|
+| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |
+| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |
+| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |         ❌          |
+| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
+| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
+| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |
+| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |         ✅          |
+| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |
+| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |
+| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |
 
 !!! note "TP Limitations for Diffusion Models"
     We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
@@ -47,7 +49,7 @@ The following table shows which models are currently supported by parallelism me
 
 ### VideoGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel |
+| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel |
 |-------|------------------|------------|---------|--------------------------|
 | **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ |
 
@@ -76,6 +78,45 @@ outputs = omni.generate(
 )
 ```
 
+### VAE Patch Parallelism
+
+VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP).
+
+!!! note "Enablement and feature gate"
+    - VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`).
+    - If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.)
+
+#### Offline Inference
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    parallel_config=DiffusionParallelConfig(
+        tensor_parallel_size=2,
+        vae_patch_parallel_size=2,
+    ),
+    vae_use_tiling=True,
+)
+
+outputs = omni.generate(
+    prompt="a cat reading a book",
+    num_inference_steps=9,
+    width=1024,
+    height=1024,
+)
+```
+
+#### How it works (method selection)
+
+VAE patch parallelism automatically selects between two internal decode methods based on whether diffusers tiling would kick in:
+
+- `_distributed_tiled_decode`: Used when the latent spatial size exceeds `vae.tile_latent_min_size` (i.e., diffusers tiled decode). Each rank decodes a subset of tiles; rank0 gathers and runs the same overlap+blend+stitch logic as diffusers. This matches the single-rank diffusers tiled output.
+
+- `_distributed_patch_decode`: Used when diffusers tiling would not kick in. Each rank decodes a grid patch expanded with a latent-space halo; then rank0 gathers the cropped core patches and stitches them into the full image. This path has no blending and can introduce small numerical differences compared to the non-parallel decode.
+
 ### Sequence Parallelism
 
 #### Ulysses-SP
@@ -233,143 +274,6 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
 | Hybrid Ulysses + Ring  |  2  |  2  |  24.3s | 1.87x |
 
 
-##### How to parallelize a new model
-
-NOTE: "Terminology: SP vs CP"
-    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.
-
----
-
-###### 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.
-
-**Requirements for `forward()` function:**
-
-- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries** (not inline Python operations)
-- If your `forward()` contains inline tensor operations (e.g., `torch.cat`, `pad_sequence`) that need sharding, **extract them into a submodule**
-
-**When to create a submodule:**
-
-```python
-# ❌ BAD: Inline operations - hooks cannot intercept
-def forward(self, x, cap_feats):
-    unified = torch.cat([x, cap_feats], dim=1)  # Cannot be sharded via _sp_plan
-    ...
-
-# ✅ GOOD: Extract into a submodule
-class UnifiedPrepare(nn.Module):
-    def forward(self, x, cap_feats):
-        return torch.cat([x, cap_feats], dim=1)  # Now can be sharded via _sp_plan
-
-class MyModel(nn.Module):
-    def __init__(self):
-        self.unified_prepare = UnifiedPrepare()  # Submodule
-
-    def forward(self, x, cap_feats):
-        unified = self.unified_prepare(x, cap_feats)  # Hook can intercept here
-```
-
----
-
-###### Defining `_sp_plan`
-
-**Type definitions** (see [diffusers `_modeling_parallel.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py) for reference):
-
-```python
-from vllm_omni.diffusion.distributed.sp_plan import (
-    SequenceParallelInput,   # Corresponds to diffusers' ContextParallelInput
-    SequenceParallelOutput,  # Corresponds to diffusers' ContextParallelOutput
-)
-```
-
-| Parameter | Description |
-|-----------|-------------|
-| `split_dim` | Dimension to split/gather (usually `1` for sequence) |
-| `expected_dims` | Expected tensor rank for validation (optional) |
-| `split_output` | `False`: shard **input** parameters; `True`: shard **output** tensors |
-| `auto_pad` | Auto-pad if sequence not divisible by world_size (Ulysses only) |
-
-**Key naming convention:**
-
-| 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` |
-| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` |
-
-**Dictionary key types:**
-
-| Key type | `split_output` | Description |
-|----------|----------------|-------------|
-| `"param_name"` (str) | `False` | Shard **input parameter** by name |
-| `0`, `1` (int) | `True` | Shard **output tuple** by index |
-
-**Example** (similar to [diffusers `transformer_wan.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py)):
-
-```python
-class MyTransformer(nn.Module):
-    _sp_plan = {
-        # Shard rope module OUTPUTS (returns tuple)
-        "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 parameter
-        "blocks.0": {
-            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
-        },
-        # Gather at final projection
-        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
-    }
-```
-
----
-
-###### Hook flow
-
-```
-Input → [SequenceParallelSplitHook: pre_forward] → Module.forward() → [post_forward] → ...
-                                                                              ↓
-... → [SequenceParallelGatherHook: post_forward] → Output
-```
-
-1. **SplitHook** shards tensors before/after the target module
-2. **Attention layers** handle Ulysses/Ring communication internally
-3. **GatherHook** collects sharded outputs
-
-The framework automatically applies these hooks when `sequence_parallel_size > 1`.
-
----
-
-###### Method 2: Intrusive modification (For complex cases)
-
-For models with dynamic sharding logic that cannot be expressed via `_sp_plan`:
-
-```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 ...
-        output = sp_gather(output, dim=1)
-    return output
-```
-
----
-
-###### Choosing the right approach
-
-| Scenario | Approach |
-|----------|----------|
-| Standard transformer | `_sp_plan` |
-| Inline tensor ops need sharding | Extract to submodule + `_sp_plan` |
-| Dynamic/conditional sharding | Intrusive modification |
-
-
 ### CFG-Parallel
 
 #### Offline Inference
@@ -413,10 +317,10 @@ python image_edit.py \
   --model "Qwen/Qwen-Image-Edit" \
   --image "qwen_image_output.png" \
   --prompt "turn this cat to a dog" \
-  --negative_prompt "low quality, blurry" \
-  --cfg_scale 4.0 \
+  --negative-prompt "low quality, blurry" \
+  --cfg-scale 4.0 \
   --output "edited_image.png" \
-  --cfg_parallel_size 2
+  --cfg-parallel-size 2
 ```
 
 #### Online Serving
@@ -426,142 +330,3 @@ You can enable CFG-Parallel in online serving for diffusion models via `--cfg-pa
 ```bash
 vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
 ```
-
-#### How to parallelize a pipeline
-
-This section describes how to add CFG-Parallel to a diffusion **pipeline**. We use the Qwen-Image pipeline (`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`) as the reference implementation.
-
-In `QwenImagePipeline`, each diffusion step runs two denoiser forward passes sequentially:
-
-- positive (prompt-conditioned)
-- negative (negative-prompt-conditioned)
-
-CFG-Parallel assigns these two branches to different ranks in the **CFG group** and synchronizes the results.
-
-vLLM-omni provides `CFGParallelMixin` base class that encapsulates the CFG parallel logic. By inheriting from this mixin and calling its methods, pipelines can easily implement CFG parallel without writing repetitive code.
-
-**Key Methods in CFGParallelMixin:**
-- `predict_noise_maybe_with_cfg()`: Automatically handles CFG parallel noise prediction
-- `scheduler_step_maybe_with_cfg()`: Scheduler step with automatic CFG rank synchronization
-
-**Example Implementation:**
-
-```python
-class QwenImageCFGParallelMixin(CFGParallelMixin):
-    """
-    Base Mixin class for Qwen Image pipelines providing shared CFG methods.
-    """
-
-    def diffuse(
-        self,
-        prompt_embeds: torch.Tensor,
-        prompt_embeds_mask: torch.Tensor,
-        negative_prompt_embeds: torch.Tensor,
-        negative_prompt_embeds_mask: torch.Tensor,
-        latents: torch.Tensor,
-        img_shapes: torch.Tensor,
-        txt_seq_lens: torch.Tensor,
-        negative_txt_seq_lens: torch.Tensor,
-        timesteps: torch.Tensor,
-        do_true_cfg: bool,
-        guidance: torch.Tensor,
-        true_cfg_scale: float,
-        image_latents: torch.Tensor | None = None,
-        cfg_normalize: bool = True,
-        additional_transformer_kwargs: dict[str, Any] | None = None,
-    ) -> torch.Tensor:
-        self.transformer.do_true_cfg = do_true_cfg
-
-        for i, t in enumerate(timesteps):
-            timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
-
-            # Prepare kwargs for positive (conditional) prediction
-            positive_kwargs = {
-                "hidden_states": latents,
-                "timestep": timestep / 1000,
-                "guidance": guidance,
-                "encoder_hidden_states_mask": prompt_embeds_mask,
-                "encoder_hidden_states": prompt_embeds,
-                "img_shapes": img_shapes,
-                "txt_seq_lens": txt_seq_lens,
-            }
-
-            # Prepare kwargs for negative (unconditional) prediction
-            if do_true_cfg:
-                negative_kwargs = {
-                    "hidden_states": latents,
-                    "timestep": timestep / 1000,
-                    "guidance": guidance,
-                    "encoder_hidden_states_mask": negative_prompt_embeds_mask,
-                    "encoder_hidden_states": negative_prompt_embeds,
-                    "img_shapes": img_shapes,
-                    "txt_seq_lens": negative_txt_seq_lens,
-                }
-            else:
-                negative_kwargs = None
-
-            # Predict noise with automatic CFG parallel handling
-            # - In CFG parallel mode: rank0 computes positive, rank1 computes negative
-            # - Automatically gathers results and combines them on rank0
-            noise_pred = self.predict_noise_maybe_with_cfg(
-                do_true_cfg=do_true_cfg,
-                true_cfg_scale=true_cfg_scale,
-                positive_kwargs=positive_kwargs,
-                negative_kwargs=negative_kwargs,
-                cfg_normalize=cfg_normalize,
-            )
-
-            # Step scheduler with automatic CFG synchronization
-            # - Only rank0 computes the scheduler step
-            # - Automatically broadcasts updated latents to all ranks
-            latents = self.scheduler_step_maybe_with_cfg(
-                noise_pred, t, latents, do_true_cfg
-            )
-
-        return latents
-```
-
-**How it works:**
-1. Prepare separate `positive_kwargs` and `negative_kwargs` for conditional and unconditional predictions
-2. Call `predict_noise_maybe_with_cfg()` which:
-   - Detects if CFG parallel is enabled (`get_classifier_free_guidance_world_size() > 1`)
-   - Distributes computation: rank0 processes positive, rank1 processes negative
-   - Gathers predictions and combines them using `combine_cfg_noise()` on rank0
-   - Returns combined noise prediction (only valid on rank0)
-3. Call `scheduler_step_maybe_with_cfg()` which:
-   - Only rank0 computes the scheduler step
-   - Broadcasts the updated latents to all ranks for synchronization
-
-**How to customize**
-
-Some pipelines may need to customize the following functions in `CFGParallelMixin`:
-1. You may need to edit `predict_noise` function for custom behaviors.
-```python
-def predict_noise(self, *args, **kwargs):
-    """
-    Forward pass through transformer to predict noise.
-
-    Subclasses should override this if they need custom behavior,
-    but the default implementation calls self.transformer.
-    """
-    return self.transformer(*args, **kwargs)[0]
-
-```
-2. The default normalization function after combining the noise predictions from both branches is as follows. You may need to customize it.
-```python
-def cfg_normalize_function(self, noise_pred, comb_pred):
-    """
-    Normalize the combined noise prediction.
-
-    Args:
-        noise_pred: positive noise prediction
-        comb_pred: combined noise prediction after CFG
-
-    Returns:
-        Normalized noise prediction tensor
-    """
-    cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
-    noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
-    noise_pred = comb_pred * (cond_norm / noise_norm)
-    return noise_pred
-```

docs/user_guide/diffusion/quantization/fp8.md

@@ -0,0 +1,77 @@
+# FP8 Quantization
+
+## Overview
+
+FP8 quantization converts BF16/FP16 weights to FP8 at model load time. No calibration or pre-quantized checkpoint needed.
+
+Depending on the model, either all layers can be quantized, or some sensitive layers should stay in BF16. See the [per-model table](#supported-models) for which case applies.
+
+Common sensitive layers in DiT-based diffusion models include **image-stream MLPs** (`img_mlp`). These are particularly vulnerable to FP8 precision loss because they process denoising latents whose dynamic range shifts significantly across timesteps, and unlike attention projections (which benefit from QK-Norm stabilization), MLPs have no built-in normalization to absorb quantization error. In deep architectures (e.g., 60+ residual blocks), small per-layer errors compound and degrade output quality. Other layers such as **attention projections** (`to_qkv`, `to_out`) and **text-stream MLPs** (`txt_mlp`) are generally more robust due to normalization or more stable input statistics.
+
+## Configuration
+
+1. **Python API**: set `quantization="fp8"`. To skip sensitive layers, use `quantization_config` with `ignored_layers`.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+# All layers quantized
+omni = Omni(model="<your-model>", quantization="fp8")
+
+# Skip sensitive layers
+omni = Omni(
+    model="<your-model>",
+    quantization_config={
+        "method": "fp8",
+        "ignored_layers": ["<layer-name>"],
+    },
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
+2. **CLI**: pass `--quantization fp8` and optionally `--ignored-layers`.
+
+```bash
+# All layers
+python text_to_image.py --model <your-model> --quantization fp8
+
+# Skip sensitive layers
+python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers "img_mlp"
+
+# Online serving
+vllm serve <your-model> --omni --quantization fp8
+```
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `method` | str | — | Quantization method (`"fp8"`) |
+| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16 |
+| `activation_scheme` | str | `"dynamic"` | `"dynamic"` (no calibration) or `"static"` |
+| `weight_block_size` | list[int] \| None | `None` | Block size for block-wise weight quantization |
+
+The available `ignored_layers` names depend on the model architecture (e.g., `to_qkv`, `to_out`, `img_mlp`, `txt_mlp`). Consult the transformer source for your target model.
+
+## Supported Models
+
+| Model | HF Models | Recommendation | `ignored_layers` |
+|-------|-----------|---------------|------------------|
+| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
+| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
+
+## Combining with Other Features
+
+FP8 quantization can be combined with cache acceleration:
+
+```python
+omni = Omni(
+    model="<your-model>",
+    quantization="fp8",
+    cache_backend="tea_cache",
+    cache_config={"rel_l1_thresh": 0.2},
+)
+```

docs/user_guide/diffusion/quantization/overview.md

@@ -0,0 +1,17 @@
+# Quantization for Diffusion Transformers
+
+vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and accelerate inference.
+
+## Supported Methods
+
+| Method | Guide |
+|--------|-------|
+| FP8 | [FP8](fp8.md) |
+
+## Device Compatibility
+
+| GPU Generation | Example GPUs | FP8 Mode |
+|---------------|-------------------|----------|
+| Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
+
+Kernel selection is automatic.

docs/user_guide/diffusion_acceleration.md

@@ -1,6 +1,6 @@
 # Diffusion Acceleration Overview
 
-vLLM-Omni supports various cache acceleration methods to speed up diffusion model inference with minimal quality degradation. These methods include **cache methods** that intelligently cache intermediate computations to avoid redundant work across diffusion timesteps, and **parallelism methods** that distribute the computation across multiple devices.
+vLLM-Omni supports various acceleration methods to speed up diffusion model inference with minimal quality degradation. These include **cache methods** that intelligently cache intermediate computations to avoid redundant work across diffusion timesteps, **parallelism methods** that distribute the computation across multiple devices, and **quantization methods** that reduce memory footprint while preserving accuracy.
 
 ## Supported Acceleration Methods
 
@@ -14,6 +14,10 @@ vLLM-Omni currently supports two main cache acceleration backends:
 
 Both methods can provide significant speedups (typically **1.5x-2.0x**) while maintaining high output quality.
 
+vLLM-Omni also supports quantization methods:
+
+3. **[FP8 Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.
+
 vLLM-Omni also supports parallelism methods for diffusion models, including:
 
 1. [Ulysses-SP](diffusion/parallelism_acceleration.md#ulysses-sp) - splits the input along the sequence dimension and uses all-to-all communication to allow each device to compute only a subset of attention heads.
@@ -22,6 +26,10 @@ vLLM-Omni also supports parallelism methods for diffusion models, including:
 
 3. [CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel) - runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
 
+4. [Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism) - shards DiT weights across devices to reduce per-GPU memory usage.
+
+5. [VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism) - shards VAE decode/encode spatially across ranks to reduce VAE peak memory (and can speed up VAE decode).
+
 ## Quick Comparison
 
 ### Cache Methods
@@ -31,33 +39,47 @@ vLLM-Omni also supports parallelism methods for diffusion models, including:
 | **TeaCache** | `cache_backend="tea_cache"` | Simple, adaptive caching with minimal configuration | Quick setup, balanced speed/quality |
 | **Cache-DiT** | `cache_backend="cache_dit"` | Advanced caching with multiple techniques (DBCache, TaylorSeer, SCM) | Maximum acceleration, fine-grained control |
 
+### Quantization Methods
+
+| Method | Configuration | Description | Best For |
+|--------|--------------|-------------|----------|
+| **FP8** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
+
 ## Supported Models
 
 The following table shows which models are currently supported by each acceleration method:
 
 ### ImageGen
 
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
-|-------|------------------|:----------:|:-----------:|:-----------:|:----------------:|:----------------:|
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ |
-| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ❌ |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel |
+|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:---------------:|:------------------:|
+| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
+| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
+| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ (TP=2 only) | ✅ |
+| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
+| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
+| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
 
 ### VideoGen
 
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention |CFG-Parallel |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
 |-------|------------------|:--------:|:---------:|:----------:|:--------------:|:----------------:|
 | **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ |
 
+### Quantization
+
+| Model | Model Identifier | FP8 |
+|-------|------------------|:---:|
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ |
+| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ |
+
 
 ## Performance Benchmarks
 
@@ -201,6 +223,47 @@ outputs = omni.generate(
 )
 ```
 
+### Using Tensor Parallelism
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
+)
+
+outputs = omni.generate(
+    prompt="a cat reading a book",
+    num_inference_steps=9,
+    width=512,
+    height=512,
+)
+```
+
+### Using VAE Patch Parallelism
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    parallel_config=DiffusionParallelConfig(
+        tensor_parallel_size=2,
+        vae_patch_parallel_size=2,
+    ),
+)
+
+outputs = omni.generate(
+    prompt="a cat reading a book",
+    num_inference_steps=9,
+    width=1024,
+    height=1024,
+)
+```
+
 ### Using CFG-Parallel
 
 Run image-to-image:
@@ -227,11 +290,31 @@ outputs = omni.generate(
 )
 ```
 
+### Using FP8 Quantization
+
+```python
+from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+omni = Omni(
+    model="<your-model>",
+    quantization="fp8",
+)
+
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
+```
+
 ## Documentation
 
 For detailed information on each acceleration method:
 
 - **[TeaCache Guide](diffusion/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
 - **[Cache-DiT Acceleration Guide](diffusion/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
+- **[FP8 Quantization Guide](diffusion/quantization/overview.md)** - FP8 quantization for DiT models with per-layer control
+- **[Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism)** - Guidance on how to enable TP for diffusion models.
 - **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
 - **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
+- **[VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism)** - Guidance on how to reduce VAE memory via patch/tile parallelism.

docs/user_guide/examples/offline_inference/image_to_image.md

@@ -22,8 +22,8 @@ python image_edit.py \
   --image qwen-bear.png \
   --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
   --output output_image_edit.png \
-  --num_inference_steps 50 \
-  --cfg_scale 4.0
+  --num-inference-steps 50 \
+  --cfg-scale 4.0
 ```
 
 ### Multiple Image Editing (Qwen-Image-Edit-2509)
@@ -36,26 +36,26 @@ python image_edit.py \
   --image img1.png img2.png \
   --prompt "Combine these images into a single scene" \
   --output output_image_edit.png \
-  --num_inference_steps 50 \
-  --cfg_scale 4.0 \
-  --guidance_scale 1.0
+  --num-inference-steps 50 \
+  --cfg-scale 4.0 \
+  --guidance-scale 1.0
 ```
 
 Key arguments:
 
 - `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support.
 - `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images.
-- `--prompt` / `--negative_prompt`: text description (string).
-- `--cfg_scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.
-- `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
-- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
+- `--prompt` / `--negative-prompt`: text description (string).
+- `--cfg-scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.
+- `--guidance-scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg-scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
+- `--num-inference-steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--output`: path to save the generated PNG.
-- `--vae_use_slicing`: enable VAE slicing for memory optimization.
-- `--vae_use_tiling`: enable VAE tiling for memory optimization.
-- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--vae-use-slicing`: enable VAE slicing for memory optimization.
+- `--vae-use-tiling`: enable VAE tiling for memory optimization.
+- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 
-> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
+> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 
 ## Example materials
 

docs/user_guide/examples/offline_inference/image_to_video.md

@@ -20,15 +20,15 @@ python image_to_video.py \
   --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
   --image cherry_blossom.jpg \
   --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative_prompt "<optional quality filter>" \
+  --negative-prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
-  --num_frames 48 \
-  --guidance_scale 5.0 \
-  --guidance_scale_high 6.0 \
-  --num_inference_steps 40 \
-  --boundary_ratio 0.875 \
-  --flow_shift 12.0 \
+  --num-frames 48 \
+  --guidance-scale 5.0 \
+  --guidance-scale-high 6.0 \
+  --num-inference-steps 40 \
+  --boundary-ratio 0.875 \
+  --flow-shift 12.0 \
   --fps 16 \
   --output i2v_output.mp4
 ```
@@ -40,13 +40,13 @@ python image_to_video.py \
   --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
   --image cherry_blossom.jpg \
   --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative_prompt "<optional quality filter>" \
+  --negative-prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
-  --num_frames 48 \
-  --guidance_scale 4.0 \
-  --num_inference_steps 40 \
-  --flow_shift 12.0 \
+  --num-frames 48 \
+  --guidance-scale 4.0 \
+  --num-inference-steps 40 \
+  --flow-shift 12.0 \
   --fps 16 \
   --output i2v_output.mp4
 ```
@@ -57,20 +57,20 @@ Key arguments:
 - `--image`: Path to input image (required).
 - `--prompt`: Text description of desired motion/animation.
 - `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
-- `--num_frames`: Number of frames (default 81).
-- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high-noise stages for MoE).
-- `--negative_prompt`: Optional list of artifacts to suppress.
-- `--boundary_ratio`: Boundary split ratio for two-stage MoE models.
-- `--flow_shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
-- `--num_inference_steps`: Number of denoising steps (default 50).
+- `--num-frames`: Number of frames (default 81).
+- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
+- `--negative-prompt`: Optional list of artifacts to suppress.
+- `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
+- `--flow-shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
+- `--num-inference-steps`: Number of denoising steps (default 50).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: Path to save the generated video.
-- `--vae_use_slicing`: Enable VAE slicing for memory optimization.
-- `--vae_use_tiling`: Enable VAE tiling for memory optimization.
-- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--vae-use-slicing`: Enable VAE slicing for memory optimization.
+- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
+- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 
-> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
+> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 
 ## Example materials
 

docs/user_guide/examples/offline_inference/lora_inference.md

@@ -26,7 +26,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
     --prompt "A piece of cheesecake" \
     --lora-path /path/to/lora/ \
     --lora-scale 1.0 \
-    --num_inference_steps 50 \
+    --num-inference-steps 50 \
     --height 1024 \
     --width 1024 \
     --output output_preloaded.png
@@ -43,7 +43,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
     --prompt "A piece of cheesecake" \
     --lora-request-path /path/to/lora/ \
     --lora-scale 1.0 \
-    --num_inference_steps 50 \
+    --num-inference-steps 50 \
     --height 1024 \
     --width 1024 \
     --output output_per_request.png
@@ -56,7 +56,7 @@ If no LoRA request is provided, we will use the base model without any LoRA adap
 ```bash
 python -m examples.offline_inference.lora_inference.lora_inference \
     --prompt "A piece of cheesecake" \
-    --num_inference_steps 50 \
+    --num-inference-steps 50 \
     --height 1024 \
     --width 1024 \
     --output output_no_lora.png
@@ -77,7 +77,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
 - `--seed`: Random seed for reproducibility (default: 42)
 - `--height`: Image height in pixels (default: 1024)
 - `--width`: Image width in pixels (default: 1024)
-- `--num_inference_steps`: Number of denoising steps (default: 50)
+- `--num-inference-steps`: Number of denoising steps (default: 50)
 - `--output`: Output file path (default: `lora_output.png`)
 
 ## How LoRA Works

docs/user_guide/examples/offline_inference/text_to_audio.md

@@ -21,22 +21,22 @@ If you use a gated model (e.g., `stabilityai/stable-audio-open-1.0`), ensure you
 python text_to_audio.py \
   --model stabilityai/stable-audio-open-1.0 \
   --prompt "The sound of a hammer hitting a wooden surface" \
-  --negative_prompt "Low quality" \
+  --negative-prompt "Low quality" \
   --seed 42 \
-  --guidance_scale 7.0 \
-  --audio_length 10.0 \
-  --num_inference_steps 100 \
+  --guidance-scale 7.0 \
+  --audio-length 10.0 \
+  --num-inference-steps 100 \
   --output stable_audio_output.wav
 ```
 
 Key arguments:
 
 - `--prompt`: text description (string).
-- `--negative_prompt`: negative prompt for classifier-free guidance.
+- `--negative-prompt`: negative prompt for classifier-free guidance.
 - `--seed`: integer seed for deterministic generation.
-- `--guidance_scale`: classifier-free guidance scale.
-- `--audio_length`: audio duration in seconds.
-- `--num_inference_steps`: diffusion sampling steps.(more steps = higher quality, slower).
+- `--guidance-scale`: classifier-free guidance scale.
+- `--audio-length`: audio duration in seconds.
+- `--num-inference-steps`: diffusion sampling steps.(more steps = higher quality, slower).
 - `--output`: path to save the generated WAV file.
 
 ## Example materials

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -51,8 +51,7 @@ if __name__ == "__main__":
 
     For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
     list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
-    batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md
-    ) to let the model accept a longer batch of prompts.
+    batched inputs, you can [modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a longer batch of prompts.
 
 Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
 This is useful for models that support negative prompts.
@@ -83,9 +82,9 @@ python text_to_image.py \
   --model Tongyi-MAI/Z-Image-Turbo \
   --prompt "a cup of coffee on the table" \
   --seed 42 \
-  --cfg_scale 4.0 \
-  --num_images_per_prompt 1 \
-  --num_inference_steps 50 \
+  --cfg-scale 4.0 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 50 \
   --height 1024 \
   --width 1024 \
   --output outputs/coffee.png
@@ -95,17 +94,17 @@ Key arguments:
 
 - `--prompt`: text description (string).
 - `--seed`: integer seed for deterministic sampling.
-- `--cfg_scale`: true CFG scale (model-specific guidance strength).
-- `--num_images_per_prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
-- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
+- `--cfg-scale`: true CFG scale (model-specific guidance strength).
+- `--num-images-per-prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
+- `--num-inference-steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--height/--width`: output resolution (defaults 1024x1024).
 - `--output`: path to save the generated PNG.
-- `--vae_use_slicing`: enable VAE slicing for memory optimization.
-- `--vae_use_tiling`: enable VAE tiling for memory optimization.
-- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--vae-use-slicing`: enable VAE slicing for memory optimization.
+- `--vae-use-tiling`: enable VAE tiling for memory optimization.
+- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 
-> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
+> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 
 > ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
 

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -10,14 +10,14 @@ The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text
 ```bash
 python text_to_video.py \
   --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-  --negative_prompt "<optional quality filter>" \
+  --negative-prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
-  --num_frames 33 \
-  --guidance_scale 4.0 \
-  --guidance_scale_high 3.0 \
-  --flow_shift 12.0 \
-  --num_inference_steps 40 \
+  --num-frames 33 \
+  --guidance-scale 4.0 \
+  --guidance-scale-high 3.0 \
+  --flow-shift 12.0 \
+  --num-inference-steps 40 \
   --fps 16 \
   --output t2v_out.mp4
 ```
@@ -26,18 +26,18 @@ Key arguments:
 
 - `--prompt`: text description (string).
 - `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
-- `--num_frames`: Number of frames (Wan default is 81).
-- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high).
-- `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
-- `--boundary_ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
+- `--num-frames`: Number of frames (Wan default is 81).
+- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high).
+- `--negative-prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
+- `--boundary-ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
-- `--vae_use_slicing`: enable VAE slicing for memory optimization.
-- `--vae_use_tiling`: enable VAE tiling for memory optimization.
-- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--vae-use-slicing`: enable VAE slicing for memory optimization.
+- `--vae-use-tiling`: enable VAE tiling for memory optimization.
+- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 
-> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
+> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 
 ## Example materials
 

docs/user_guide/examples/online_serving/image_to_image.md

@@ -15,6 +15,9 @@ For **multi-image** input editing, use **Qwen-Image-Edit-2509** (QwenImageEditPl
 vllm serve Qwen/Qwen-Image-Edit --omni --port 8092
 ```
 
+!!! note
+    If you encounter Out-of-Memory (OOM) issues or have limited GPU memory, you can enable VAE slicing and tiling to reduce memory usage, --vae-use-slicing --vae-use-tiling
+
 ### Multi-Image Edit (Qwen-Image-Edit-2509)
 
 ```bash

docs/user_guide/examples/online_serving/qwen3_tts.md

@@ -1,9 +1,11 @@
-# Qwen3-TTS Online Serving
+# Qwen3-TTS
 
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_tts>.
 
 
-This directory contains examples for running Qwen3-TTS models with vLLM-Omni's online serving API.
+## 🛠️ Installation
+
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
 
 ## Supported Models
 
@@ -12,34 +14,77 @@ This directory contains examples for running Qwen3-TTS models with vLLM-Omni's o
 | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
 | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
 | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
 
-## Quick Start
+## Run examples (Qwen3-TTS)
 
-### 1. Start the Server
+### Launch the Server
 
 ```bash
-# CustomVoice model (default)
-./run_server.sh
+# CustomVoice model (predefined speakers)
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
 
-# Or specify task type
-./run_server.sh CustomVoice
-./run_server.sh VoiceDesign
-./run_server.sh Base
+# VoiceDesign model
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
+
+# Base model (voice cloning)
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
 ```
 
-### 2. Run the Client
+If you have custom stage configs file, launch the server with command below
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --stage-configs-path /path/to/stage_configs_file \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager
+```
+
+Alternatively, use the convenience script:
+```bash
+./run_server.sh                  # Default: CustomVoice model
+./run_server.sh CustomVoice      # CustomVoice model
+./run_server.sh VoiceDesign      # VoiceDesign model
+./run_server.sh Base             # Base (voice clone) model
+```
+
+### Send TTS Request
+
+Get into the example folder
+```bash
+cd examples/online_serving/qwen3_tts
+```
+
+####  Send request via python
 
 ```bash
 # CustomVoice: Use predefined speaker
 python openai_speech_client.py \
     --text "你好，我是通义千问" \
-    --voice Vivian \
+    --voice vivian \
     --language Chinese
 
 # CustomVoice with style instruction
 python openai_speech_client.py \
     --text "今天天气真好" \
-    --voice Ryan \
+    --voice ryan \
     --instructions "用开心的语气说"
 
 # VoiceDesign: Describe the voice style
@@ -58,29 +103,86 @@ python openai_speech_client.py \
     --ref-text "Original transcript of the reference audio"
 ```
 
-### 3. Using curl
+The Python client supports the following command-line arguments:
+
+- `--api-base`: API base URL (default: `http://localhost:8091`)
+- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice`)
+- `--task-type` (or `-t`): TTS task type. Options: `CustomVoice`, `VoiceDesign`, `Base`
+- `--text`: Text to synthesize (required)
+- `--voice`: Speaker/voice name (default: `vivian`). Options: `vivian`, `ryan`, `aiden`, etc.
+- `--language`: Language. Options: `Auto`, `Chinese`, `English`, `Japanese`, `Korean`, `German`, `French`, `Russian`, `Portuguese`, `Spanish`, `Italian`
+- `--instructions`: Voice style/emotion instructions
+- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task)
+- `--ref-text`: Reference audio transcript for voice cloning (Base task)
+- `--response-format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`, `aac`, `opus`
+- `--output` (or `-o`): Output audio file path (default: `tts_output.wav`)
+
+####  Send request via curl
 
 ```bash
 # Simple TTS request
-curl -X POST http://localhost:8000/v1/audio/speech \
+curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
         "input": "Hello, how are you?",
-        "voice": "Vivian",
+        "voice": "vivian",
         "language": "English"
     }' --output output.wav
 
 # With style instruction
-curl -X POST http://localhost:8000/v1/audio/speech \
+curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
         "input": "I am so excited!",
-        "voice": "Vivian",
+        "voice": "vivian",
         "instructions": "Speak with great enthusiasm"
     }' --output excited.wav
 
 # List available voices in CustomVoice models
-curl http://localhost:8000/v1/audio/voices
+curl http://localhost:8091/v1/audio/voices
+```
+
+### Using OpenAI SDK
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
+
+response = client.audio.speech.create(
+    model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+    voice="vivian",
+    input="Hello, how are you?",
+)
+
+response.stream_to_file("output.wav")
+```
+
+### Using Python httpx
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/speech",
+    json={
+        "input": "Hello, how are you?",
+        "voice": "vivian",
+        "language": "English",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+### FAQ
+
+If you encounter error about backend of librosa, try to install ffmpeg with command below.
+```
+sudo apt update
+sudo apt install ffmpeg
 ```
 
 ## API Reference
@@ -89,16 +191,31 @@ curl http://localhost:8000/v1/audio/voices
 
 ```
 POST /v1/audio/speech
+Content-Type: application/json
 ```
 
 This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
 
+### Voices Endpoint
+
+```
+GET /v1/audio/voices
+```
+
+Lists available voices for the loaded model:
+
+```json
+{
+    "voices": ["aiden", "dylan", "eric", "one_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
+}
+```
+
 ### Request Body
 
 ```json
 {
     "input": "Text to synthesize",
-    "voice": "Vivian",
+    "voice": "vivian",
     "response_format": "wav",
     "task_type": "CustomVoice",
     "language": "Auto",
@@ -114,56 +231,38 @@ This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/
 
 ### Response
 
-Returns audio data in the requested format (default: WAV).
+Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
 
 ## Parameters
 
-### Standard OpenAI Parameters
+### OpenAI Standard Parameters
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `input` | string | required | Text to synthesize |
-| `voice` | string | "Vivian" | Speaker/voice name |
+| `input` | string | **required** | Text to synthesize |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
 | `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
 | `speed` | float | 1.0 | Playback speed (0.25-4.0) |
-| `model` | string | optional | Model name (optional when serving single model) |
 
-### Qwen3-TTS Parameters
+### vLLM-Omni Extension Parameters
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base |
-| `language` | string | "Auto" | Language: Auto, Chinese, English, Japanese, Korean |
+| `language` | string | "Auto" | Language (see supported languages below) |
 | `instructions` | string | "" | Voice style/emotion instructions |
 | `max_new_tokens` | int | 2048 | Maximum tokens to generate |
 
+**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
+
 ### Voice Clone Parameters (Base task)
 
 | Parameter | Type | Required | Description |
 |-----------|------|----------|-------------|
-| `ref_audio` | string | Yes* | Reference audio (file path, URL, or base64) |
+| `ref_audio` | string | **Yes** | Reference audio (URL or base64 data URL) |
 | `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
-| `x_vector_only_mode` | bool | false | Use speaker embedding only (no ICL) |
-
-## Python Usage
-
-```python
-import httpx
-
-# Simple request
-response = httpx.post(
-    "http://localhost:8000/v1/audio/speech",
-    json={
-        "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
-        "input": "Hello world",
-        "voice": "Vivian",
-    },
-    timeout=300.0,
-)
-
-with open("output.wav", "wb") as f:
-    f.write(response.content)
-```
+| `x_vector_only_mode` | bool | No | Use speaker embedding only (no ICL) |
 
 ## Limitations
 
@@ -172,10 +271,11 @@ with open("output.wav", "wb") as f:
 
 ## Troubleshooting
 
-1. **Connection refused**: Make sure the server is running on the correct port
-2. **Out of memory**: Reduce `--gpu-memory-utilization` in run_server.sh
-3. **Unsupported speaker**: Check supported speakers via model documentation
-4. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning
+1. **TTS model did not produce audio output**: Ensure you're using the correct model variant for your task type (CustomVoice task → CustomVoice model, etc.)
+2. **Connection refused**: Make sure the server is running on the correct port
+3. **Out of memory**: Use smaller model variant (`Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`) or reduce `--gpu-memory-utilization`
+4. **Unsupported speaker**: Use `/v1/audio/voices` to list available voices for the loaded model
+5. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning
 
 ## Example materials
 

examples/offline_inference/bagel/README.md

@@ -151,6 +151,24 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 
 ------
 
+#### Tensor Parallelism (TP)
+
+For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) by modifying the stage configuration (e.g., [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)).
+
+1. **Set `tensor_parallel_size`**: Increase this value (e.g., to `2` or `4`).
+2. **Set `devices`**: Specify the comma-separated GPU IDs to be used for the stage (e.g., `"0,1"`).
+
+Example configuration for TP=2 on GPUs 0 and 1:
+```yaml
+    engine_args:
+      tensor_parallel_size: 2
+      ...
+    runtime:
+      devices: "0,1"
+```
+
+------
+
 #### 🔗 Runtime Configuration
 
 | Parameter             | Value   | Description                      |

examples/offline_inference/bagel/end2end.py

@@ -19,7 +19,7 @@ def parse_args():
         default=None,
         help="Path to a .txt file with one prompt per line (preferred).",
     )
-    parser.add_argument("--prompt_type", default="text", choices=["text"])
+    parser.add_argument("--prompt-type", default="text", choices=["text"])
 
     parser.add_argument(
         "--modality",
@@ -36,7 +36,7 @@ def parse_args():
     )
 
     # OmniLLM init args
-    parser.add_argument("--enable-stats", action="store_true", default=False)
+    parser.add_argument("--log-stats", action="store_true", default=False)
     parser.add_argument("--init-sleep-seconds", type=int, default=20)
     parser.add_argument("--batch-timeout", type=int, default=5)
     parser.add_argument("--init-timeout", type=int, default=300)
@@ -46,6 +46,12 @@ def parse_args():
     parser.add_argument("--stage-configs-path", type=str, default=None)
     parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
 
+    parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)")
+    parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)")
+    parser.add_argument(
+        "--negative-prompt", type=str, default=None, help="Negative prompt (not yet supported, reserved for future)"
+    )
+
     args = parser.parse_args()
     return args
 
@@ -102,6 +108,10 @@ def main():
                 seed=52,
                 need_kv_receive=False,
                 num_inference_steps=args.steps,
+                extra_args={
+                    "cfg_text_scale": args.cfg_text_scale,
+                    "cfg_img_scale": args.cfg_img_scale,
+                },
             ),
         )
 
@@ -120,7 +130,7 @@ def main():
 
         omni_kwargs.update(
             {
-                "log_stats": args.enable_stats,
+                "log_stats": args.log_stats,
                 "init_sleep_seconds": args.init_sleep_seconds,
                 "batch_timeout": args.batch_timeout,
                 "init_timeout": args.init_timeout,
@@ -158,7 +168,12 @@ def main():
         if args.modality == "text2img":
             params_list[0].max_tokens = 1  # type: ignore # The first stage is a SamplingParam (vllm)
             if len(params_list) > 1:
-                params_list[1].num_inference_steps = args.steps  # type: ignore # The second stage is an OmniDiffusionSamplingParam
+                diffusion_params = params_list[1]
+                diffusion_params.num_inference_steps = args.steps  # type: ignore
+                diffusion_params.extra_args = {  # type: ignore
+                    "cfg_text_scale": args.cfg_text_scale,
+                    "cfg_img_scale": args.cfg_img_scale,
+                }
 
         omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 

examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import logging
+from typing import Any
+
+import PIL.Image
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+logger = logging.getLogger(__name__)
+
+
+class CustomPipeline(QwenImageEditPipeline):
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        image: PIL.Image.Image | torch.Tensor | None = None,
+        true_cfg_scale: float = 4.0,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 50,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 1.0,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_embeds_mask: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds_mask: torch.Tensor | None = None,
+        output_type: str | None = "pil",
+        attention_kwargs: dict[str, Any] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        max_sequence_length: int = 512,
+    ) -> DiffusionOutput:
+        """Forward pass for image editing with dummy trajectory data."""
+        # Call parent's forward to get the normal output
+        output = super().forward(
+            req=req,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            true_cfg_scale=true_cfg_scale,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            sigmas=sigmas,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=num_images_per_prompt,
+            generator=generator,
+            latents=latents,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+            output_type=output_type,
+            attention_kwargs=attention_kwargs,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        # Get actual num_inference_steps used
+        actual_num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+
+        # Create dummy trajectory data
+        dummy_trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32)
+
+        # Inject dummy trajectory data into output
+        output.trajectory_latents = dummy_trajectory_latents
+
+        return output

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+======================================================
+ Qwen Image Editor (via vLLM-Omni)
+======================================================
+
+Example CLI for Qwen-Image-Edit and compatible models.
+
+This script edits, combines, or layers images according to a text prompt
+using the vLLM Omni diffusion backend.
+
+Examples
+---------
+Single image edit:
+    python image_edit.py \
+        --image input.png \
+        --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles saying 'Be Kind'" \
+        --output output.png
+
+Multiple image composition:
+    python image_edit.py \
+        --image input1.png input2.png \
+        --prompt "Combine these into a single magical composition" \
+        --output combined.png
+
+Accelerated with cache_dit:
+    python image_edit.py \
+        --image input.png \
+        --prompt "Edit description" \
+        --cache_backend cache_dit \
+        --cache_dit_enable_taylorseer
+
+Layered RGBA output:
+    python image_edit.py \
+        --model "Qwen/Qwen-Image-Layered" \
+        --image input.png \
+        --prompt "" \
+        --output layered \
+        --layers 4 \
+        --color-format RGBA
+"""
+
+import argparse
+import asyncio
+import os
+import time
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+
+# ===========================
+# Argument Parser
+# ===========================
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Edit or generate images using Qwen-Image-Edit.")
+    parser.add_argument("--model", default="Qwen/Qwen-Image-Edit", help="Model name or local path.")
+    parser.add_argument("--image", type=str, nargs="+", required=True, help="Input image file(s).")
+    parser.add_argument("--prompt", type=str, required=True, help="Edit description prompt.")
+    parser.add_argument("--negative_prompt", type=str, default=None)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--cfg_scale", type=float, default=4.0)
+    parser.add_argument("--guidance_scale", type=float, default=1.0)
+    parser.add_argument("--output", type=str, default="output.png")
+    parser.add_argument("--num_outputs_per_prompt", type=int, default=1)
+    parser.add_argument("--num_inference_steps", type=int, default=50)
+    parser.add_argument("--cache_backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
+    parser.add_argument("--ulysses_degree", type=int, default=1)
+    parser.add_argument("--ring_degree", type=int, default=1)
+    parser.add_argument("--tensor_parallel_size", type=int, default=1)
+    parser.add_argument("--layers", type=int, default=4)
+    parser.add_argument("--resolution", type=int, default=640)
+    parser.add_argument("--color-format", type=str, default="RGB")
+
+    # Acceleration + Optimization Options
+    parser.add_argument("--cache_dit_fn_compute_blocks", type=int, default=1)
+    parser.add_argument("--cache_dit_bn_compute_blocks", type=int, default=0)
+    parser.add_argument("--cache_dit_max_warmup_steps", type=int, default=4)
+    parser.add_argument("--cache_dit_residual_diff_threshold", type=float, default=0.24)
+    parser.add_argument("--cache_dit_max_continuous_cached_steps", type=int, default=3)
+    parser.add_argument("--cache_dit_enable_taylorseer", action="store_true", default=False)
+    parser.add_argument("--cache_dit_taylorseer_order", type=int, default=1)
+    parser.add_argument(
+        "--cache_dit_scm_steps_mask_policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
+    )
+    parser.add_argument("--cache_dit_scm_steps_policy", type=str, default="dynamic", choices=["dynamic", "static"])
+    parser.add_argument("--tea_cache_rel_l1_thresh", type=float, default=0.2)
+    parser.add_argument("--cfg_parallel_size", type=int, default=1, choices=[1, 2])
+    parser.add_argument("--enforce_eager", action="store_true")
+    parser.add_argument("--vae_use_slicing", action="store_true")
+    parser.add_argument("--vae_use_tiling", action="store_true")
+    parser.add_argument("--enable-cpu-offload", action="store_true")
+
+    return parser.parse_args()
+
+
+# ===========================
+# Core Logic
+# ===========================
+async def main():
+    args = parse_args()
+
+    # ---- Load input images ----
+    input_images: list[Image.Image] = []
+    for image_path in args.image:
+        if not os.path.exists(image_path):
+            raise FileNotFoundError(f"Input image not found: {image_path}")
+        input_images.append(Image.open(image_path).convert(args.color_format))
+
+    # Single or multi-image input
+    input_image: Image.Image | list[Image.Image] = input_images[0] if len(input_images) == 1 else input_images
+
+    # ---- Torch setup ----
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
+    )
+
+    # ---- Cache Config ----
+    if args.cache_backend == "cache_dit":
+        cache_config = {
+            "Fn_compute_blocks": args.cache_dit_fn_compute_blocks,
+            "Bn_compute_blocks": args.cache_dit_bn_compute_blocks,
+            "max_warmup_steps": args.cache_dit_max_warmup_steps,
+            "residual_diff_threshold": args.cache_dit_residual_diff_threshold,
+            "max_continuous_cached_steps": args.cache_dit_max_continuous_cached_steps,
+            "enable_taylorseer": args.cache_dit_enable_taylorseer,
+            "taylorseer_order": args.cache_dit_taylorseer_order,
+            "scm_steps_mask_policy": args.cache_dit_scm_steps_mask_policy,
+            "scm_steps_policy": args.cache_dit_scm_steps_policy,
+        }
+    elif args.cache_backend == "tea_cache":
+        cache_config = {"rel_l1_thresh": args.tea_cache_rel_l1_thresh}
+    else:
+        cache_config = None
+
+    # ---- Initialize Omni ----
+    omni = Omni(
+        model=args.model,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
+        cache_backend=args.cache_backend,
+        cache_config=cache_config,
+        parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
+        enable_cpu_offload=args.enable_cpu_offload,
+        diffusion_load_format="dummy",
+        custom_pipeline_args={"pipeline_class": "custom_pipeline.CustomPipeline"},
+    )
+
+    print(">>> Pipeline loaded successfully")
+
+    # ---- Profiling + Info ----
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    print(f"\n{'=' * 60}")
+    print("Generation Configuration")
+    print(f"Model: {args.model}")
+    print(f"Inference steps: {args.num_inference_steps}")
+    print(f"Cache: {args.cache_backend or 'None'}")
+    print(f"Input(s): {len(input_images)} image(s)")
+    for idx, img in enumerate(input_images):
+        print(f"  • Image {idx + 1}: size={img.size}")
+    print(
+        f"Parallel config: Ulysses={args.ulysses_degree}, Ring={args.ring_degree}, "
+        f"CFG Parallel={args.cfg_parallel_size}, Tensor Parallel={args.tensor_parallel_size}"
+    )
+    print(f"{'=' * 60}\n")
+
+    if profiler_enabled:
+        omni.start_profile()
+
+    # ---- Generation ----
+    t0 = time.perf_counter()
+    outputs = omni.generate(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+            "multi_modal_data": {"image": input_image},
+        },
+        OmniDiffusionSamplingParams(
+            generator=generator,
+            true_cfg_scale=args.cfg_scale,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_outputs_per_prompt,
+            layers=args.layers,
+            resolution=args.resolution,
+        ),
+    )
+    t1 = time.perf_counter()
+    print(f"Generation completed in {t1 - t0:.2f}s")
+
+    if profiler_enabled:
+        omni.stop_profile()
+
+    # ---- Output Validation ----
+    if not outputs:
+        raise ValueError("No output produced from omni.generate()")
+
+    first_out = outputs[0].request_output[0]
+    req_out: OmniRequestOutput = first_out
+
+    # Verify trajectory data (from custom pipeline)
+    print(f"\n{'=' * 60}")
+    print("TRAJECTORY VERIFICATION:")
+    assert hasattr(req_out, "latents") and req_out.latents is not None
+    print(f"  ✓ trajectory_latents shape: {req_out.latents.shape}")
+    print(f"  ✓ trajectory_latents dtype: {req_out.latents.dtype}")
+    print(f"  ✓ trajectory_latents mean: {req_out.latents.mean().item():.6f}")
+    print(f"  ✓ trajectory_latents std: {req_out.latents.std().item():.6f}")
+    assert req_out.latents.shape[0] == args.num_inference_steps, (
+        f"Expected {args.num_inference_steps} latent snapshots, got {req_out.latents.shape[0]}"
+    )
+    print(f"{'=' * 60}\n")
+
+    # ---- Save images ----
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    suffix = output_path.suffix or ".png"
+    stem = output_path.stem or "output"
+
+    images = req_out.images
+    if args.num_outputs_per_prompt <= 1:
+        img = images[0]
+        if isinstance(img, list):  # Layered
+            for i, sub_img in enumerate(img):
+                save_path = output_path.parent / f"{stem}_{i}{suffix}"
+                sub_img.save(save_path)
+                print(f"Saved layer {i}: {save_path.resolve()}")
+        else:
+            img.save(output_path)
+            print(f"Saved: {output_path.resolve()}")
+    else:
+        for idx, img in enumerate(images):
+            if isinstance(img, list):
+                for sub_idx, sub_img in enumerate(img):
+                    save_path = output_path.parent / f"{stem}_{idx}_{sub_idx}{suffix}"
+                    sub_img.save(save_path)
+                    print(f"Saved composite: {save_path.resolve()}")
+            else:
+                save_path = output_path.parent / f"{stem}_{idx}{suffix}"
+                img.save(save_path)
+                print(f"Saved: {save_path.resolve()}")
+
+
+# ===========================
+# Entrypoint
+# ===========================
+if __name__ == "__main__":
+    asyncio.run(main())

examples/offline_inference/custom_pipeline/image_to_image/run.sh

@@ -0,0 +1,6 @@
+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 and poetic bubbles such as 'Be Kind'" \
+        --output output_image_edit.png \
+        --num_inference_steps 10

examples/offline_inference/image_to_image/image_edit.py

@@ -9,34 +9,34 @@ Usage (single image):
         --image input.png \
         --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
         --output output_image_edit.png \
-        --num_inference_steps 50 \
-        --cfg_scale 4.0 \
-        --guidance_scale 1.0
+        --num-inference-steps 50 \
+        --cfg-scale 4.0 \
+        --guidance-scale 1.0
 
 Usage (multiple images):
     python image_edit.py \
         --image input1.png input2.png input3.png \
         --prompt "Combine these images into a single scene" \
         --output output_image_edit.png \
-        --num_inference_steps 50 \
-        --cfg_scale 4.0 \
-        --guidance_scale 1.0
+        --num-inference-steps 50 \
+        --cfg-scale 4.0 \
+        --guidance-scale 1.0
 
 Usage (with cache-dit acceleration):
     python image_edit.py \
         --image input.png \
         --prompt "Edit description" \
-        --cache_backend cache_dit \
-        --cache_dit_max_continuous_cached_steps 3 \
-        --cache_dit_residual_diff_threshold 0.24 \
-        --cache_dit_enable_taylorseer
+        --cache-backend cache_dit \
+        --cache-dit-max-continuous-cached-steps 3 \
+        --cache-dit-residual-diff-threshold 0.24 \
+        --cache-dit-enable-taylorseer
 
 Usage (with tea_cache acceleration):
     python image_edit.py \
         --image input.png \
         --prompt "Edit description" \
-        --cache_backend tea_cache \
-        --tea_cache_rel_l1_thresh 0.25
+        --cache-backend tea_cache \
+        --tea-cache-rel-l1-thresh 0.25
 
 Usage (layered):
     python image_edit.py \
@@ -44,8 +44,8 @@ Usage (layered):
         --image input.png \
         --prompt "" \
         --output "layered" \
-        --num_inference_steps 50 \
-        --cfg_scale 4.0 \
+        --num-inference-steps 50 \
+        --cfg-scale 4.0 \
         --layers 4 \
         --color-format "RGBA"
 
@@ -53,17 +53,17 @@ Usage (with CFG Parallel):
     python image_edit.py \
         --image input.png \
         --prompt "Edit description" \
-        --cfg_parallel_size 2 \
-        --num_inference_steps 50 \
-        --cfg_scale 4.0
+        --cfg-parallel-size 2 \
+        --num-inference-steps 50 \
+        --cfg-scale 4.0
 
 Usage (disable torch.compile):
     python image_edit.py \
         --image input.png \
         --prompt "Edit description" \
-        --enforce_eager \
-        --num_inference_steps 50 \
-        --cfg_scale 4.0
+        --enforce-eager \
+        --num-inference-steps 50 \
+        --cfg-scale 4.0
 
 For more options, run:
     python image_edit.py --help
@@ -109,7 +109,7 @@ def parse_args() -> argparse.Namespace:
         help="Text prompt describing the edit to make to the image.",
     )
     parser.add_argument(
-        "--negative_prompt",
+        "--negative-prompt",
         type=str,
         default=None,
         required=False,
@@ -121,7 +121,7 @@ def parse_args() -> argparse.Namespace:
         help="Random seed for deterministic results.",
     )
     parser.add_argument(
-        "--cfg_scale",
+        "--cfg-scale",
         type=float,
         default=4.0,
         help=(
@@ -132,12 +132,12 @@ def parse_args() -> argparse.Namespace:
         ),
     )
     parser.add_argument(
-        "--guidance_scale",
+        "--guidance-scale",
         type=float,
         default=1.0,
         help=(
             "Guidance scale for guidance-distilled models (default: 1.0, disabled). "
-            "Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale "
+            "Unlike classifier-free guidance (--cfg-scale), guidance-distilled models take the guidance scale "
             "directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models."
         ),
     )
@@ -148,19 +148,19 @@ def parse_args() -> argparse.Namespace:
         help=("Path to save the edited image (PNG). Or prefix for Qwen-Image-Layered model save images(PNG)."),
     )
     parser.add_argument(
-        "--num_outputs_per_prompt",
+        "--num-outputs-per-prompt",
         type=int,
         default=1,
         help="Number of images to generate for the given prompt.",
     )
     parser.add_argument(
-        "--num_inference_steps",
+        "--num-inference-steps",
         type=int,
         default=50,
         help="Number of denoising steps for the diffusion sampler.",
     )
     parser.add_argument(
-        "--cache_backend",
+        "--cache-backend",
         type=str,
         default=None,
         choices=["cache_dit", "tea_cache"],
@@ -171,19 +171,19 @@ def parse_args() -> argparse.Namespace:
         ),
     )
     parser.add_argument(
-        "--ulysses_degree",
+        "--ulysses-degree",
         type=int,
         default=1,
         help="Number of GPUs used for ulysses sequence parallelism.",
     )
     parser.add_argument(
-        "--ring_degree",
+        "--ring-degree",
         type=int,
         default=1,
         help="Number of GPUs used for ring sequence parallelism.",
     )
     parser.add_argument(
-        "--tensor_parallel_size",
+        "--tensor-parallel-size",
         type=int,
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
@@ -205,56 +205,56 @@ def parse_args() -> argparse.Namespace:
 
     # Cache-DiT specific parameters
     parser.add_argument(
-        "--cache_dit_fn_compute_blocks",
+        "--cache-dit-fn-compute-blocks",
         type=int,
         default=1,
         help="[cache-dit] Number of forward compute blocks. Optimized for single-transformer models.",
     )
     parser.add_argument(
-        "--cache_dit_bn_compute_blocks",
+        "--cache-dit-bn-compute-blocks",
         type=int,
         default=0,
         help="[cache-dit] Number of backward compute blocks.",
     )
     parser.add_argument(
-        "--cache_dit_max_warmup_steps",
+        "--cache-dit-max-warmup-steps",
         type=int,
         default=4,
         help="[cache-dit] Maximum warmup steps (works for few-step models).",
     )
     parser.add_argument(
-        "--cache_dit_residual_diff_threshold",
+        "--cache-dit-residual-diff-threshold",
         type=float,
         default=0.24,
         help="[cache-dit] Residual diff threshold. Higher values enable more aggressive caching.",
     )
     parser.add_argument(
-        "--cache_dit_max_continuous_cached_steps",
+        "--cache-dit-max-continuous-cached-steps",
         type=int,
         default=3,
         help="[cache-dit] Maximum continuous cached steps to prevent precision degradation.",
     )
     parser.add_argument(
-        "--cache_dit_enable_taylorseer",
+        "--cache-dit-enable-taylorseer",
         action="store_true",
         default=False,
         help="[cache-dit] Enable TaylorSeer acceleration (not suitable for few-step models).",
     )
     parser.add_argument(
-        "--cache_dit_taylorseer_order",
+        "--cache-dit-taylorseer-order",
         type=int,
         default=1,
         help="[cache-dit] TaylorSeer polynomial order.",
     )
     parser.add_argument(
-        "--cache_dit_scm_steps_mask_policy",
+        "--cache-dit-scm-steps-mask-policy",
         type=str,
         default=None,
         choices=[None, "slow", "medium", "fast", "ultra"],
         help="[cache-dit] SCM mask policy: None (disabled), slow, medium, fast, ultra.",
     )
     parser.add_argument(
-        "--cache_dit_scm_steps_policy",
+        "--cache-dit-scm-steps-policy",
         type=str,
         default="dynamic",
         choices=["dynamic", "static"],
@@ -263,30 +263,30 @@ def parse_args() -> argparse.Namespace:
 
     # TeaCache specific parameters
     parser.add_argument(
-        "--tea_cache_rel_l1_thresh",
+        "--tea-cache-rel-l1-thresh",
         type=float,
         default=0.2,
         help="[tea_cache] Threshold for accumulated relative L1 distance.",
     )
     parser.add_argument(
-        "--cfg_parallel_size",
+        "--cfg-parallel-size",
         type=int,
         default=1,
         choices=[1, 2],
         help="Number of GPUs used for classifier free guidance parallel size.",
     )
     parser.add_argument(
-        "--enforce_eager",
+        "--enforce-eager",
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
     parser.add_argument(
-        "--vae_use_slicing",
+        "--vae-use-slicing",
         action="store_true",
         help="Enable VAE slicing for memory optimization.",
     )
     parser.add_argument(
-        "--vae_use_tiling",
+        "--vae-use-tiling",
         action="store_true",
         help="Enable VAE tiling for memory optimization.",
     )
@@ -300,12 +300,6 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
-    parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
-    )
     return parser.parse_args()
 
 
@@ -362,7 +356,6 @@ def main():
     omni = Omni(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,

examples/offline_inference/image_to_image/image_to_image.md

@@ -19,8 +19,8 @@ python image_edit.py \
   --image qwen-bear.png \
   --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
   --output output_image_edit.png \
-  --num_inference_steps 50 \
-  --cfg_scale 4.0
+  --num-inference-steps 50 \
+  --cfg-scale 4.0
 ```
 
 ### Multiple Image Editing (Qwen-Image-Edit-2509)
@@ -33,23 +33,23 @@ python image_edit.py \
   --image img1.png img2.png \
   --prompt "Combine these images into a single scene" \
   --output output_image_edit.png \
-  --num_inference_steps 50 \
-  --cfg_scale 4.0 \
-  --guidance_scale 1.0
+  --num-inference-steps 50 \
+  --cfg-scale 4.0 \
+  --guidance-scale 1.0
 ```
 
 Key arguments:
 
 - `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support.
 - `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images.
-- `--prompt` / `--negative_prompt`: text description (string).
-- `--cfg_scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.
-- `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
-- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
+- `--prompt` / `--negative-prompt`: text description (string).
+- `--cfg-scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality.
+- `--guidance-scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg-scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
+- `--num-inference-steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--output`: path to save the generated PNG.
-- `--vae_use_slicing`: enable VAE slicing for memory optimization.
-- `--vae_use_tiling`: enable VAE tiling for memory optimization.
-- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--vae-use-slicing`: enable VAE slicing for memory optimization.
+- `--vae-use-tiling`: enable VAE tiling for memory optimization.
+- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 
-> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
+> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/image_to_image/run_qwen_image_edit_2511.sh

@@ -3,6 +3,6 @@ python image_edit.py \
     --image qwen_bear.png \
     --prompt "Add a white art board written with colorful text 'vLLM-Omni' on grassland. Add a paintbrush in the bear's hands. position the bear standing in front of the art board as if painting" \
     --output output_image_edit.png \
-    --num_inference_steps 50 \
-    --cfg_scale 4.0 \
-    --cache_backend  cache_dit \
+    --num-inference-steps 50 \
+    --cfg-scale 4.0 \
+    --cache-backend  cache_dit \

examples/offline_inference/image_to_video/README.md

@@ -17,15 +17,15 @@ python image_to_video.py \
   --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
   --image cherry_blossom.jpg \
   --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative_prompt "<optional quality filter>" \
+  --negative-prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
-  --num_frames 48 \
-  --guidance_scale 5.0 \
-  --guidance_scale_high 6.0 \
-  --num_inference_steps 40 \
-  --boundary_ratio 0.875 \
-  --flow_shift 12.0 \
+  --num-frames 48 \
+  --guidance-scale 5.0 \
+  --guidance-scale-high 6.0 \
+  --num-inference-steps 40 \
+  --boundary-ratio 0.875 \
+  --flow-shift 12.0 \
   --fps 16 \
   --output i2v_output.mp4
 ```
@@ -37,13 +37,13 @@ python image_to_video.py \
   --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
   --image cherry_blossom.jpg \
   --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
-  --negative_prompt "<optional quality filter>" \
+  --negative-prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
-  --num_frames 48 \
-  --guidance_scale 4.0 \
-  --num_inference_steps 40 \
-  --flow_shift 12.0 \
+  --num-frames 48 \
+  --guidance-scale 4.0 \
+  --num-inference-steps 40 \
+  --flow-shift 12.0 \
   --fps 16 \
   --output i2v_output.mp4
 ```
@@ -54,17 +54,17 @@ Key arguments:
 - `--image`: Path to input image (required).
 - `--prompt`: Text description of desired motion/animation.
 - `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
-- `--num_frames`: Number of frames (default 81).
-- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high-noise stages for MoE).
-- `--negative_prompt`: Optional list of artifacts to suppress.
-- `--boundary_ratio`: Boundary split ratio for two-stage MoE models.
-- `--flow_shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
-- `--num_inference_steps`: Number of denoising steps (default 50).
+- `--num-frames`: Number of frames (default 81).
+- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
+- `--negative-prompt`: Optional list of artifacts to suppress.
+- `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
+- `--flow-shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
+- `--num-inference-steps`: Number of denoising steps (default 50).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: Path to save the generated video.
-- `--vae_use_slicing`: Enable VAE slicing for memory optimization.
-- `--vae_use_tiling`: Enable VAE tiling for memory optimization.
-- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--vae-use-slicing`: Enable VAE slicing for memory optimization.
+- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
+- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 
-> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
+> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/image_to_video/image_to_video.py

@@ -42,31 +42,31 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--image", required=True, help="Path to input image.")
     parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
-    parser.add_argument("--negative_prompt", default="", help="Negative prompt.")
+    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    parser.add_argument("--guidance_scale", type=float, default=5.0, help="CFG scale.")
+    parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
     parser.add_argument(
-        "--guidance_scale_high", type=float, default=None, help="Optional separate CFG for high-noise (MoE only)."
+        "--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise (MoE only)."
     )
     parser.add_argument(
         "--height", type=int, default=None, help="Video height (auto-calculated from image if not set)."
     )
     parser.add_argument("--width", type=int, default=None, help="Video width (auto-calculated from image if not set).")
-    parser.add_argument("--num_frames", type=int, default=81, help="Number of frames.")
-    parser.add_argument("--num_inference_steps", type=int, default=50, help="Sampling steps.")
-    parser.add_argument("--boundary_ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
+    parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
+    parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps.")
+    parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
     parser.add_argument(
-        "--flow_shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
+        "--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
     )
     parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
     parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
     parser.add_argument(
-        "--vae_use_slicing",
+        "--vae-use-slicing",
         action="store_true",
         help="Enable VAE slicing for memory optimization.",
     )
     parser.add_argument(
-        "--vae_use_tiling",
+        "--vae-use-tiling",
         action="store_true",
         help="Enable VAE tiling for memory optimization.",
     )
@@ -81,20 +81,14 @@ def parse_args() -> argparse.Namespace:
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
     parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
-    )
-    parser.add_argument(
-        "--cfg_parallel_size",
+        "--cfg-parallel-size",
         type=int,
         default=1,
         choices=[1, 2],
         help="Number of GPUs used for classifier free guidance parallel size.",
     )
     parser.add_argument(
-        "--enforce_eager",
+        "--enforce-eager",
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
@@ -139,7 +133,6 @@ def main():
     omni = Omni(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,

examples/offline_inference/lora_inference/README.md

@@ -23,7 +23,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
     --prompt "A piece of cheesecake" \
     --lora-path /path/to/lora/ \
     --lora-scale 1.0 \
-    --num_inference_steps 50 \
+    --num-inference-steps 50 \
     --height 1024 \
     --width 1024 \
     --output output_preloaded.png
@@ -40,7 +40,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
     --prompt "A piece of cheesecake" \
     --lora-request-path /path/to/lora/ \
     --lora-scale 1.0 \
-    --num_inference_steps 50 \
+    --num-inference-steps 50 \
     --height 1024 \
     --width 1024 \
     --output output_per_request.png
@@ -53,7 +53,7 @@ If no LoRA request is provided, we will use the base model without any LoRA adap
 ```bash
 python -m examples.offline_inference.lora_inference.lora_inference \
     --prompt "A piece of cheesecake" \
-    --num_inference_steps 50 \
+    --num-inference-steps 50 \
     --height 1024 \
     --width 1024 \
     --output output_no_lora.png
@@ -74,7 +74,7 @@ python -m examples.offline_inference.lora_inference.lora_inference \
 - `--seed`: Random seed for reproducibility (default: 42)
 - `--height`: Image height in pixels (default: 1024)
 - `--width`: Image width in pixels (default: 1024)
-- `--num_inference_steps`: Number of denoising steps (default: 50)
+- `--num-inference-steps`: Number of denoising steps (default: 50)
 - `--output`: Output file path (default: `lora_output.png`)
 
 ## How LoRA Works

examples/offline_inference/lora_inference/lora_inference.py

@@ -18,7 +18,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
     parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")
     parser.add_argument(
-        "--num_inference_steps",
+        "--num-inference-steps",
         type=int,
         default=50,
         help="Number of denoising steps for the diffusion sampler.",

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -322,7 +322,7 @@ def main(args):
         query_result = query_func()
     omni_llm = Omni(
         model=model_name,
-        log_stats=args.enable_stats,
+        log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
         batch_timeout=args.batch_timeout,
         init_timeout=args.init_timeout,
@@ -439,7 +439,7 @@ def parse_args():
         help="Query type.",
     )
     parser.add_argument(
-        "--enable-stats",
+        "--log-stats",
         action="store_true",
         default=False,
         help="Enable writing detailed statistics (default: disabled)",

examples/offline_inference/qwen3_omni/end2end.py

@@ -226,37 +226,6 @@ def get_multi_audios_query() -> QueryResult:
     )
 
 
-# def get_use_audio_in_video_query(video_path: str | None = None) -> QueryResult:
-# question = (
-#     "Describe the content of the video in details, then convert what the "
-#     "baby say into text."
-# )
-# prompt = (
-#     f"<|im_start|>system\n{default_system}<|im_end|>\n"
-#     "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
-#     f"{question}<|im_end|>\n"
-#     f"<|im_start|>assistant\n"
-# )
-# if video_path:
-#     if not os.path.exists(video_path):
-#         raise FileNotFoundError(f"Video file not found: {video_path}")
-#     video_frames = video_to_ndarrays(video_path, num_frames=16)
-# else:
-#     video_frames = VideoAsset(name="baby_reading", num_frames=16).np_ndarrays
-# audio = extract_video_audio(video_path, sampling_rate=16000)
-# return QueryResult(
-#     inputs={
-#         "prompt": prompt,
-#         "multi_modal_data": {
-#             "video": video_frames,
-#             "audio": audio,
-#         },
-#         "mm_processor_kwargs": {
-#             "use_audio_in_video": True,
-#         },
-#     },
-#     limit_mm_per_prompt={"audio": 1, "video": 1},
-# )
 def get_use_audio_in_video_query() -> QueryResult:
     question = "Describe the content of the video in details, then convert what the baby say into text."
     prompt = (
@@ -328,7 +297,7 @@ def main(args):
     omni_llm = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
-        log_stats=args.enable_stats,
+        log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
 
@@ -455,7 +424,7 @@ def parse_args():
         help="Query type.",
     )
     parser.add_argument(
-        "--enable-stats",
+        "--log-stats",
         action="store_true",
         default=False,
         help="Enable writing detailed statistics (default: disabled)",

examples/offline_inference/qwen3_tts/end2end.py

@@ -219,7 +219,7 @@ def main(args):
     omni = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
-        log_stats=args.enable_stats,
+        log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
 
@@ -275,7 +275,7 @@ def parse_args():
         help="Query type.",
     )
     parser.add_argument(
-        "--enable-stats",
+        "--log-stats",
         action="store_true",
         default=False,
         help="Enable writing detailed statistics (default: disabled)",

examples/offline_inference/text_to_audio/README.md

@@ -18,20 +18,20 @@ If you use a gated model (e.g., `stabilityai/stable-audio-open-1.0`), ensure you
 python text_to_audio.py \
   --model stabilityai/stable-audio-open-1.0 \
   --prompt "The sound of a hammer hitting a wooden surface" \
-  --negative_prompt "Low quality" \
+  --negative-prompt "Low quality" \
   --seed 42 \
-  --guidance_scale 7.0 \
-  --audio_length 10.0 \
-  --num_inference_steps 100 \
+  --guidance-scale 7.0 \
+  --audio-length 10.0 \
+  --num-inference-steps 100 \
   --output stable_audio_output.wav
 ```
 
 Key arguments:
 
 - `--prompt`: text description (string).
-- `--negative_prompt`: negative prompt for classifier-free guidance.
+- `--negative-prompt`: negative prompt for classifier-free guidance.
 - `--seed`: integer seed for deterministic generation.
-- `--guidance_scale`: classifier-free guidance scale.
-- `--audio_length`: audio duration in seconds.
-- `--num_inference_steps`: diffusion sampling steps.(more steps = higher quality, slower).
+- `--guidance-scale`: classifier-free guidance scale.
+- `--audio-length`: audio duration in seconds.
+- `--num-inference-steps`: diffusion sampling steps.(more steps = higher quality, slower).
 - `--output`: path to save the generated WAV file.

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -9,8 +9,8 @@ the Stable Audio Open model with vLLM-Omni.
 
 Usage:
     python text_to_audio.py --prompt "The sound of a dog barking"
-    python text_to_audio.py --prompt "A piano playing a gentle melody" --audio_length 10.0
-    python text_to_audio.py --prompt "Thunder and rain sounds" --negative_prompt "Low quality"
+    python text_to_audio.py --prompt "A piano playing a gentle melody" --audio-length 10.0
+    python text_to_audio.py --prompt "Thunder and rain sounds" --negative-prompt "Low quality"
 """
 
 import argparse
@@ -38,7 +38,7 @@ def parse_args() -> argparse.Namespace:
         help="Text prompt for audio generation.",
     )
     parser.add_argument(
-        "--negative_prompt",
+        "--negative-prompt",
         default="Low quality.",
         help="Negative prompt for classifier-free guidance.",
     )
@@ -49,31 +49,31 @@ def parse_args() -> argparse.Namespace:
         help="Random seed for deterministic results.",
     )
     parser.add_argument(
-        "--guidance_scale",
+        "--guidance-scale",
         type=float,
         default=7.0,
         help="Classifier-free guidance scale.",
     )
     parser.add_argument(
-        "--audio_start",
+        "--audio-start",
         type=float,
         default=0.0,
         help="Audio start time in seconds.",
     )
     parser.add_argument(
-        "--audio_length",
+        "--audio-length",
         type=float,
         default=10.0,
         help="Audio length in seconds (max ~47s for stable-audio-open-1.0).",
     )
     parser.add_argument(
-        "--num_inference_steps",
+        "--num-inference-steps",
         type=int,
         default=100,
         help="Number of denoising steps for the diffusion sampler.",
     )
     parser.add_argument(
-        "--num_waveforms",
+        "--num-waveforms",
         type=int,
         default=1,
         help="Number of audio waveforms to generate for the given prompt.",
@@ -85,7 +85,7 @@ def parse_args() -> argparse.Namespace:
         help="Path to save the generated audio (WAV format).",
     )
     parser.add_argument(
-        "--sample_rate",
+        "--sample-rate",
         type=int,
         default=44100,
         help="Sample rate for output audio (Stable Audio uses 44100 Hz).",

examples/offline_inference/text_to_image/README.md

@@ -79,9 +79,9 @@ python text_to_image.py \
   --model Tongyi-MAI/Z-Image-Turbo \
   --prompt "a cup of coffee on the table" \
   --seed 42 \
-  --cfg_scale 4.0 \
-  --num_images_per_prompt 1 \
-  --num_inference_steps 50 \
+  --cfg-scale 4.0 \
+  --num-images-per-prompt 1 \
+  --num-inference-steps 50 \
   --height 1024 \
   --width 1024 \
   --output outputs/coffee.png
@@ -91,17 +91,17 @@ Key arguments:
 
 - `--prompt`: text description (string).
 - `--seed`: integer seed for deterministic sampling.
-- `--cfg_scale`: true CFG scale (model-specific guidance strength).
-- `--num_images_per_prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
-- `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
+- `--cfg-scale`: true CFG scale (model-specific guidance strength).
+- `--num-images-per-prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
+- `--num-inference-steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--height/--width`: output resolution (defaults 1024x1024).
 - `--output`: path to save the generated PNG.
-- `--vae_use_slicing`: enable VAE slicing for memory optimization.
-- `--vae_use_tiling`: enable VAE tiling for memory optimization.
-- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--vae-use-slicing`: enable VAE slicing for memory optimization.
+- `--vae-use-tiling`: enable VAE tiling for memory optimization.
+- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 
-> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
+> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
 
 > ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
 

examples/offline_inference/text_to_image/gradio_demo.py

@@ -40,7 +40,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--default-seed", type=int, default=42, help="Initial seed shown in UI.")
     parser.add_argument("--default-cfg-scale", type=float, default=4.0, help="Initial CFG scale shown in UI.")
     parser.add_argument(
-        "--num_inference_steps",
+        "--num-inference-steps",
         type=int,
         default=50,
         help="Default number of denoising steps shown in the UI.",

examples/offline_inference/text_to_image/text_to_image.py

@@ -5,6 +5,7 @@ import argparse
 import os
 import time
 from pathlib import Path
+from typing import Any
 
 import torch
 
@@ -25,19 +26,19 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
     parser.add_argument(
-        "--negative_prompt",
+        "--negative-prompt",
         default=None,
         help="negative prompt for classifier-free conditional guidance.",
     )
     parser.add_argument("--seed", type=int, default=142, help="Random seed for deterministic results.")
     parser.add_argument(
-        "--cfg_scale",
+        "--cfg-scale",
         type=float,
         default=4.0,
         help="True classifier-free guidance scale specific to Qwen-Image.",
     )
     parser.add_argument(
-        "--guidance_scale",
+        "--guidance-scale",
         type=float,
         default=1.0,
         help="Classifier-free guidance scale.",
@@ -51,19 +52,19 @@ def parse_args() -> argparse.Namespace:
         help="Path to save the generated image (PNG).",
     )
     parser.add_argument(
-        "--num_images_per_prompt",
+        "--num-images-per-prompt",
         type=int,
         default=1,
         help="Number of images to generate for the given prompt.",
     )
     parser.add_argument(
-        "--num_inference_steps",
+        "--num-inference-steps",
         type=int,
         default=50,
         help="Number of denoising steps for the diffusion sampler.",
     )
     parser.add_argument(
-        "--cache_backend",
+        "--cache-backend",
         type=str,
         default=None,
         choices=["cache_dit", "tea_cache"],
@@ -79,26 +80,26 @@ def parse_args() -> argparse.Namespace:
         help="Enable cache-dit summary logging after diffusion forward passes.",
     )
     parser.add_argument(
-        "--ulysses_degree",
+        "--ulysses-degree",
         type=int,
         default=1,
         help="Number of GPUs used for ulysses sequence parallelism.",
     )
     parser.add_argument(
-        "--ring_degree",
+        "--ring-degree",
         type=int,
         default=1,
         help="Number of GPUs used for ring sequence parallelism.",
     )
     parser.add_argument(
-        "--cfg_parallel_size",
+        "--cfg-parallel-size",
         type=int,
         default=1,
         choices=[1, 2],
         help="Number of GPUs used for classifier free guidance parallel size.",
     )
     parser.add_argument(
-        "--enforce_eager",
+        "--enforce-eager",
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
@@ -113,26 +114,44 @@ def parse_args() -> argparse.Namespace:
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
     parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8"],
+        help="Quantization method for the transformer. "
+        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs). "
+        "Default: None (no quantization, uses BF16).",
     )
     parser.add_argument(
-        "--tensor_parallel_size",
+        "--ignored-layers",
+        type=str,
+        default=None,
+        help="Comma-separated list of layer name patterns to skip quantization. "
+        "Only used when --quantization is set. "
+        "Available layers: to_qkv, to_out, add_kv_proj, to_add_out, img_mlp, txt_mlp, proj_out. "
+        "Example: --ignored-layers 'add_kv_proj,to_add_out'",
+    )
+    parser.add_argument(
+        "--vae-use-slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae-use-tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
+    parser.add_argument(
+        "--tensor-parallel-size",
         type=int,
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
     parser.add_argument(
-        "--vae_use_slicing",
-        action="store_true",
-        help="Enable VAE slicing for memory optimization.",
-    )
-    parser.add_argument(
-        "--vae_use_tiling",
-        action="store_true",
-        help="Enable VAE tiling for memory optimization.",
+        "--vae-patch-parallel-size",
+        type=int,
+        default=1,
+        help="Number of ranks used for VAE patch/tile parallelism (decode/encode).",
     )
     return parser.parse_args()
 
@@ -176,15 +195,27 @@ def main():
         ring_degree=args.ring_degree,
         cfg_parallel_size=args.cfg_parallel_size,
         tensor_parallel_size=args.tensor_parallel_size,
+        vae_patch_parallel_size=args.vae_patch_parallel_size,
     )
 
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
+    # Build quantization kwargs: use quantization_config dict when
+    # ignored_layers is specified so the list flows through OmniDiffusionConfig
+    quant_kwargs: dict[str, Any] = {}
+    ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None
+    if args.quantization and ignored_layers:
+        quant_kwargs["quantization_config"] = {
+            "method": args.quantization,
+            "ignored_layers": ignored_layers,
+        }
+    elif args.quantization:
+        quant_kwargs["quantization"] = args.quantization
+
     omni = Omni(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
@@ -193,6 +224,7 @@ def main():
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
+        **quant_kwargs,
     )
 
     if profiler_enabled:
@@ -205,10 +237,15 @@ def main():
     print(f"  Model: {args.model}")
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
+    print(f"  Quantization: {args.quantization if args.quantization else 'None (BF16)'}")
+    if ignored_layers:
+        print(f"  Ignored layers: {ignored_layers}")
     print(
         f"  Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
-        f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
+        f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
+        f"vae_patch_parallel_size={args.vae_patch_parallel_size}"
     )
+    print(f"  CPU offload: {args.enable_cpu_offload}")
     print(f"  Image size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
 

examples/offline_inference/text_to_video/text_to_video.md

@@ -7,14 +7,14 @@ The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text
 ```bash
 python text_to_video.py \
   --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-  --negative_prompt "<optional quality filter>" \
+  --negative-prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
-  --num_frames 33 \
-  --guidance_scale 4.0 \
-  --guidance_scale_high 3.0 \
-  --flow_shift 12.0 \
-  --num_inference_steps 40 \
+  --num-frames 33 \
+  --guidance-scale 4.0 \
+  --guidance-scale-high 3.0 \
+  --flow-shift 12.0 \
+  --num-inference-steps 40 \
   --fps 16 \
   --output t2v_out.mp4
 ```
@@ -23,15 +23,15 @@ Key arguments:
 
 - `--prompt`: text description (string).
 - `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
-- `--num_frames`: Number of frames (Wan default is 81).
-- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high).
-- `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
-- `--boundary_ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
+- `--num-frames`: Number of frames (Wan default is 81).
+- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high).
+- `--negative-prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
+- `--boundary-ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
-- `--vae_use_slicing`: enable VAE slicing for memory optimization.
-- `--vae_use_tiling`: enable VAE tiling for memory optimization.
-- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--vae-use-slicing`: enable VAE slicing for memory optimization.
+- `--vae-use-tiling`: enable VAE tiling for memory optimization.
+- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 
-> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
+> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/text_to_video/text_to_video.py

@@ -24,25 +24,25 @@ def parse_args() -> argparse.Namespace:
         help="Diffusers Wan2.2 model ID or local path.",
     )
     parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
-    parser.add_argument("--negative_prompt", default="", help="Negative prompt.")
+    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    parser.add_argument("--guidance_scale", type=float, default=4.0, help="CFG scale (applied to low/high).")
-    parser.add_argument("--guidance_scale_high", type=float, default=None, help="Optional separate CFG for high-noise.")
+    parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (applied to low/high).")
+    parser.add_argument("--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise.")
     parser.add_argument("--height", type=int, default=720, help="Video height.")
     parser.add_argument("--width", type=int, default=1280, help="Video width.")
-    parser.add_argument("--num_frames", type=int, default=81, help="Number of frames (Wan default is 81).")
-    parser.add_argument("--num_inference_steps", type=int, default=40, help="Sampling steps.")
+    parser.add_argument("--num-frames", type=int, default=81, help="Number of frames (Wan default is 81).")
+    parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.")
     parser.add_argument(
-        "--boundary_ratio",
+        "--boundary-ratio",
         type=float,
         default=0.875,
         help="Boundary split ratio for low/high DiT. Default 0.875 uses both transformers for best quality. Set to 1.0 to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited).",
     )
     parser.add_argument(
-        "--flow_shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
+        "--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
     )
     parser.add_argument(
-        "--cache_backend",
+        "--cache-backend",
         type=str,
         default=None,
         choices=["cache_dit"],
@@ -60,17 +60,17 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
     parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
     parser.add_argument(
-        "--vae_use_slicing",
+        "--vae-use-slicing",
         action="store_true",
         help="Enable VAE slicing for memory optimization.",
     )
     parser.add_argument(
-        "--vae_use_tiling",
+        "--vae-use-tiling",
         action="store_true",
         help="Enable VAE tiling for memory optimization.",
     )
     parser.add_argument(
-        "--enforce_eager",
+        "--enforce-eager",
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
@@ -85,32 +85,26 @@ def parse_args() -> argparse.Namespace:
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
     parser.add_argument(
-        "--layerwise-num-gpu-layers",
-        type=int,
-        default=1,
-        help="Number of ready layers (blocks) to keep on GPU during generation.",
-    )
-    parser.add_argument(
-        "--ulysses_degree",
+        "--ulysses-degree",
         type=int,
         default=1,
         help="Number of GPUs used for ulysses sequence parallelism.",
     )
     parser.add_argument(
-        "--ring_degree",
+        "--ring-degree",
         type=int,
         default=1,
         help="Number of GPUs used for ring sequence parallelism.",
     )
     parser.add_argument(
-        "--cfg_parallel_size",
+        "--cfg-parallel-size",
         type=int,
         default=1,
         choices=[1, 2],
         help="Number of GPUs used for classifier free guidance parallel size.",
     )
     parser.add_argument(
-        "--tensor_parallel_size",
+        "--tensor-parallel-size",
         type=int,
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
@@ -155,7 +149,6 @@ def main():
     omni = Omni(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
-        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,

examples/online_serving/bagel/README.md

@@ -22,12 +22,29 @@ cd /workspace/vllm-omni/examples/online_serving/bagel
 bash run_server.sh
 ```
 
-If you have a custom stage configs file, launch the server with the command below:
-
 ```bash
 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
 ```
 
+#### 🚀 Tensor Parallelism (TP)
+
+For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) for the server.
+
+1. **Modify Stage Config**: Create or modify a stage configuration yaml (e.g., [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)). Set `tensor_parallel_size` to `2` (or more) and update `devices` to include multiple GPU IDs (e.g., `"0,1"`).
+
+```yaml
+    engine_args:
+      tensor_parallel_size: 2
+      ...
+    runtime:
+      devices: "0,1"
+```
+
+2. **Launch Server**:
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/your/custom_bagel.yaml
+```
+
 ### Send Multi-modal Request
 
 Get into the bagel folder:

examples/online_serving/image_to_image/openai_chat_client.py

@@ -76,6 +76,10 @@ def edit_image(
 
     # Build extra_body with generation parameters
     extra_body = {}
+    if height is not None:
+        extra_body["height"] = height
+    if width is not None:
+        extra_body["width"] = width
     if steps is not None:
         extra_body["num_inference_steps"] = steps
     if guidance_scale is not None:

examples/online_serving/image_to_image/run_curl_image_edit.sh

@@ -19,26 +19,25 @@ if [[ ! -f "$INPUT_IMG" ]]; then
   exit 1
 fi
 
-IMG_B64=$(base64 -w0 "$INPUT_IMG")
+REQUEST_JSON_FILE=$(mktemp)
+trap 'rm -f "$REQUEST_JSON_FILE"' EXIT
 
-REQUEST_JSON=$(
-  jq -n --arg prompt "$PROMPT" --arg img "$IMG_B64" '{
+# Pipe base64 into jq via stdin to avoid ARG_MAX limit on large images
+base64 -w0 "$INPUT_IMG" \
+  | jq -Rs --arg prompt "$PROMPT" '{
     messages: [{
       role: "user",
       content: [
         {"type": "text", "text": $prompt},
-        {"type": "image_url", "image_url": {"url": ("data:image/png;base64," + $img)}}
+        {"type": "image_url", "image_url": {"url": ("data:image/png;base64," + .)}}
       ]
     }],
     extra_body: {
-      height: 1024,
-      width: 1024,
       num_inference_steps: 50,
       guidance_scale: 1,
       seed: 42
     }
-  }'
-)
+  }' > "$REQUEST_JSON_FILE"
 
 echo "Generating edited image..."
 echo "Server: $SERVER"
@@ -48,7 +47,7 @@ echo "Output: $OUTPUT"
 
 curl -s "$SERVER/v1/chat/completions" \
   -H "Content-Type: application/json" \
-  -d "$REQUEST_JSON" \
+  -d @"$REQUEST_JSON_FILE" \
   | jq -r '.choices[0].message.content[0].image_url.url' \
   | cut -d',' -f2 \
   | base64 -d > "$OUTPUT"

examples/online_serving/image_to_video/README.md

@@ -0,0 +1,83 @@
+# Image-To-Video
+
+This example demonstrates how to deploy the Wan2.2 image-to-video model for online video generation using vLLM-Omni.
+
+## Start Server
+
+### Basic Start
+
+```bash
+vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --port 8091
+```
+
+### Start with Parameters
+
+Or use the startup script:
+
+```bash
+bash run_server.sh
+```
+
+The script allows overriding:
+- `MODEL` (default: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`)
+- `PORT` (default: `8091`)
+- `BOUNDARY_RATIO` (default: `0.875`)
+- `FLOW_SHIFT` (default: `12.0`)
+- `CACHE_BACKEND` (default: `none`)
+- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
+
+## API Calls
+
+### Method 1: Using curl
+
+```bash
+# Basic image-to-video generation
+bash run_curl_image_to_video.sh
+
+# Or execute directly (OpenAI-style multipart)
+curl -X POST http://localhost:8091/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "input_reference=@/path/to/qwen-bear.png" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=12.0" \
+  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
+```
+
+## Request Format
+
+### Required Fields
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "input_reference=@/path/to/qwen-bear.png"
+```
+
+### Generation with Parameters
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "input_reference=@/path/to/qwen-bear.png" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=12.0" \
+  -F "seed=42"
+```

examples/online_serving/image_to_video/run_curl_image_to_video.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Wan2.2 image-to-video curl example (OpenAI-style multipart)
+
+INPUT_IMAGE="${INPUT_IMAGE:-../../offline_inference/image_to_video/qwen-bear.png}"
+OUTPUT_PATH="${OUTPUT_PATH:-wan22_i2v_output.mp4}"
+NEGATIVE_PROMPT="${NEGATIVE_PROMPT:-}"
+
+if [ ! -f "$INPUT_IMAGE" ]; then
+    echo "Input image not found: $INPUT_IMAGE"
+    exit 1
+fi
+
+NEGATIVE_PROMPT_FLAG=""
+if [ -n "$NEGATIVE_PROMPT" ]; then
+    NEGATIVE_PROMPT_FLAG="-F negative_prompt=${NEGATIVE_PROMPT}"
+fi
+
+curl -X POST http://localhost:8099/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  $NEGATIVE_PROMPT_FLAG \
+  -F "input_reference=@${INPUT_IMAGE}" \
+  -F "seconds=2" \
+  -F "size=832x480" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=1.0" \
+  -F "guidance_scale_2=1.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=12.0" \
+  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
+
+echo "Saved video to ${OUTPUT_PATH}"

examples/online_serving/image_to_video/run_server.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+# Wan2.2 image-to-video server start script
+
+MODEL="${MODEL:-Wan-AI/Wan2.2-I2V-A14B-Diffusers}"
+PORT="${PORT:-8099}"
+CACHE_BACKEND="${CACHE_BACKEND:-none}"
+ENABLE_CACHE_DIT_SUMMARY="${ENABLE_CACHE_DIT_SUMMARY:-0}"
+
+echo "Starting Wan2.2 I2V server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+echo "Cache backend: $CACHE_BACKEND"
+if [ "$ENABLE_CACHE_DIT_SUMMARY" != "0" ]; then
+    echo "Cache-DiT summary: enabled"
+fi
+
+CACHE_BACKEND_FLAG=""
+if [ "$CACHE_BACKEND" != "none" ]; then
+    CACHE_BACKEND_FLAG="--cache-backend $CACHE_BACKEND"
+fi
+
+vllm serve "$MODEL" --omni \
+    --port "$PORT" \
+    $CACHE_BACKEND_FLAG \
+    $(if [ "$ENABLE_CACHE_DIT_SUMMARY" != "0" ]; then echo "--enable-cache-dit-summary"; fi)

examples/online_serving/qwen2_5_omni/README.md

@@ -27,7 +27,7 @@ cd examples/online_serving/qwen2_5_omni
 #### Send request via python
 
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
+python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities --port 8091 --host "localhost"
 ```
 
 The Python client supports the following command-line arguments:

examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -6,16 +6,6 @@ from openai import OpenAI
 from vllm.assets.audio import AudioAsset
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-# Modify OpenAI's API key and API base to use vLLM's API server.
-openai_api_key = "EMPTY"
-openai_api_base = "http://localhost:8091/v1"
-
-client = OpenAI(
-    # defaults to os.environ.get("OPENAI_API_KEY")
-    api_key=openai_api_key,
-    base_url=openai_api_base,
-)
-
 SEED = 42
 
 
@@ -276,7 +266,7 @@ query_map = {
 }
 
 
-def run_multimodal_generation(args) -> None:
+def run_multimodal_generation(args, client: OpenAI) -> None:
     model_name = "Qwen/Qwen2.5-Omni-7B"
     thinker_sampling_params = {
         "temperature": 0.0,  # Deterministic - no randomness
@@ -442,10 +432,28 @@ def parse_args():
         action="store_true",
         help="Stream the response.",
     )
-
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8091,
+        help="Port of the vLLM Omni API server.",
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="localhost",
+        help="Host/IP of the vLLM Omni API server.",
+    )
     return parser.parse_args()
 
 
 if __name__ == "__main__":
     args = parse_args()
-    run_multimodal_generation(args)
+    host = args.host
+    port = args.port
+    openai_api_base = f"http://{host}:{port}/v1"
+    client = OpenAI(
+        api_key="EMPTY",
+        base_url=openai_api_base,
+    )
+    run_multimodal_generation(args, client)

examples/online_serving/qwen3_omni/README.md

@@ -33,7 +33,7 @@ cd examples/online_serving/qwen3_omni
 ####  Send request via python
 
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
+python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image --port 8091 --host "localhost"
 ```
 
 The Python client supports the following command-line arguments:

examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -8,16 +8,6 @@ from openai import OpenAI
 from vllm.assets.audio import AudioAsset
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
-# Modify OpenAI's API key and API base to use vLLM's API server.
-openai_api_key = "EMPTY"
-openai_api_base = "http://localhost:8091/v1"
-
-client = OpenAI(
-    # defaults to os.environ.get("OPENAI_API_KEY")
-    api_key=openai_api_key,
-    base_url=openai_api_base,
-)
-
 SEED = 42
 
 
@@ -342,7 +332,7 @@ query_map = {
 }
 
 
-def run_multimodal_generation(args) -> None:
+def run_multimodal_generation(args, client: OpenAI) -> None:
     model_name = args.model
     thinker_sampling_params = {
         "temperature": 0.4,  # Deterministic
@@ -544,10 +534,28 @@ def parse_args():
         default=1,
         help="Number of concurrent requests to send. Default is 1.",
     )
-
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8091,
+        help="Port of the vLLM Omni API server.",
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="localhost",
+        help="Host/IP of the vLLM Omni API server.",
+    )
     return parser.parse_args()
 
 
 if __name__ == "__main__":
     args = parse_args()
-    run_multimodal_generation(args)
+    host = args.host
+    port = args.port
+    openai_api_base = f"http://{host}:{port}/v1"
+    client = OpenAI(
+        api_key="EMPTY",
+        base_url=openai_api_base,
+    )
+    run_multimodal_generation(args, client)

examples/online_serving/qwen3_tts/README.md

@@ -9,6 +9,8 @@ This directory contains examples for running Qwen3-TTS models with vLLM-Omni's o
 | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
 | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
 | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
+| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
 
 ## Quick Start
 
@@ -24,19 +26,27 @@ This directory contains examples for running Qwen3-TTS models with vLLM-Omni's o
 ./run_server.sh Base
 ```
 
+Or launch directly with vllm serve:
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
+    --omni --port 8091 --trust-remote-code --enforce-eager
+```
+
 ### 2. Run the Client
 
 ```bash
 # CustomVoice: Use predefined speaker
 python openai_speech_client.py \
     --text "你好，我是通义千问" \
-    --voice Vivian \
+    --voice vivian \
     --language Chinese
 
 # CustomVoice with style instruction
 python openai_speech_client.py \
     --text "今天天气真好" \
-    --voice Ryan \
+    --voice ryan \
     --instructions "用开心的语气说"
 
 # VoiceDesign: Describe the voice style
@@ -59,25 +69,25 @@ python openai_speech_client.py \
 
 ```bash
 # Simple TTS request
-curl -X POST http://localhost:8000/v1/audio/speech \
+curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
         "input": "Hello, how are you?",
-        "voice": "Vivian",
+        "voice": "vivian",
         "language": "English"
     }' --output output.wav
 
 # With style instruction
-curl -X POST http://localhost:8000/v1/audio/speech \
+curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
         "input": "I am so excited!",
-        "voice": "Vivian",
+        "voice": "vivian",
         "instructions": "Speak with great enthusiasm"
     }' --output excited.wav
 
 # List available voices in CustomVoice models
-curl http://localhost:8000/v1/audio/voices
+curl http://localhost:8091/v1/audio/voices
 ```
 
 ## API Reference
@@ -90,12 +100,26 @@ POST /v1/audio/speech
 
 This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
 
+### Voices Endpoint
+
+```
+GET /v1/audio/voices
+```
+
+Lists available voices for the loaded model:
+
+```json
+{
+    "voices": ["aiden", "dylan", "eric", "one_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
+}
+```
+
 ### Request Body
 
 ```json
 {
     "input": "Text to synthesize",
-    "voice": "Vivian",
+    "voice": "vivian",
     "response_format": "wav",
     "task_type": "CustomVoice",
     "language": "Auto",
@@ -115,45 +139,64 @@ Returns audio data in the requested format (default: WAV).
 
 ## Parameters
 
-### Standard OpenAI Parameters
+### OpenAI Standard Parameters
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `input` | string | required | Text to synthesize |
-| `voice` | string | "Vivian" | Speaker/voice name |
+| `input` | string | **required** | Text to synthesize |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
 | `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
 | `speed` | float | 1.0 | Playback speed (0.25-4.0) |
-| `model` | string | optional | Model name (optional when serving single model) |
 
-### Qwen3-TTS Parameters
+### vLLM-Omni Extension Parameters
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
 | `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base |
-| `language` | string | "Auto" | Language: Auto, Chinese, English, Japanese, Korean |
+| `language` | string | "Auto" | Language (see supported languages below) |
 | `instructions` | string | "" | Voice style/emotion instructions |
 | `max_new_tokens` | int | 2048 | Maximum tokens to generate |
 
+**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
+
 ### Voice Clone Parameters (Base task)
 
-| Parameter | Type | Required | Description |
+| Parameter | Type | Default | Description |
 |-----------|------|----------|-------------|
-| `ref_audio` | string | Yes* | Reference audio (file path, URL, or base64) |
-| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
-| `x_vector_only_mode` | bool | false | Use speaker embedding only (no ICL) |
+| `ref_audio` | string | null | Reference audio (URL or base64 data URL) |
+| `ref_text` | string | null | Transcript of reference audio |
+| `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
 
 ## Python Usage
 
+### Using OpenAI SDK
+
+```python
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
+
+response = client.audio.speech.create(
+    model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+    voice="vivian",
+    input="Hello, how are you?",
+)
+
+response.stream_to_file("output.wav")
+```
+
+### Using httpx
+
 ```python
 import httpx
 
-# Simple request
 response = httpx.post(
-    "http://localhost:8000/v1/audio/speech",
+    "http://localhost:8091/v1/audio/speech",
     json={
-        "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
-        "input": "Hello world",
-        "voice": "Vivian",
+        "input": "Hello, how are you?",
+        "voice": "vivian",
+        "language": "English",
     },
     timeout=300.0,
 )
@@ -169,7 +212,8 @@ with open("output.wav", "wb") as f:
 
 ## Troubleshooting
 
-1. **Connection refused**: Make sure the server is running on the correct port
-2. **Out of memory**: Reduce `--gpu-memory-utilization` in run_server.sh
-3. **Unsupported speaker**: Check supported speakers via model documentation
-4. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning
+1. **"TTS model did not produce audio output"**: Ensure you're using the correct model variant for your task type (CustomVoice task → CustomVoice model, etc.)
+2. **Connection refused**: Make sure the server is running on the correct port
+3. **Out of memory**: Use smaller model variant (`Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`) or reduce `--gpu-memory-utilization`
+4. **Unsupported speaker**: Use `/v1/audio/voices` to list available voices for the loaded model
+5. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning

examples/online_serving/qwen3_tts/openai_speech_client.py

@@ -5,10 +5,10 @@ to generate audio from text using Qwen3-TTS models.
 
 Examples:
     # CustomVoice task (predefined speaker)
-    python openai_speech_client.py --text "Hello, how are you?" --voice Vivian
+    python openai_speech_client.py --text "Hello, how are you?" --voice vivian
 
     # CustomVoice with emotion instruction
-    python openai_speech_client.py --text "I'm so happy!" --voice Vivian \
+    python openai_speech_client.py --text "I'm so happy!" --voice vivian \
         --instructions "Speak with excitement"
 
     # VoiceDesign task (voice from description)
@@ -30,7 +30,7 @@ import os
 import httpx
 
 # Default server configuration
-DEFAULT_API_BASE = "http://localhost:8000"
+DEFAULT_API_BASE = "http://localhost:8091"
 DEFAULT_API_KEY = "EMPTY"
 
 
@@ -178,8 +178,8 @@ def parse_args():
     parser.add_argument(
         "--voice",
         type=str,
-        default="Vivian",
-        help="Speaker/voice name (default: Vivian). Options: Vivian, Ryan, etc.",
+        default="vivian",
+        help="Speaker/voice name (default: vivian). Options: vivian, ryan, aiden, etc.",
     )
     parser.add_argument(
         "--language",

examples/online_serving/qwen3_tts/run_server.sh

@@ -33,7 +33,7 @@ echo "Starting Qwen3-TTS server with model: $MODEL"
 vllm-omni serve "$MODEL" \
     --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
     --host 0.0.0.0 \
-    --port 8000 \
+    --port 8091 \
     --gpu-memory-utilization 0.9 \
     --trust-remote-code \
     --enforce-eager \

examples/online_serving/text_to_video/README.md

@@ -0,0 +1,116 @@
+# Text-To-Video
+
+This example demonstrates how to deploy the Wan2.2 text-to-video model for online video generation using vLLM-Omni.
+
+## Start Server
+
+### Basic Start
+
+```bash
+vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091
+```
+
+### Start with Parameters
+
+Or use the startup script:
+
+```bash
+bash run_server.sh
+```
+
+The script allows overriding:
+- `MODEL` (default: `Wan-AI/Wan2.2-T2V-A14B-Diffusers`)
+- `PORT` (default: `8091`)
+- `BOUNDARY_RATIO` (default: `0.875`)
+- `FLOW_SHIFT` (default: `5.0`)
+- `CACHE_BACKEND` (default: `none`)
+- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
+
+## API Calls
+
+### Method 1: Using curl
+
+```bash
+# Basic text-to-video generation
+bash run_curl_text_to_video.sh
+
+# Or execute directly (OpenAI-style multipart)
+curl -s http://localhost:8091/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "negative_prompt=色调艳丽 ，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.0" \
+  -F "guidance_scale_2=4.0" \
+  -F "boundary_ratio=0.875" \
+  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+```
+
+## Request Format
+
+### Simple Text-to-Video Generation
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A cinematic view of a futuristic city at sunset"
+```
+
+### Generation with Parameters
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A cinematic view of a futuristic city at sunset" \
+  -F "width=832" \
+  -F "height=480" \
+  -F "num_frames=33" \
+  -F "negative_prompt=low quality, blurry, static" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.0" \
+  -F "guidance_scale_2=4.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=5.0" \
+  -F "seed=42"
+```
+
+## Generation Parameters
+
+| Parameter             | Type   | Default | Description                                      |
+| --------------------- | ------ | ------- | ------------------------------------------------ |
+| `prompt`              | str    | -       | Text description of the desired video            |
+| `negative_prompt`     | str    | None    | Negative prompt                                  |
+| `n`                   | int    | 1       | Number of videos to generate                     |
+| `width`               | int    | None    | Video width in pixels                            |
+| `height`              | int    | None    | Video height in pixels                           |
+| `num_frames`          | int    | None    | Number of frames to generate                     |
+| `fps`                 | int    | None    | Frames per second for output video               |
+| `num_inference_steps` | int    | None    | Number of denoising steps                        |
+| `guidance_scale`      | float  | None    | CFG guidance scale (low-noise stage)             |
+| `guidance_scale_2`    | float  | None    | CFG guidance scale (high-noise stage, Wan2.2)     |
+| `boundary_ratio`      | float  | None    | Boundary split ratio for low/high DiT (Wan2.2)   |
+| `flow_shift`          | float  | None    | Scheduler flow shift (Wan2.2)                    |
+| `seed`                | int    | None    | Random seed (reproducible)                       |
+| `lora`                | object | None    | LoRA configuration                               |
+| `extra_body`          | object | None    | Model-specific extra parameters                  |
+
+## Response Format
+
+```json
+{
+  "created": 1234567890,
+  "data": [
+    { "b64_json": "<base64-mp4>" }
+  ]
+}
+```
+
+## Extract Video
+
+```bash
+# Extract base64 from response and decode to video
+cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+```

examples/online_serving/text_to_video/run_curl_text_to_video.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+# Wan2.2 text-to-video curl example
+
+OUTPUT_PATH="wan22_output.mp4"
+
+curl -X POST http://localhost:8098/v1/videos \
+  -H "Accept: application/json" \
+  -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+  -F "seconds=2" \
+  -F "size=832x480" \
+  -F "negative_prompt=色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
+  -F "fps=16" \
+  -F "num_inference_steps=40" \
+  -F "guidance_scale=4.0" \
+  -F "guidance_scale_2=4.0" \
+  -F "boundary_ratio=0.875" \
+  -F "flow_shift=5.0" \
+  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
+
+echo "Saved video to ${OUTPUT_PATH}"

examples/online_serving/text_to_video/run_server.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+# Wan2.2 online serving startup script
+
+MODEL="${MODEL:-Wan-AI/Wan2.2-T2V-A14B-Diffusers}"
+PORT="${PORT:-8098}"
+BOUNDARY_RATIO="${BOUNDARY_RATIO:-0.875}"
+FLOW_SHIFT="${FLOW_SHIFT:-5.0}"
+CACHE_BACKEND="${CACHE_BACKEND:-none}"
+ENABLE_CACHE_DIT_SUMMARY="${ENABLE_CACHE_DIT_SUMMARY:-0}"
+
+echo "Starting Wan2.2 server..."
+echo "Model: $MODEL"
+echo "Port: $PORT"
+echo "Boundary ratio: $BOUNDARY_RATIO"
+echo "Flow shift: $FLOW_SHIFT"
+echo "Cache backend: $CACHE_BACKEND"
+if [ "$ENABLE_CACHE_DIT_SUMMARY" != "0" ]; then
+    echo "Cache-DiT summary: enabled"
+fi
+
+CACHE_BACKEND_FLAG=""
+if [ "$CACHE_BACKEND" != "none" ]; then
+    CACHE_BACKEND_FLAG="--cache-backend $CACHE_BACKEND"
+fi
+
+vllm serve "$MODEL" --omni \
+    --port "$PORT" \
+    --boundary-ratio "$BOUNDARY_RATIO" \
+    --flow-shift "$FLOW_SHIFT" \
+    $CACHE_BACKEND_FLAG \
+    $(if [ "$ENABLE_CACHE_DIT_SUMMARY" != "0" ]; then echo "--enable-cache-dit-summary"; fi)

mkdocs.yml

@@ -8,7 +8,7 @@ repo_url: https://github.com/vllm-project/vllm-omni
 edit_uri: edit/main/docs/
 
 # Copyright
-copyright: Copyright © 2025 vLLM-Omni Team
+copyright: Copyright © 2026 vLLM-Omni Team
 
 # Theme
 theme:
@@ -94,6 +94,7 @@ plugins:
       exclude:
         - "re:vllm_omni\\._.*"  # Internal modules
         - "vllm_omni.diffusion.models.qwen_image"  # avoid importing vllm in mkdocs building
+        - "vllm_omni.diffusion.quantization"  # avoid importing vllm in mkdocs building
         - "vllm_omni.entrypoints.async_diffusion"  # avoid importing vllm in mkdocs building
         - "vllm_omni.entrypoints.openai"  # avoid importing vllm in mkdocs building
         - "vllm_omni.entrypoints.openai.protocol"  # avoid importing vllm in mkdocs building

pyproject.toml

@@ -137,6 +137,7 @@ warn_unreachable = true
 strict_equality = true
 
 [tool.pytest.ini_options]
+asyncio_mode = "auto"
 testpaths = ["tests"]
 python_files = ["test_*.py", "*_test.py"]
 python_classes = ["Test*"]
@@ -147,7 +148,8 @@ addopts = [
 ]
 markers = [
     # ci/cd required
-    "core_model: Core model tests (run in each PR)",
+    "core_model: L1&L2 tests (run in each PR)",
+    "advanced_model: L3&L4 level tests (run in each merge or nightly)",
     # function module markers
     "diffusion: Diffusion model tests",
     "omni: Omni model tests",
@@ -175,6 +177,10 @@ markers = [
     "slow: Slow tests (may skip in quick CI)",
     "benchmark: Benchmark tests",
 ]
+filterwarnings = [
+    "ignore:.*does not have '__test__' attribute.*:UserWarning",
+    "ignore:.*does not have '__bases__' attribute.*:UserWarning",
+]
 
 [tool.typos.default]
 extend-ignore-identifiers-re = [
@@ -184,9 +190,10 @@ extend-ignore-identifiers-re = [
     ".*arange",
     ".*MoBA",
     ".*temperal_downsample",
+    ".*_nd",
     ".*nothink.*",
     ".*NOTHINK.*",
     ".*nin.*",
-    "Ono_Anna",
+    ".*[Oo]no_[Aa]nna.*",
     ".*cann.*",
 ]

requirements/common.txt

@@ -12,3 +12,5 @@ torchsde>=0.2.6
 openai-whisper>=20250625
 imageio[ffmpeg]>=2.37.2
 sox>=1.5.0
+prettytable>=3.8.0
+aenum==3.1.16

requirements/cuda.txt

@@ -1,3 +1,5 @@
 -r common.txt
 onnxruntime>=1.23.2
-fa3-fwd==0.0.1
+fa3-fwd==0.0.2
+# nvidia-cublas-cu12 is pinned to 12.9.1.4 via force-reinstall in Dockerfile.ci
+# because torch pins its own cublas version, making it unsatisfiable at resolve time.

tests/benchmarks/patch/test_patch.py

@@ -0,0 +1,578 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for patch.py
+"""
+
+import asyncio
+import json
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
+
+from vllm_omni.benchmarks.patch.patch import MixRequestFuncOutput, async_request_openai_chat_omni_completions
+
+pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
+
+
+class MockResponse:
+    """Mock aiohttp response for testing"""
+
+    def __init__(self, status, chunks, delay_between_chunks=0):
+        self.status = status
+        self.reason = "OK" if status == 200 else "Error"
+        self._chunks = chunks
+        self._delay = delay_between_chunks
+        self.content = self
+
+    async def iter_any(self):
+        for chunk in self._chunks:
+            if self._delay > 0:
+                await asyncio.sleep(self._delay)
+            yield chunk
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+
+def create_sse_chunk(data_dict):
+    """Helper to create SSE formatted chunk"""
+    return f"data: {json.dumps(data_dict)}\n\n".encode()
+
+
+# ============================================================================
+# output_tokens Tests
+# ============================================================================
+
+
+@pytest.mark.asyncio
+async def test_output_tokens_assigned_with_metrics():
+    """Test that output.output_tokens is assigned when metrics are present"""
+    # Arrange
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=10,
+        output_len=20,
+    )
+
+    # Create response chunks with metrics
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "Hello"}}],
+                "modality": "text",
+            }
+        ),
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": " world"}}],
+                "modality": "text",
+                "metrics": {"num_tokens_out": 42, "num_tokens_in": 10},
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+
+    mock_response = MockResponse(200, chunks)
+    mock_session = AsyncMock()
+    mock_session.post = MagicMock(return_value=mock_response)
+
+    # Act
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    # Assert
+    assert output.success is True
+    assert output.output_tokens == 42, "output_tokens should be assigned from metrics"
+    assert output.generated_text == "Hello world"
+
+
+@pytest.mark.asyncio
+async def test_output_tokens_not_assigned_without_metrics():
+    """Test that output.output_tokens defaults to 0 when no metrics present"""
+    # Arrange
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=10,
+        output_len=20,
+    )
+
+    # Create response chunks without metrics
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "Hello"}}],
+                "modality": "text",
+            }
+        ),
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": " world"}}],
+                "modality": "text",
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+
+    mock_response = MockResponse(200, chunks)
+    mock_session = AsyncMock()
+    mock_session.post = MagicMock(return_value=mock_response)
+
+    # Act
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    # Assert
+    assert output.success is True
+    assert output.output_tokens == 0, "output_tokens should default to 0 when no metrics"
+    assert output.generated_text == "Hello world"
+
+
+@pytest.mark.asyncio
+async def test_output_tokens_assigned_multiple_metrics():
+    """Test that output.output_tokens is updated with the latest metrics value"""
+    # Arrange
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=10,
+        output_len=20,
+    )
+
+    # Create response chunks with multiple metrics updates
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "Hello"}}],
+                "modality": "text",
+                "metrics": {"num_tokens_out": 5},
+            }
+        ),
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": " world"}}],
+                "modality": "text",
+                "metrics": {"num_tokens_out": 10},
+            }
+        ),
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "!"}}],
+                "modality": "text",
+                "metrics": {"num_tokens_out": 15},
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+
+    mock_response = MockResponse(200, chunks)
+    mock_session = AsyncMock()
+    mock_session.post = MagicMock(return_value=mock_response)
+
+    # Act
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    # Assert
+    assert output.success is True
+    assert output.output_tokens == 15, "output_tokens should be updated to the latest value"
+    assert output.generated_text == "Hello world!"
+
+
+@pytest.mark.asyncio
+async def test_output_tokens_with_audio_and_text():
+    """Test output_tokens assignment in mixed audio and text response"""
+    # Arrange
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=10,
+        output_len=20,
+    )
+
+    # Create response chunks with both audio and text, with metrics
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "Text response"}}],
+                "modality": "text",
+            }
+        ),
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": ""}}],
+                "modality": "audio",
+            }
+        ),
+        create_sse_chunk(
+            {
+                "modality": "text",
+                "metrics": {"num_tokens_out": 25, "num_tokens_in": 10},
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+
+    mock_response = MockResponse(200, chunks)
+    mock_session = AsyncMock()
+    mock_session.post = MagicMock(return_value=mock_response)
+
+    # Act
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    # Assert
+    assert output.success is True
+    assert output.output_tokens == 25, "output_tokens should be assigned even with audio modality"
+
+
+@pytest.mark.asyncio
+async def test_output_tokens_with_missing_num_tokens_out():
+    """Test that output_tokens defaults to 0 when num_tokens_out is missing in metrics"""
+    # Arrange
+    request_input = RequestFuncInput(
+        model="test-model",
+        model_name="test-model",
+        prompt="test prompt",
+        api_url="http://test.com/v1/chat/completions",
+        prompt_len=10,
+        output_len=20,
+    )
+
+    # Create response chunks with metrics but without num_tokens_out
+    chunks = [
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": "Hello"}}],
+                "modality": "text",
+            }
+        ),
+        create_sse_chunk(
+            {
+                "choices": [{"delta": {"content": " world"}}],
+                "modality": "text",
+                "metrics": {"num_tokens_in": 10},  # Missing num_tokens_out
+            }
+        ),
+        b"data: [DONE]\n\n",
+    ]
+
+    mock_response = MockResponse(200, chunks)
+    mock_session = AsyncMock()
+    mock_session.post = MagicMock(return_value=mock_response)
+
+    # Act
+    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+    # Assert
+    assert output.success is True
+    assert output.output_tokens == 0, "output_tokens should default to 0 when num_tokens_out is missing"
+
+
+@pytest.mark.asyncio
+async def test_output_tokens_initialization():
+    """Test that MixRequestFuncOutput initializes output_tokens correctly"""
+    # Arrange & Act
+    output = MixRequestFuncOutput()
+
+    # Assert
+    assert hasattr(output, "output_tokens"), "MixRequestFuncOutput should have output_tokens attribute"
+    assert output.output_tokens == 0, "output_tokens should be initialized to 0"
+
+
+# ============================================================================
+# text_latency Tests
+# ============================================================================
+
+
+class TestTextLatencyAttribute:
+    """Tests for text_latency attribute existence and assignment"""
+
+    def test_mix_request_func_output_has_text_latency(self):
+        """Test that MixRequestFuncOutput has text_latency attribute"""
+        output = MixRequestFuncOutput()
+        assert hasattr(output, "text_latency"), "MixRequestFuncOutput should have text_latency attribute"
+
+    def test_text_latency_initial_value(self):
+        """Test that text_latency initializes to a default value"""
+        output = MixRequestFuncOutput()
+        # Check if attribute exists and has a value (should be 0.0 or similar default)
+        text_latency = getattr(output, "text_latency", None)
+        assert text_latency is not None or hasattr(output, "text_latency"), "text_latency attribute should exist"
+
+    @pytest.mark.asyncio
+    async def test_text_latency_assigned_with_text_response(self):
+        """Test that text_latency is assigned when text response is received"""
+        request_input = RequestFuncInput(
+            model="test-model",
+            model_name="test-model",
+            prompt="test prompt",
+            api_url="http://test.com/v1/chat/completions",
+            prompt_len=10,
+            output_len=20,
+        )
+
+        # Create response chunks with text modality
+        chunks = [
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": "Hello"}}],
+                    "modality": "text",
+                }
+            ),
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": " world"}}],
+                    "modality": "text",
+                }
+            ),
+            b"data: [DONE]\n\n",
+        ]
+
+        mock_response = MockResponse(200, chunks)
+        mock_session = AsyncMock()
+        mock_session.post = MagicMock(return_value=mock_response)
+
+        # Act
+        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+        # Assert
+        assert output.success is True
+        assert hasattr(output, "text_latency"), "Output should have text_latency attribute"
+        assert output.text_latency > 0, "text_latency should be greater than 0 for text response"
+
+    @pytest.mark.asyncio
+    async def test_text_latency_updated_with_multiple_text_chunks(self):
+        """Test that text_latency is updated as more text chunks arrive"""
+        request_input = RequestFuncInput(
+            model="test-model",
+            model_name="test-model",
+            prompt="test prompt",
+            api_url="http://test.com/v1/chat/completions",
+            prompt_len=10,
+            output_len=20,
+        )
+
+        # Create response chunks with small delays to simulate streaming
+        chunks = [
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": "First"}}],
+                    "modality": "text",
+                }
+            ),
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": " second"}}],
+                    "modality": "text",
+                }
+            ),
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": " third"}}],
+                    "modality": "text",
+                }
+            ),
+            b"data: [DONE]\n\n",
+        ]
+
+        mock_response = MockResponse(200, chunks)
+        mock_session = AsyncMock()
+        mock_session.post = MagicMock(return_value=mock_response)
+
+        # Act
+        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+        # Assert
+        assert output.success is True
+        assert hasattr(output, "text_latency"), "Output should have text_latency attribute"
+        assert output.text_latency > 0, "text_latency should accumulate"
+        assert output.generated_text == "First second third"
+
+    @pytest.mark.asyncio
+    async def test_text_latency_with_only_audio_response(self):
+        """Test text_latency behavior when only audio is received"""
+        request_input = RequestFuncInput(
+            model="test-model",
+            model_name="test-model",
+            prompt="test prompt",
+            api_url="http://test.com/v1/chat/completions",
+            prompt_len=10,
+            output_len=20,
+        )
+
+        # Create response chunks with only audio modality (no text)
+        chunks = [
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": ""}}],
+                    "modality": "audio",
+                }
+            ),
+            b"data: [DONE]\n\n",
+        ]
+
+        mock_response = MockResponse(200, chunks)
+        mock_session = AsyncMock()
+        mock_session.post = MagicMock(return_value=mock_response)
+
+        # Act
+        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+        # Assert
+        assert output.success is True
+        assert hasattr(output, "text_latency"), "Output should have text_latency attribute even with audio-only"
+        # text_latency should either be 0 or the initial value since no text was processed
+        assert output.text_latency >= 0, "text_latency should be non-negative"
+
+    @pytest.mark.asyncio
+    async def test_text_latency_not_affected_by_metrics(self):
+        """Test that text_latency is independent of metrics data"""
+        request_input = RequestFuncInput(
+            model="test-model",
+            model_name="test-model",
+            prompt="test prompt",
+            api_url="http://test.com/v1/chat/completions",
+            prompt_len=10,
+            output_len=20,
+        )
+
+        # Create response with text and metrics
+        chunks = [
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": "Response text"}}],
+                    "modality": "text",
+                }
+            ),
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": ""}}],
+                    "modality": "text",
+                    "metrics": {"num_tokens_out": 100, "num_tokens_in": 20},
+                }
+            ),
+            b"data: [DONE]\n\n",
+        ]
+
+        mock_response = MockResponse(200, chunks)
+        mock_session = AsyncMock()
+        mock_session.post = MagicMock(return_value=mock_response)
+
+        # Act
+        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+        # Assert
+        assert output.success is True
+        assert hasattr(output, "text_latency"), "text_latency should exist"
+        assert output.text_latency > 0, "text_latency should be set for text response"
+        # Verify metrics are also present
+        assert output.output_tokens == 100, "metrics should not affect text_latency"
+
+    @pytest.mark.asyncio
+    async def test_text_latency_mixed_modalities(self):
+        """Test text_latency with mixed text and audio modalities"""
+        request_input = RequestFuncInput(
+            model="test-model",
+            model_name="test-model",
+            prompt="test prompt",
+            api_url="http://test.com/v1/chat/completions",
+            prompt_len=10,
+            output_len=20,
+        )
+
+        # Create response with both text and audio
+        chunks = [
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": "Text"}}],
+                    "modality": "text",
+                }
+            ),
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": ""}}],
+                    "modality": "audio",
+                }
+            ),
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": " more text"}}],
+                    "modality": "text",
+                }
+            ),
+            b"data: [DONE]\n\n",
+        ]
+
+        mock_response = MockResponse(200, chunks)
+        mock_session = AsyncMock()
+        mock_session.post = MagicMock(return_value=mock_response)
+
+        # Act
+        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+        # Assert
+        assert output.success is True
+        assert hasattr(output, "text_latency"), "text_latency should exist with mixed modalities"
+        assert output.text_latency > 0, "text_latency should be set when text is present"
+        assert output.generated_text == "Text more text"
+
+    @pytest.mark.asyncio
+    async def test_text_latency_value_consistency(self):
+        """Test that text_latency matches latency minus ttft relationship"""
+        request_input = RequestFuncInput(
+            model="test-model",
+            model_name="test-model",
+            prompt="test prompt",
+            api_url="http://test.com/v1/chat/completions",
+            prompt_len=10,
+            output_len=20,
+        )
+
+        chunks = [
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": "Hello"}}],
+                    "modality": "text",
+                }
+            ),
+            create_sse_chunk(
+                {
+                    "choices": [{"delta": {"content": " world"}}],
+                    "modality": "text",
+                }
+            ),
+            b"data: [DONE]\n\n",
+        ]
+
+        mock_response = MockResponse(200, chunks)
+        mock_session = AsyncMock()
+        mock_session.post = MagicMock(return_value=mock_response)
+
+        # Act
+        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
+
+        # Assert
+        assert output.success is True
+        assert hasattr(output, "text_latency"), "text_latency should exist"
+        assert hasattr(output, "ttft"), "ttft should exist"
+        assert hasattr(output, "latency"), "latency should exist"
+        # text_latency should be between ttft and total latency
+        assert output.ttft <= output.text_latency <= output.latency, (
+            "text_latency should be between ttft and total latency"
+        )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "-s"])

tests/benchmarks/test_serve_cli.py

@@ -4,6 +4,7 @@ from pathlib import Path
 import pytest
 
 from tests.conftest import OmniServer
+from tests.utils import hardware_test
 
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
 stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen3_omni_ci.yaml")]
@@ -29,6 +30,9 @@ def omni_server(request):
         print("OmniServer stopped")
 
 
+@pytest.mark.core_model
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_bench_serve_chat(omni_server):
     command = [

tests/conftest.py

@@ -10,11 +10,14 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 if "VLLM_TARGET_DEVICE" not in os.environ:
     os.environ["VLLM_TARGET_DEVICE"] = "cpu"
 
+import concurrent.futures
 import gc
 import socket
 import subprocess
 import sys
+import threading
 import time
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
@@ -23,12 +26,22 @@ import psutil
 import pytest
 import torch
 import yaml
+from openai import OpenAI
+from vllm import TextPrompt
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 from vllm.logger import init_logger
 from vllm.utils.network_utils import get_open_port
 
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+
 logger = init_logger(__name__)
 
+PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
+PromptImageInput = list[Any] | Any | None
+PromptVideoInput = list[Any] | Any | None
+
 
 @pytest.fixture(autouse=True)
 def default_vllm_config():
@@ -292,6 +305,11 @@ def generate_synthetic_audio(
     if max_amp > 0:
         audio_data = audio_data / max_amp * 0.8
 
+    audio_array = audio_data.copy()
+    result = {
+        "np_array": audio_array,
+    }
+
     # Handle file saving
     audio_bytes = None
 
@@ -318,16 +336,14 @@ def generate_synthetic_audio(
 
     # Return result
     base64_audio = base64.b64encode(audio_bytes).decode("utf-8")
-    result = {
-        "base64": base64_audio,
-    }
+    result["base64"] = base64_audio
     if save_to_file and output_path:
         result["file_path"] = output_path
 
     return result
 
 
-def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_file: bool = False) -> str:
+def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_file: bool = False) -> dict[str, Any]:
     """Generate synthetic video with bouncing balls and return base64 string."""
 
     import cv2
@@ -397,6 +413,10 @@ def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_f
         frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
         video_frames.append(frame_rgb)
 
+    video_array = np.array(video_frames)
+    result = {
+        "np_array": video_array,
+    }
     video_bytes = None
     saved_file_path = None
 
@@ -451,16 +471,14 @@ def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_f
 
     base64_video = base64.b64encode(video_bytes).decode("utf-8")
 
-    result = {
-        "base64": base64_video,
-    }
+    result["base64"] = base64_video
     if save_to_file and saved_file_path:
         result["file_path"] = saved_file_path
 
     return result
 
 
-def generate_synthetic_image(width: int, height: int, save_to_file: bool = False) -> Any:
+def generate_synthetic_image(width: int, height: int, save_to_file: bool = False) -> dict[str, Any]:
     """Generate synthetic image with randomly colored squares and return base64 string."""
     from PIL import Image, ImageDraw
 
@@ -488,6 +506,9 @@ def generate_synthetic_image(width: int, height: int, save_to_file: bool = False
         # Draw square
         draw.rectangle([x, y, x + square_size, y + square_size], fill=color, outline=(0, 0, 0), width=border_width)
 
+    image_array = np.array(image)
+    result = {"np_array": image_array.copy()}
+
     # Handle file saving
     image_bytes = None
     saved_file_path = None
@@ -521,9 +542,7 @@ def generate_synthetic_image(width: int, height: int, save_to_file: bool = False
     base64_image = base64.b64encode(image_bytes).decode("utf-8")
 
     # Return result
-    result = {
-        "base64": base64_image,
-    }
+    result["base64"] = base64_image
     if save_to_file and saved_file_path:
         result["file_path"] = saved_file_path
 
@@ -588,22 +607,30 @@ def convert_audio_to_text(audio_data):
     """
     Convert base64 encoded audio data to text using speech recognition.
     """
-    import whisper
-
     audio_data = base64.b64decode(audio_data)
     output_path = f"./test_{int(time.time())}"
     with open(output_path, "wb") as audio_file:
         audio_file.write(audio_data)
 
     print(f"audio data is saved: {output_path}")
+    text = convert_audio_file_to_text(output_path=output_path)
+    return text
 
-    model = whisper.load_model("base")
+
+def convert_audio_file_to_text(output_path):
+    import whisper
+
+    model = whisper.load_model("small")
     text = model.transcribe(
         output_path,
         temperature=0.0,
         word_timestamps=True,
         condition_on_previous_text=False,
     )["text"]
+    del model
+    if torch.cuda.is_available():
+        gc.collect()
+        torch.cuda.empty_cache()
     if text:
         return text
     else:
@@ -614,7 +641,6 @@ def merge_base64_and_convert_to_text(base64_list):
     """
     Merge a list of base64 encoded audio data and convert to text.
     """
-    import whisper
     from pydub import AudioSegment
 
     merged_audio = None
@@ -627,22 +653,13 @@ def merge_base64_and_convert_to_text(base64_list):
             merged_audio += seg
     output_path = f"./test_{int(time.time())}"
     merged_audio.export(output_path, format="wav")
-    model = whisper.load_model("base")
-    text = model.transcribe(
-        output_path,
-        temperature=0.0,
-        word_timestamps=True,
-        condition_on_previous_text=False,
-    )["text"]
-    if text:
-        return text
-    else:
-        return ""
+    text = convert_audio_file_to_text(output_path)
+    return text
 
 
 def modify_stage_config(
     yaml_path: str,
-    updates: dict[str, Any],
+    updates: dict[str, Any] = None,
     deletes: dict[str, Any] = None,
 ) -> str:
     """
@@ -793,7 +810,7 @@ def modify_stage_config(
         elif isinstance(current, dict) and last_key in current:
             del current[last_key]
         else:
-            raise KeyError(f"Path {path} does not exist")
+            print(f"Path {path} does not exist")
 
     # Apply deletions first
     if deletes:
@@ -811,7 +828,7 @@ def modify_stage_config(
                         # Find stage by ID
                         target_stage = None
                         for stage in stage_args:
-                            if stage.get("stage_id") == stage_id:
+                            if stage.get("stage_id") == int(stage_id):
                                 target_stage = stage
                                 break
 
@@ -842,7 +859,7 @@ def modify_stage_config(
                     # Find stage by ID
                     target_stage = None
                     for stage in stage_args:
-                        if stage.get("stage_id") == stage_id:
+                        if stage.get("stage_id") == int(stage_id):
                             target_stage = stage
                             break
 
@@ -886,6 +903,7 @@ class OmniServer:
         model: str,
         serve_args: list[str],
         *,
+        port: int | None = None,
         env_dict: dict[str, str] | None = None,
     ) -> None:
         _run_pre_test_cleanup(enable_force=True)
@@ -896,7 +914,10 @@ class OmniServer:
         self.env_dict = env_dict
         self.proc: subprocess.Popen | None = None
         self.host = "127.0.0.1"
-        self.port = get_open_port()
+        if port is None:
+            self.port = get_open_port()
+        else:
+            self.port = port
 
     def _start_server(self) -> None:
         """Start the vLLM-Omni server subprocess."""
@@ -1009,3 +1030,641 @@ class OmniServer:
         _run_pre_test_cleanup(enable_force=True)
         _run_post_test_cleanup(enable_force=True)
         cleanup_dist_env_and_memory()
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--run-level",
+        action="store",
+        default="core_model",
+        choices=["core_model", "advanced_model"],
+        help="Test level to run: L2, L3",
+    )
+
+
+@pytest.fixture(scope="session")
+def run_level(request):
+    return request.config.getoption("--run-level")
+
+
+_omni_server_lock = threading.Lock()
+
+
+@pytest.fixture(scope="module")
+def omni_server(request, run_level):
+    """Start vLLM-Omni server as a subprocess with actual model weights.
+    Uses session scope so the server starts only once for the entire test session.
+    Multi-stage initialization can take 10-20+ minutes.
+    """
+    with _omni_server_lock:
+        model, stage_config_path = request.param
+        if run_level == "advanced_model":
+            stage_config_path = modify_stage_config(
+                stage_config_path,
+                deletes={
+                    "stage_args": {
+                        0: ["engine_args.load_format"],
+                        1: ["engine_args.load_format"],
+                        2: ["engine_args.load_format"],
+                    }
+                },
+            )
+
+        with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
+            print("OmniServer started successfully")
+            yield server
+            print("OmniServer stopping...")
+
+        print("OmniServer stopped")
+
+
+@dataclass
+class OmniResponse:
+    text_content: str | None = None
+    audio_data: list[str] | None = None
+    audio_content: str | None = None
+    similarity: float | None = None
+    e2e_latency: float | None = None
+    success: bool = False
+    error_message: str | None = None
+
+
+def assert_omni_response(response: OmniResponse, request_config: dict[str, Any], run_level):
+    """
+    Validate response results.
+
+    Args:
+        response: OmniResponse object
+
+    Raises:
+        AssertionError: When the response does not meet validation criteria
+    """
+    assert response.success, "The request failed."
+    e2e_latency = response.e2e_latency
+    if e2e_latency is not None:
+        print(f"the avg e2e is: {e2e_latency}")
+
+    modalities = request_config.get("modalities", ["text", "audio"])
+
+    if "audio" in modalities:
+        assert response.audio_content is not None, "No audio output is generated"
+        print(f"audio content is: {response.audio_content}")
+
+    if "text" in modalities:
+        assert response.text_content is not None, "No text output is generated"
+        print(f"text content is: {response.text_content}")
+
+    if run_level == "advanced_model":
+        # Verify image description
+        word_types = ["text", "image", "audio", "video"]
+        keywords_dict = request_config.get("key_words", {})
+        for word_type in word_types:
+            keywords = keywords_dict.get(word_type)
+            if "text" in modalities:
+                if keywords:
+                    assert any(keyword in response.text_content.lower() for keyword in keywords), (
+                        "The output does not contain any of the keywords."
+                    )
+            else:
+                if keywords:
+                    assert any(keyword in response.audio_content.lower() for keyword in keywords), (
+                        "The output does not contain any of the keywords."
+                    )
+
+        # Verify similarity
+        if "text" in modalities and "audio" in modalities:
+            assert response.similarity > 0.9, "The audio content is not same as the text"
+            print(f"similarity is: {response.similarity}")
+
+
+class OpenAIClientHandler:
+    """
+    OpenAI client handler class, encapsulating both streaming and non-streaming response processing logic.
+
+    This class integrates OpenAI API request sending, response handling, and validation functionality,
+    supporting both single request and concurrent request modes.
+    """
+
+    def __init__(
+        self, host: str = "127.0.0.1", port: int = get_open_port(), api_key: str = "EMPTY", run_level: str = None
+    ):
+        """
+        Initialize the OpenAI client.
+
+        Args:
+            host: vLLM-Omni server host address
+            port: vLLM-Omni server port
+            api_key: API key (defaults to "EMPTY")
+        """
+        self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key)
+        self.run_level = run_level
+
+    def _process_stream_response(self, chat_completion) -> OmniResponse:
+        """
+        Process streaming responses.
+
+        Args:
+            chat_completion: OpenAI streaming response object
+            request_config: Request configuration dictionary
+
+        Returns:
+            OmniResponse: Processed response object
+        """
+        result = OmniResponse()
+        start_time = time.perf_counter()
+
+        try:
+            text_content = ""
+            audio_data = []
+
+            for chunk in chat_completion:
+                for choice in chunk.choices:
+                    # Get content data
+                    if hasattr(choice, "delta"):
+                        content = getattr(choice.delta, "content", None)
+                    else:
+                        content = None
+
+                    # Get modality type
+                    modality = getattr(chunk, "modality", None)
+
+                    # Process content based on modality type
+                    if modality == "audio" and content:
+                        audio_data.append(content)
+                    elif modality == "text" and content:
+                        text_content += content if content else ""
+
+            # Calculate end-to-end latency
+            result.e2e_latency = time.perf_counter() - start_time
+
+            # Process audio and text content
+            audio_content = None
+            similarity = None
+
+            if audio_data or text_content:
+                if audio_data:
+                    audio_content = merge_base64_and_convert_to_text(audio_data)
+                if audio_content and text_content:
+                    similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
+
+            # Populate result object
+            result.text_content = text_content
+            result.audio_data = audio_data
+            result.audio_content = audio_content
+            result.similarity = similarity
+            result.success = True
+
+        except Exception as e:
+            result.error_message = f"Stream processing error: {str(e)}"
+            print(f"Error: {result.error_message}")
+
+        return result
+
+    def _process_non_stream_response(self, chat_completion) -> OmniResponse:
+        """
+        Process non-streaming responses.
+
+        Args:
+            chat_completion: OpenAI non-streaming response object
+            request_config: Request configuration dictionary
+
+        Returns:
+            OmniResponse: Processed response object
+        """
+        result = OmniResponse()
+        start_time = time.perf_counter()
+
+        try:
+            audio_data = None
+            text_content = None
+
+            # Iterate through all choices
+            for choice in chat_completion.choices:
+                # Process audio data
+                if hasattr(choice.message, "audio") and choice.message.audio is not None:
+                    audio_message = choice.message
+                    audio_data = audio_message.audio.data
+
+                # Process text content
+                if hasattr(choice.message, "content") and choice.message.content is not None:
+                    text_content = choice.message.content
+
+            # Calculate end-to-end latency
+            result.e2e_latency = time.perf_counter() - start_time
+
+            # Process audio and text content
+            audio_content = None
+            similarity = None
+
+            if audio_data or text_content:
+                if audio_data:
+                    audio_content = convert_audio_to_text(audio_data)
+                if audio_content and text_content:
+                    similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
+
+            # Populate result object
+            result.text_content = text_content
+            result.audio_content = audio_content
+            result.similarity = similarity
+            result.success = True
+
+        except Exception as e:
+            result.error_message = f"Non-stream processing error: {str(e)}"
+            print(f"Error: {result.error_message}")
+
+        return result
+
+    def send_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]:
+        """
+        Send OpenAI requests.
+
+        Args:
+            request_config: Request configuration dictionary containing parameters like model, messages, stream
+            request_num: Number of requests, defaults to 1 (single request)
+
+        Returns:
+            List[OmniResponse]: List of response objects
+        """
+
+        responses = []
+        stream = request_config.get("stream", False)
+
+        if request_num == 1:
+            # Send single request
+            chat_completion = self.client.chat.completions.create(
+                model=request_config.get("model"), messages=request_config.get("messages"), stream=stream
+            )
+
+            if stream:
+                response = self._process_stream_response(chat_completion)
+            else:
+                response = self._process_non_stream_response(chat_completion)
+
+            assert_omni_response(response, request_config, run_level=self.run_level)
+            responses.append(response)
+
+        else:
+            # Send concurrent requests
+            with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor:
+                futures = []
+
+                # Submit all request tasks
+                for _ in range(request_num):
+                    future = executor.submit(
+                        self.client.chat.completions.create,
+                        model=request_config.get("model"),
+                        messages=request_config.get("messages"),
+                        stream=stream,
+                    )
+                    futures.append(future)
+
+                # Process completed tasks
+                for future in concurrent.futures.as_completed(futures):
+                    chat_completion = future.result()
+
+                    if stream:
+                        response = self._process_stream_response(chat_completion)
+                    else:
+                        response = self._process_non_stream_response(chat_completion)
+
+                    assert_omni_response(response, request_config, run_level=self.run_level)
+                    responses.append(response)
+
+        return responses
+
+
+@pytest.fixture
+def openai_client(omni_server, run_level):
+    """Create OpenAIClientHandler fixture"""
+    return OpenAIClientHandler(host=omni_server.host, port=omni_server.port, api_key="EMPTY", run_level=run_level)
+
+
+class OmniRunner:
+    """
+    Test runner for Omni models.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        seed: int = 42,
+        stage_init_timeout: int = 300,
+        batch_timeout: int = 10,
+        init_timeout: int = 300,
+        shm_threshold_bytes: int = 65536,
+        log_stats: bool = False,
+        stage_configs_path: str | None = None,
+        **kwargs,
+    ) -> None:
+        """
+        Initialize an OmniRunner for testing.
+
+        Args:
+            model_name: The model name or path
+            seed: Random seed for reproducibility
+            stage_init_timeout: Timeout for initializing a single stage in seconds
+            batch_timeout: Timeout for batching in seconds
+            init_timeout: Timeout for initializing stages in seconds
+            shm_threshold_bytes: Threshold for using shared memory
+            log_stats: Enable detailed statistics logging
+            stage_configs_path: Optional path to YAML stage config file
+            **kwargs: Additional arguments passed to Omni
+        """
+        cleanup_dist_env_and_memory()
+        _run_pre_test_cleanup(enable_force=True)
+        _run_post_test_cleanup(enable_force=True)
+        self.model_name = model_name
+        self.seed = seed
+
+        self.omni = Omni(
+            model=model_name,
+            log_stats=log_stats,
+            stage_init_timeout=stage_init_timeout,
+            batch_timeout=batch_timeout,
+            init_timeout=init_timeout,
+            shm_threshold_bytes=shm_threshold_bytes,
+            stage_configs_path=stage_configs_path,
+            **kwargs,
+        )
+
+    def get_default_sampling_params_list(self) -> list[OmniSamplingParams]:
+        """
+        Get a list of default sampling parameters for all stages.
+
+        Returns:
+            List of SamplingParams with default decoding for each stage
+        """
+        return [st.default_sampling_params for st in self.omni.stage_list]
+
+    def get_omni_inputs(
+        self,
+        prompts: list[str] | str,
+        system_prompt: str | None = None,
+        audios: PromptAudioInput = None,
+        images: PromptImageInput = None,
+        videos: PromptVideoInput = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+        modalities: list[str] | None = None,
+    ) -> list[TextPrompt]:
+        """
+        Construct Omni input format from prompts and multimodal data.
+
+        Args:
+            prompts: Text prompt(s) - either a single string or list of strings
+            system_prompt: Optional system prompt (defaults to Qwen system prompt)
+            audios: Audio input(s) - tuple of (audio_array, sample_rate) or list of tuples
+            images: Image input(s) - PIL Image or list of PIL Images
+            videos: Video input(s) - numpy array or list of numpy arrays
+            mm_processor_kwargs: Optional processor kwargs (e.g., use_audio_in_video)
+
+        Returns:
+            List of prompt dictionaries suitable for Omni.generate()
+        """
+        if system_prompt is None:
+            system_prompt = (
+                "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."
+            )
+
+        video_padding_token = "<|VIDEO|>"
+        image_padding_token = "<|IMAGE|>"
+        audio_padding_token = "<|AUDIO|>"
+
+        if "Qwen3-Omni-30B-A3B-Instruct" in self.model_name:
+            video_padding_token = "<|video_pad|>"
+            image_padding_token = "<|image_pad|>"
+            audio_padding_token = "<|audio_pad|>"
+
+        if isinstance(prompts, str):
+            prompts = [prompts]
+
+        def _normalize_mm_input(mm_input, num_prompts):
+            if mm_input is None:
+                return [None] * num_prompts
+            if isinstance(mm_input, list):
+                if len(mm_input) != num_prompts:
+                    raise ValueError(
+                        f"Multimodal input list length ({len(mm_input)}) must match prompts length ({num_prompts})"
+                    )
+                return mm_input
+            return [mm_input] * num_prompts
+
+        num_prompts = len(prompts)
+        audios_list = _normalize_mm_input(audios, num_prompts)
+        images_list = _normalize_mm_input(images, num_prompts)
+        videos_list = _normalize_mm_input(videos, num_prompts)
+
+        omni_inputs = []
+        for i, prompt_text in enumerate(prompts):
+            user_content = ""
+            multi_modal_data = {}
+
+            audio = audios_list[i]
+            if audio is not None:
+                if isinstance(audio, list):
+                    for _ in audio:
+                        user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
+                    multi_modal_data["audio"] = audio
+                else:
+                    user_content += f"<|audio_bos|>{audio_padding_token}<|audio_eos|>"
+                    multi_modal_data["audio"] = audio
+
+            image = images_list[i]
+            if image is not None:
+                if isinstance(image, list):
+                    for _ in image:
+                        user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
+                    multi_modal_data["image"] = image
+                else:
+                    user_content += f"<|vision_bos|>{image_padding_token}<|vision_eos|>"
+                    multi_modal_data["image"] = image
+
+            video = videos_list[i]
+            if video is not None:
+                if isinstance(video, list):
+                    for _ in video:
+                        user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
+                    multi_modal_data["video"] = video
+                else:
+                    user_content += f"<|vision_bos|>{video_padding_token}<|vision_eos|>"
+                    multi_modal_data["video"] = video
+
+            user_content += prompt_text
+
+            full_prompt = (
+                f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
+                f"<|im_start|>user\n{user_content}<|im_end|>\n"
+                f"<|im_start|>assistant\n"
+            )
+
+            input_dict: TextPrompt = {"prompt": full_prompt}
+            if multi_modal_data:
+                input_dict["multi_modal_data"] = multi_modal_data
+            if modalities:
+                input_dict["modalities"] = modalities
+            if mm_processor_kwargs:
+                input_dict["mm_processor_kwargs"] = mm_processor_kwargs
+
+            omni_inputs.append(input_dict)
+
+        return omni_inputs
+
+    def generate(
+        self,
+        prompts: list[TextPrompt],
+        sampling_params_list: list[OmniSamplingParams] | None = None,
+    ) -> list[OmniRequestOutput]:
+        """
+        Generate outputs for the given prompts.
+
+        Args:
+            prompts: List of prompt dictionaries with 'prompt' and optionally
+                    'multi_modal_data' keys
+            sampling_params_list: List of sampling parameters for each stage.
+                                 If None, uses default parameters.
+
+        Returns:
+            List of OmniRequestOutput objects from stages with final_output=True
+        """
+        if sampling_params_list is None:
+            sampling_params_list = self.get_default_sampling_params_list()
+
+        return self.omni.generate(prompts, sampling_params_list)
+
+    def generate_multimodal(
+        self,
+        prompts: list[str] | str,
+        sampling_params_list: list[OmniSamplingParams] | None = None,
+        system_prompt: str | None = None,
+        audios: PromptAudioInput = None,
+        images: PromptImageInput = None,
+        videos: PromptVideoInput = None,
+        mm_processor_kwargs: dict[str, Any] | None = None,
+        modalities: list[str] | None = None,
+    ) -> list[OmniRequestOutput]:
+        """
+        Convenience method to generate with multimodal inputs.
+
+        Args:
+            prompts: Text prompt(s)
+            sampling_params_list: List of sampling parameters for each stage
+            system_prompt: Optional system prompt
+            audios: Audio input(s)
+            images: Image input(s)
+            videos: Video input(s)
+            mm_processor_kwargs: Optional processor kwargs
+
+        Returns:
+            List of OmniRequestOutput objects from stages with final_output=True
+        """
+        omni_inputs = self.get_omni_inputs(
+            prompts=prompts,
+            system_prompt=system_prompt,
+            audios=audios,
+            images=images,
+            videos=videos,
+            mm_processor_kwargs=mm_processor_kwargs,
+            modalities=modalities,
+        )
+        return self.generate(omni_inputs, sampling_params_list)
+
+    def _cleanup_process(self):
+        try:
+            keywords = ["enginecore"]
+
+            for proc in psutil.process_iter(["pid", "name", "cmdline", "username"]):
+                try:
+                    cmdline = " ".join(proc.cmdline()).lower() if proc.cmdline() else ""
+                    name = proc.name().lower()
+
+                    is_process = any(keyword in cmdline for keyword in keywords) or any(
+                        keyword in name for keyword in keywords
+                    )
+
+                    if is_process:
+                        print(f"Found vllm process: PID={proc.pid}, cmd={cmdline[:100]}")
+
+                        try:
+                            proc.terminate()
+                            time.sleep(2)
+                        except Exception:
+                            proc.kill()
+
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    pass
+
+        except Exception as e:
+            print(f"Error in psutil vllm cleanup: {e}")
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - cleanup resources."""
+        if hasattr(self.omni, "close"):
+            self.omni.close()
+        self._cleanup_process()
+        _run_pre_test_cleanup(enable_force=True)
+        _run_post_test_cleanup(enable_force=True)
+        cleanup_dist_env_and_memory()
+
+
+@pytest.fixture(scope="module")
+def omni_runner(request):
+    with _omni_server_lock:
+        model, stage_config_path = request.param
+        with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, stage_init_timeout=300) as runner:
+            print("OmniRunner started successfully")
+            yield runner
+            print("OmniRunner stopping...")
+
+        print("OmniRunner stopped")
+
+
+class OmniRunnerHandler:
+    def __init__(self, omni_runner):
+        self.runner = omni_runner
+
+    def _process_output(self, outputs: list[Any]) -> OmniResponse:
+        result = OmniResponse()
+        try:
+            text_content = None
+            audio_content = None
+            for stage_output in outputs:
+                if getattr(stage_output, "final_output_type", None) == "text":
+                    text_content = stage_output.request_output[0].outputs[0].text
+                if getattr(stage_output, "final_output_type", None) == "audio":
+                    audio_content = stage_output.request_output[0].outputs[0].multimodal_output["audio"]
+
+            result.audio_content = audio_content
+            result.text_content = text_content
+            result.success = True
+
+        except Exception as e:
+            result.error_message = f"Output processing error: {str(e)}"
+            result.success = False
+            print(f"Error: {result.error_message}")
+
+        return result
+
+    def send_request(self, request_config: dict[str, Any] | None = None) -> OmniResponse:
+        if request_config is None:
+            request_config = {}
+        prompts = request_config.get("prompts")
+        videos = request_config.get("videos")
+        images = request_config.get("images")
+        audios = request_config.get("audios")
+        modalities = request_config.get("modalities", ["text", "audio"])
+        outputs = self.runner.generate_multimodal(
+            prompts=prompts, videos=videos, images=images, audios=audios, modalities=modalities
+        )
+        response = self._process_output(outputs)
+        assert_omni_response(response, request_config, run_level="L2")
+        return response
+
+
+@pytest.fixture
+def omni_runner_handler(omni_runner):
+    return OmniRunnerHandler(omni_runner)

tests/diffusion/cache/test_cache_backends.py

@@ -22,6 +22,8 @@ from vllm_omni.diffusion.cache.selector import get_cache_backend
 from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class TestCacheDiTBackend:
     """Test CacheDiTBackend implementation."""

tests/diffusion/distributed/test_vae_patch_parallel.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for VAE patch/tile parallelism helpers (CPU-only)."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.distributed import vae_patch_parallel as vae_patch_parallel
+
+
+class _DummyConfig:
+    def __init__(self, **attrs):
+        for k, v in attrs.items():
+            setattr(self, k, v)
+
+
+class _DummyVae:
+    def __init__(self, *, config=None, **attrs):
+        self.config = config
+        for k, v in attrs.items():
+            setattr(self, k, v)
+
+
+def test_get_vae_spatial_scale_factor_uses_block_out_channels_len_minus_1():
+    vae = _DummyVae(config=_DummyConfig(block_out_channels=[128, 256, 512, 512]))
+    assert vae_patch_parallel._get_vae_spatial_scale_factor(vae) == 8
+
+    vae = _DummyVae(config=_DummyConfig(block_out_channels=[1, 2, 3, 4, 5]))
+    assert vae_patch_parallel._get_vae_spatial_scale_factor(vae) == 16
+
+
+def test_get_vae_spatial_scale_factor_defaults_to_8_on_missing_or_empty():
+    assert vae_patch_parallel._get_vae_spatial_scale_factor(_DummyVae(config=_DummyConfig())) == 8
+    assert vae_patch_parallel._get_vae_spatial_scale_factor(_DummyVae(config=_DummyConfig(block_out_channels=[]))) == 8
+    assert vae_patch_parallel._get_vae_spatial_scale_factor(_DummyVae(config=None)) == 8
+
+
+def test_get_vae_spatial_scale_factor_defaults_to_8_on_exception():
+    class _BrokenConfig:
+        @property
+        def block_out_channels(self):
+            raise RuntimeError("boom")
+
+    assert vae_patch_parallel._get_vae_spatial_scale_factor(_DummyVae(config=_BrokenConfig())) == 8
+
+
+@pytest.mark.parametrize(
+    ("pp_size", "expected"),
+    [
+        (0, (1, 1)),
+        (1, (1, 1)),
+        (2, (1, 2)),
+        (3, (1, 3)),
+        (4, (2, 2)),
+        (6, (2, 3)),
+        (8, (2, 4)),
+        (12, (3, 4)),
+        (16, (4, 4)),
+    ],
+)
+def test_factor_pp_grid(pp_size: int, expected: tuple[int, int]):
+    assert vae_patch_parallel._factor_pp_grid(pp_size) == expected
+
+
+def test_get_world_rank_pp_size(monkeypatch):
+    monkeypatch.setattr(vae_patch_parallel.dist, "get_world_size", lambda _: 8)
+    monkeypatch.setattr(vae_patch_parallel.dist, "get_rank", lambda _: 3)
+
+    world_size, rank, pp_size = vae_patch_parallel._get_world_rank_pp_size(object(), 4)
+    assert (world_size, rank, pp_size) == (8, 3, 4)
+
+    world_size, rank, pp_size = vae_patch_parallel._get_world_rank_pp_size(object(), 16)
+    assert (world_size, rank, pp_size) == (8, 3, 8)
+
+
+def test_get_vae_out_channels_defaults_to_3():
+    assert vae_patch_parallel._get_vae_out_channels(_DummyVae(config=None)) == 3
+    assert vae_patch_parallel._get_vae_out_channels(_DummyVae(config=_DummyConfig())) == 3
+
+
+def test_get_vae_out_channels_reads_config():
+    assert vae_patch_parallel._get_vae_out_channels(_DummyVae(config=_DummyConfig(out_channels=4))) == 4
+    assert vae_patch_parallel._get_vae_out_channels(_DummyVae(config=_DummyConfig(out_channels="5"))) == 5
+
+
+def test_get_vae_tile_params_returns_none_if_missing():
+    assert (
+        vae_patch_parallel._get_vae_tile_params(_DummyVae(tile_latent_min_size=None, tile_overlap_factor=0.25)) is None
+    )
+    assert (
+        vae_patch_parallel._get_vae_tile_params(_DummyVae(tile_latent_min_size=128, tile_overlap_factor=None)) is None
+    )
+
+
+def test_get_vae_tile_params_parses_types():
+    vae = _DummyVae(tile_latent_min_size="128", tile_overlap_factor="0.25")
+    assert vae_patch_parallel._get_vae_tile_params(vae) == (128, 0.25)
+
+
+def test_get_vae_tiling_params_returns_none_if_missing():
+    vae = _DummyVae(tile_latent_min_size=128, tile_overlap_factor=0.25, tile_sample_min_size=None)
+    assert vae_patch_parallel._get_vae_tiling_params(vae) is None
+
+    vae = _DummyVae(tile_latent_min_size=None, tile_overlap_factor=0.25, tile_sample_min_size=1024)
+    assert vae_patch_parallel._get_vae_tiling_params(vae) is None
+
+
+def test_get_vae_tiling_params_parses_types():
+    vae = _DummyVae(tile_latent_min_size="128", tile_overlap_factor="0.25", tile_sample_min_size="1024")
+    assert vae_patch_parallel._get_vae_tiling_params(vae) == (128, 0.25, 1024)
+
+
+def test_distributed_tiled_decode_stitches_tiles(monkeypatch):
+    class _TinyConfig:
+        def __init__(self):
+            self.out_channels = 1
+            self.use_post_quant_conv = False
+
+    class _TinyVae:
+        def __init__(self):
+            self.config = _TinyConfig()
+            self.tile_latent_min_size = 2
+            self.tile_overlap_factor = 0.0
+            self.tile_sample_min_size = 2
+
+        def decoder(self, x: torch.Tensor) -> torch.Tensor:
+            return x
+
+        def blend_v(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
+            return b
+
+        def blend_h(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
+            return b
+
+    def _collect_local_tiles(
+        *,
+        vae: _TinyVae,
+        z: torch.Tensor,
+        rank: int,
+        pp_size: int,
+    ) -> tuple[list[torch.Tensor], list[tuple[int, int, int]]]:
+        tile_latent_min_size = vae.tile_latent_min_size
+        overlap_size = int(tile_latent_min_size * (1 - vae.tile_overlap_factor))
+        h_starts = list(range(0, z.shape[2], overlap_size))
+        w_starts = list(range(0, z.shape[3], overlap_size))
+
+        local_tiles: list[torch.Tensor] = []
+        local_meta: list[tuple[int, int, int]] = []
+        tile_id = 0
+        for i in h_starts:
+            for j in w_starts:
+                tile_rank = (tile_id + 1) % pp_size
+                if tile_rank == rank:
+                    tile = z[:, :, i : i + tile_latent_min_size, j : j + tile_latent_min_size]
+                    decoded = vae.decoder(tile)
+                    local_tiles.append(decoded)
+                    local_meta.append((tile_id, int(decoded.shape[-2]), int(decoded.shape[-1])))
+                tile_id += 1
+        return local_tiles, local_meta
+
+    vae = _TinyVae()
+    z = torch.arange(16, dtype=torch.float32).reshape(1, 1, 4, 4)
+
+    rank0_tiles, rank0_meta = _collect_local_tiles(vae=vae, z=z, rank=0, pp_size=2)
+    rank1_tiles, rank1_meta = _collect_local_tiles(vae=vae, z=z, rank=1, pp_size=2)
+    max_count = max(len(rank0_tiles), len(rank1_tiles))
+
+    def _pack_meta_and_tiles(
+        tiles: list[torch.Tensor],
+        meta: list[tuple[int, int, int]],
+        max_count: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        meta_tensor = torch.full((max_count, 3), -1, dtype=torch.int64)
+        tile_tensor = torch.zeros(
+            (max_count, z.shape[0], vae.config.out_channels, vae.tile_sample_min_size, vae.tile_sample_min_size),
+            dtype=z.dtype,
+        )
+        for idx, (tile_id, h, w) in enumerate(meta):
+            meta_tensor[idx, 0] = tile_id
+            meta_tensor[idx, 1] = h
+            meta_tensor[idx, 2] = w
+            tile_tensor[idx, :, :, :h, :w] = tiles[idx]
+        return meta_tensor, tile_tensor
+
+    rank1_meta_tensor, rank1_tile_tensor = _pack_meta_and_tiles(rank1_tiles, rank1_meta, max_count)
+    rank1_count_tensor = torch.tensor([len(rank1_tiles)], dtype=torch.int64)
+
+    def _fake_gather(tensor, gather_list=None, dst=0, group=None):
+        if gather_list is None:
+            return
+        if tensor.ndim == 1 and tensor.numel() == 1:
+            gather_list[0].copy_(tensor)
+            gather_list[1].copy_(rank1_count_tensor)
+            return
+        if tensor.ndim == 2 and tensor.shape[1] == 3:
+            gather_list[0].copy_(tensor)
+            gather_list[1].copy_(rank1_meta_tensor)
+            return
+        if tensor.ndim == 5:
+            gather_list[0].copy_(tensor)
+            gather_list[1].copy_(rank1_tile_tensor)
+            return
+        raise AssertionError("Unexpected gather payload for test.")
+
+    def _fake_broadcast(_tensor, src=0, group=None):
+        return
+
+    monkeypatch.setattr(vae_patch_parallel.dist, "get_world_size", lambda _group: 2)
+    monkeypatch.setattr(vae_patch_parallel.dist, "get_rank", lambda _group: 0)
+    monkeypatch.setattr(vae_patch_parallel.dist, "gather", _fake_gather)
+    monkeypatch.setattr(vae_patch_parallel.dist, "broadcast", _fake_broadcast)
+
+    output = vae_patch_parallel._distributed_tiled_decode(
+        vae=vae,
+        orig_decode=lambda z, return_dict=False: (z,),
+        z=z,
+        group=object(),
+        vae_patch_parallel_size=2,
+    )
+
+    assert torch.equal(output, z)

tests/diffusion/lora/test_base_linear.py

@@ -5,10 +5,13 @@ from __future__ import annotations
 
 from dataclasses import dataclass
 
+import pytest
 import torch
 
 from vllm_omni.diffusion.lora.layers.base_linear import DiffusionBaseLinearLayerWithLoRA
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 @dataclass
 class _DummyLoRAConfig:

tests/diffusion/lora/test_lora_manager.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import pytest
 import torch
 from vllm.lora.lora_weights import LoRALayerWeights
 from vllm.lora.utils import get_supported_lora_modules
@@ -11,6 +12,8 @@ from vllm.model_executor.layers.linear import LinearBase
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.lora.request import LoRARequest
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class _DummyLoRALayer:
     def __init__(self, n_slices: int, output_slices: tuple[int, ...]):

tests/diffusion/quantization/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/diffusion/quantization/test_fp8_config.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for FP8 quantization config."""
+
+import pytest
+
+
+def test_fp8_config_creation():
+    """Test that FP8 config can be created."""
+    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+
+    config = get_diffusion_quant_config("fp8")
+    assert config is not None
+    assert config.get_name() == "fp8"
+
+
+def test_vllm_config_extraction():
+    """Test that vLLM config can be extracted from diffusion config."""
+    from vllm_omni.diffusion.quantization import (
+        get_diffusion_quant_config,
+        get_vllm_quant_config_for_layers,
+    )
+
+    diff_config = get_diffusion_quant_config("fp8")
+    vllm_config = get_vllm_quant_config_for_layers(diff_config)
+    assert vllm_config is not None
+    assert vllm_config.activation_scheme == "dynamic"
+
+
+def test_none_quantization():
+    """Test that None quantization returns None config."""
+    from vllm_omni.diffusion.quantization import (
+        get_diffusion_quant_config,
+        get_vllm_quant_config_for_layers,
+    )
+
+    config = get_diffusion_quant_config(None)
+    assert config is None
+    vllm_config = get_vllm_quant_config_for_layers(config)
+    assert vllm_config is None
+
+
+def test_invalid_quantization():
+    """Test that invalid quantization method raises error."""
+    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+
+    with pytest.raises(ValueError, match="Unknown quantization method"):
+        get_diffusion_quant_config("invalid_method")
+
+
+def test_fp8_config_with_custom_params():
+    """Test FP8 config with custom parameters."""
+    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+
+    config = get_diffusion_quant_config(
+        "fp8",
+        activation_scheme="static",
+        ignored_layers=["proj_out"],
+    )
+    assert config is not None
+    assert config.activation_scheme == "static"
+    assert "proj_out" in config.ignored_layers
+
+
+def test_supported_methods():
+    """Test that supported methods list is correct."""
+    from vllm_omni.diffusion.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    assert "fp8" in SUPPORTED_QUANTIZATION_METHODS
+
+
+def test_quantization_integration():
+    """Test end-to-end quantization flow through OmniDiffusionConfig."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    # Test with quantization string only
+    config = OmniDiffusionConfig(model="test", quantization="fp8")
+    assert config.quantization_config is not None
+    assert config.quantization_config.get_name() == "fp8"
+
+    # Test with quantization_config dict
+    config2 = OmniDiffusionConfig(
+        model="test",
+        quantization_config={"method": "fp8", "activation_scheme": "static"},
+    )
+    assert config2.quantization_config is not None
+    assert config2.quantization_config.get_name() == "fp8"
+    assert config2.quantization_config.activation_scheme == "static"
+
+    # Test that vLLM config can be extracted
+    vllm_config = config.quantization_config.get_vllm_quant_config()
+    assert vllm_config is not None
+
+
+def test_quantization_dict_not_mutated():
+    """Test that passing a dict to quantization_config doesn't mutate it."""
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    original_dict = {"method": "fp8", "activation_scheme": "static"}
+    dict_copy = original_dict.copy()
+
+    OmniDiffusionConfig(model="test", quantization_config=original_dict)
+
+    # Original dict should be unchanged
+    assert original_dict == dict_copy
+
+
+def test_quantization_conflicting_methods_warning(caplog):
+    """Test warning when quantization and quantization_config['method'] conflict."""
+    import logging
+
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    with caplog.at_level(logging.WARNING):
+        config = OmniDiffusionConfig(
+            model="test",
+            quantization="fp8",  # This should be overridden
+            quantization_config={"method": "fp8", "activation_scheme": "static"},
+        )
+    # No warning when methods match
+    assert config.quantization_config is not None
+
+
+def test_fp8_delegates_to_vllm_config():
+    """Test that DiffusionFp8Config delegates to vLLM's Fp8Config."""
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    from vllm_omni.diffusion.quantization import DiffusionFp8Config
+
+    # Test that quant_config_cls is set correctly
+    assert DiffusionFp8Config.quant_config_cls is Fp8Config
+
+    # Test that get_name() delegates to vLLM
+    assert DiffusionFp8Config.get_name() == Fp8Config.get_name()
+
+    # Test that get_min_capability() delegates to vLLM
+    assert DiffusionFp8Config.get_min_capability() == Fp8Config.get_min_capability()

tests/diffusion/test_diffusion_worker.py

@@ -17,6 +17,8 @@ import torch
 
 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 @pytest.fixture
 def mock_od_config():
@@ -34,7 +36,7 @@ def mock_od_config():
 @pytest.fixture
 def mock_gpu_worker(mock_od_config):
     """Create a DiffusionWorker with mocked initialization."""
-    with patch.object(DiffusionWorker, "init_device"):
+    with patch.object(DiffusionWorker, "init_device"), patch.object(DiffusionWorker, "load_model"):
         worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
         # Mock the model_runner with pipeline
         worker.model_runner = Mock()
@@ -79,18 +81,18 @@ class TestDiffusionWorkerLoadWeights:
 class TestDiffusionWorkerSleep:
     """Test DiffusionWorker.sleep method."""
 
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
     @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_gpu_worker):
+    def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker):
         """Test sleep mode level 1 (offload weights only)."""
-        # Setup memory info mocks
-        # Before sleep: 1GB free
-        # After sleep: 3GB free (freed 2GB)
-        mock_platform.get_free_memory.side_effect = [
-            1 * 1024**3,  # Before sleep
-            3 * 1024**3,  # After sleep
+        # Setup process-scoped memory mocks
+        # Before sleep: 3GB used
+        # After sleep: 1GB used (freed 2GB)
+        mock_get_process_memory.side_effect = [
+            3 * 1024**3,
+            1 * 1024**3,
         ]
-        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         # Setup allocator mock
         mock_allocator = Mock()
@@ -106,16 +108,16 @@ class TestDiffusionWorkerSleep:
         # Verify buffers were NOT saved (level 1 doesn't save buffers)
         assert len(mock_gpu_worker._sleep_saved_buffers) == 0
 
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
     @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_gpu_worker):
+    def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker):
         """Test sleep mode level 2 (offload all, save buffers)."""
-        # Setup memory info mocks
-        mock_platform.get_free_memory.side_effect = [
-            1 * 1024**3,  # Before sleep
-            5 * 1024**3,  # After sleep (freed 4GB)
+        # Setup process-scoped memory mocks
+        mock_get_process_memory.side_effect = [
+            5 * 1024**3,  # Before sleep
+            1 * 1024**3,  # After sleep (freed 4GB)
         ]
-        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         # Setup allocator mock
         mock_allocator = Mock()
@@ -144,16 +146,22 @@ class TestDiffusionWorkerSleep:
         assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
         assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
 
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
     @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
     @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_sleep_memory_freed_validation(self, mock_allocator_class, mock_platform, mock_gpu_worker):
+    def test_sleep_memory_freed_validation(
+        self,
+        mock_allocator_class,
+        mock_platform,
+        mock_get_process_memory,
+        mock_gpu_worker,
+    ):
         """Test that sleep validates memory was actually freed."""
-        # Simulate memory increase (should trigger assertion error)
-        mock_platform.get_free_memory.side_effect = [
-            3 * 1024**3,  # Before sleep: 3GB free
-            1 * 1024**3,  # After sleep: 1GB free (negative freed!)
+        # Simulate process memory increase (should trigger assertion error)
+        mock_get_process_memory.side_effect = [
+            1 * 1024**3,  # Before sleep: 1GB used
+            3 * 1024**3,  # After sleep: 3GB used (negative freed)
         ]
-        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
 
         mock_allocator = Mock()
         mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
@@ -163,6 +171,33 @@ class TestDiffusionWorkerSleep:
         with pytest.raises(AssertionError, match="Memory usage increased after sleeping"):
             mock_gpu_worker.sleep(level=1)
 
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
+    @patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+    @patch("vllm.device_allocator.cumem.CuMemAllocator")
+    def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(
+        self,
+        mock_allocator_class,
+        mock_platform,
+        mock_get_process_memory,
+        mock_gpu_worker,
+    ):
+        """Test sleep uses device-scoped fallback when NVML is unavailable."""
+        mock_get_process_memory.side_effect = [None, None]
+        mock_platform.get_free_memory.side_effect = [
+            1 * 1024**3,  # Before sleep
+            3 * 1024**3,  # After sleep
+        ]
+        mock_platform.get_device_total_memory.return_value = 8 * 1024**3
+
+        mock_allocator = Mock()
+        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
+        mock_allocator.sleep = Mock()
+
+        result = mock_gpu_worker.sleep(level=1)
+
+        mock_allocator.sleep.assert_called_once_with(offload_tags=("weights",))
+        assert result is True
+
 
 class TestDiffusionWorkerWakeUp:
     """Test DiffusionWorker.wake_up method."""

tests/diffusion/test_worker_wrapper_base.py

@@ -0,0 +1,540 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for WorkerWrapperBase class.
+
+This module tests the WorkerWrapperBase implementation:
+- Initialization with and without worker extensions
+- Custom pipeline initialization
+- Method delegation via execute_method
+- Attribute delegation via __getattr__
+- Dynamic worker class extension
+"""
+
+from typing import Any
+from unittest.mock import Mock, patch
+
+import pytest
+
+from vllm_omni.diffusion.worker.diffusion_worker import (
+    CustomPipelineWorkerExtension,
+    DiffusionWorker,
+    WorkerWrapperBase,
+)
+
+# -------------------------------------------------------------------------
+# Fixtures
+# -------------------------------------------------------------------------
+
+
+@pytest.fixture
+def mock_od_config():
+    """Create a mock OmniDiffusionConfig for use in tests."""
+    config = Mock()
+    config.num_gpus = 1
+    config.master_port = 12345
+    config.enable_sleep_mode = False
+    config.cache_backend = None
+    config.cache_config = None
+    config.model = "test-model"
+    config.diffusion_load_format = None
+    config.dtype = "float32"
+    config.max_cpu_loras = 0
+    config.lora_path = None
+    config.lora_scale = 1.0
+    return config
+
+
+class TestExtension:
+    """Simple test extension adding one custom method."""
+
+    def custom_method(self):
+        return "extension_method"
+
+
+class MockCustomPipeline:
+    """Mock custom pipeline for testing."""
+
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+
+    def __call__(self, *args, **kwargs):
+        return "pipeline_output"
+
+
+# -------------------------------------------------------------------------
+# Tests: Initialization
+# -------------------------------------------------------------------------
+
+
+class TestWorkerWrapperBaseInitialization:
+    """Test WorkerWrapperBase initialization behavior."""
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_basic_initialization(self, mock_worker_init, mock_od_config):
+        """Test basic initialization without extensions."""
+        wrapper = WorkerWrapperBase(
+            gpu_id=0,
+            od_config=mock_od_config,
+            base_worker_class=DiffusionWorker,
+        )
+
+        assert wrapper.gpu_id == 0
+        assert wrapper.od_config == mock_od_config
+        assert wrapper.base_worker_class == DiffusionWorker
+        assert wrapper.worker_extension_cls is None
+        assert wrapper.custom_pipeline_args is None
+        assert wrapper.worker is not None
+
+        mock_worker_init.assert_called_once_with(
+            local_rank=0,
+            rank=0,
+            od_config=mock_od_config,
+        )
+
+
+# -------------------------------------------------------------------------
+# Tests: Worker Extension Functionality
+# -------------------------------------------------------------------------
+
+
+class TestWorkerWrapperBaseExtension:
+    """Test WorkerWrapperBase worker extension functionality."""
+
+    def test_prepare_worker_class_without_extension(self, mock_od_config):
+        """Test _prepare_worker_class without a worker extension."""
+        with patch.object(DiffusionWorker, "__init__", return_value=None):
+            wrapper = WorkerWrapperBase(
+                gpu_id=0,
+                od_config=mock_od_config,
+                base_worker_class=DiffusionWorker,
+            )
+            worker_class = wrapper._prepare_worker_class()
+            assert worker_class == DiffusionWorker
+
+    def test_prepare_worker_class_with_extension_class(self, mock_od_config):
+        """Test _prepare_worker_class with an explicit extension class."""
+
+        class TestExtension:
+            def custom_method(self):
+                return "extension_method"
+
+        with patch.object(DiffusionWorker, "__init__", return_value=None):
+            wrapper = WorkerWrapperBase(
+                gpu_id=0,
+                od_config=mock_od_config,
+                base_worker_class=DiffusionWorker,
+                worker_extension_cls=TestExtension,
+            )
+
+            assert hasattr(wrapper.worker.__class__, "custom_method")
+            assert TestExtension in wrapper.worker.__class__.__bases__
+
+    @patch("vllm.utils.import_utils.resolve_obj_by_qualname")
+    def test_prepare_worker_class_with_extension_string(self, mock_resolve, mock_od_config):
+        """Test _prepare_worker_class with worker extension as string."""
+        mock_resolve.return_value = TestExtension
+
+        with patch.object(DiffusionWorker, "__init__", return_value=None):
+            wrapper = WorkerWrapperBase(
+                gpu_id=0,
+                od_config=mock_od_config,
+                base_worker_class=DiffusionWorker,
+                worker_extension_cls="tests.diffusion.test_worker_wrapper_base.TestExtension",
+            )
+
+            assert hasattr(wrapper.worker.__class__, "custom_method")
+
+
+# -------------------------------------------------------------------------
+# Tests: Method Delegation
+# -------------------------------------------------------------------------
+
+
+class TestWorkerWrapperBaseDelegation:
+    """Test WorkerWrapperBase delegation to wrapped worker."""
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_generate_delegation(self, mock_worker_init, mock_od_config):
+        """Test that generate() delegates to worker.generate()."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        mock_output = Mock()
+        wrapper.worker.generate = Mock(return_value=mock_output)
+
+        mock_requests = [Mock()]
+        result = wrapper.generate(mock_requests)
+
+        wrapper.worker.generate.assert_called_once_with(mock_requests)
+        assert result == mock_output
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_execute_model_delegation(self, mock_worker_init, mock_od_config):
+        """Test that execute_model() delegates to worker.execute_model()."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        mock_output = Mock()
+        wrapper.worker.execute_model = Mock(return_value=mock_output)
+
+        mock_reqs = [Mock()]
+        result = wrapper.execute_model(mock_reqs, mock_od_config)
+
+        wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config)
+        assert result == mock_output
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_load_weights_delegation(self, mock_worker_init, mock_od_config):
+        """Test that load_weights() delegates to worker.load_weights()."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        expected_result = {"weight1", "weight2"}
+        wrapper.worker.load_weights = Mock(return_value=expected_result)
+
+        mock_weights = [("weight1", Mock()), ("weight2", Mock())]
+        result = wrapper.load_weights(mock_weights)
+
+        wrapper.worker.load_weights.assert_called_once_with(mock_weights)
+        assert result == expected_result
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_sleep_delegation(self, mock_worker_init, mock_od_config):
+        """Test that sleep() delegates to worker.sleep()."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        wrapper.worker.sleep = Mock(return_value=True)
+        result = wrapper.sleep(level=1)
+
+        wrapper.worker.sleep.assert_called_once_with(1)
+        assert result is True
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_wake_up_delegation(self, mock_worker_init, mock_od_config):
+        """Test that wake_up() delegates to worker.wake_up()."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        wrapper.worker.wake_up = Mock(return_value=True)
+
+        result = wrapper.wake_up(tags=["weights"])
+        wrapper.worker.wake_up.assert_called_once_with(["weights"])
+        assert result is True
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_shutdown_delegation(self, mock_worker_init, mock_od_config):
+        """Test that shutdown() delegates to worker.shutdown()."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        wrapper.worker.shutdown = Mock(return_value=None)
+
+        result = wrapper.shutdown()
+        wrapper.worker.shutdown.assert_called_once()
+        assert result is None
+
+
+# -------------------------------------------------------------------------
+# Tests: execute_method
+# -------------------------------------------------------------------------
+
+
+class TestWorkerWrapperBaseExecuteMethod:
+    """Test WorkerWrapperBase.execute_method functionality."""
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_execute_method_success(self, mock_worker_init, mock_od_config):
+        """Test execute_method successfully calls worker method."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        wrapper.worker.test_method = Mock(return_value="method_result")
+
+        result = wrapper.execute_method("test_method", "arg1", kwarg1="value1")
+
+        wrapper.worker.test_method.assert_called_once_with("arg1", kwarg1="value1")
+        assert result == "method_result"
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_execute_method_with_no_args(self, mock_worker_init, mock_od_config):
+        """Test execute_method with no arguments."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        wrapper.worker.no_args_method = Mock(return_value="no_args_result")
+
+        result = wrapper.execute_method("no_args_method")
+        wrapper.worker.no_args_method.assert_called_once_with()
+        assert result == "no_args_result"
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_execute_method_error(self, mock_worker_init, mock_od_config):
+        """Test execute_method raises exception on error."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        wrapper.worker.error_method = Mock(side_effect=RuntimeError("Test error"))
+
+        with pytest.raises(RuntimeError, match="Test error"):
+            wrapper.execute_method("error_method")
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_execute_method_invalid_type(self, mock_worker_init, mock_od_config):
+        """Test execute_method with invalid method type."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+
+        with pytest.raises(AssertionError, match="Method must be str"):
+            wrapper.execute_method(b"bytes_method")
+
+
+# -------------------------------------------------------------------------
+# Tests: __getattr__ delegation
+# -------------------------------------------------------------------------
+
+
+class TestWorkerWrapperBaseGetAttr:
+    """Test WorkerWrapperBase.__getattr__ delegation."""
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_getattr_delegation(self, mock_worker_init, mock_od_config):
+        """Test __getattr__ delegates to worker attributes."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        wrapper.worker.custom_attribute = "test_value"
+        assert wrapper.custom_attribute == "test_value"
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_getattr_method_access(self, mock_worker_init, mock_od_config):
+        """Test __getattr__ delegates to worker methods."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        wrapper.worker.custom_method = Mock(return_value="method_result")
+
+        result = wrapper.custom_method()
+        wrapper.worker.custom_method.assert_called_once()
+        assert result == "method_result"
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_getattr_missing_attribute(self, mock_worker_init, mock_od_config):
+        """Test __getattr__ raises AttributeError for missing attributes."""
+        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
+        with pytest.raises(AttributeError):
+            _ = wrapper.nonexistent_attribute
+
+
+# -------------------------------------------------------------------------
+# Tests: Edge Cases
+# -------------------------------------------------------------------------
+
+
+class TestWorkerWrapperBaseEdgeCases:
+    """Test WorkerWrapperBase edge cases and special scenarios."""
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_extension_conflict_warning(self, mock_worker_init, mock_od_config, caplog):
+        """Test a warning is logged when an extension conflicts with worker."""
+
+        class ConflictExtension:
+            def load_model(self):
+                return "extension_load_model"
+
+        with patch.object(DiffusionWorker, "load_model"):
+            wrapper = WorkerWrapperBase(
+                gpu_id=0,
+                od_config=mock_od_config,
+                base_worker_class=DiffusionWorker,
+                worker_extension_cls=ConflictExtension,
+            )
+            assert wrapper.worker is not None
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_multiple_extensions_same_class(self, mock_worker_init, mock_od_config):
+        """Test that applying same extension twice doesn't duplicate it."""
+
+        class TestExtension:
+            def custom_method(self):
+                return "extension"
+
+        with patch.object(DiffusionWorker, "__init__", return_value=None):
+            wrapper1 = WorkerWrapperBase(
+                gpu_id=0,
+                od_config=mock_od_config,
+                base_worker_class=DiffusionWorker,
+                worker_extension_cls=TestExtension,
+            )
+            wrapper2 = WorkerWrapperBase(
+                gpu_id=0,
+                od_config=mock_od_config,
+                base_worker_class=DiffusionWorker,
+                worker_extension_cls=TestExtension,
+            )
+
+            assert hasattr(wrapper1.worker, "custom_method")
+            assert hasattr(wrapper2.worker, "custom_method")
+
+
+# -------------------------------------------------------------------------
+# Tests: CustomPipelineWorkerExtension
+# -------------------------------------------------------------------------
+
+
+class TestCustomPipelineWorkerExtension:
+    """Test CustomPipelineWorkerExtension functionality."""
+
+    @patch("torch.cuda.empty_cache")
+    @patch("gc.collect")
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_re_init_pipeline_basic(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
+        """Test basic re_init_pipeline functionality."""
+        wrapper = WorkerWrapperBase(
+            gpu_id=0,
+            od_config=mock_od_config,
+            base_worker_class=DiffusionWorker,
+            worker_extension_cls=CustomPipelineWorkerExtension,
+        )
+
+        # Setup mock model_runner and pipeline
+        mock_model_runner = Mock()
+        mock_pipeline = Mock()
+        mock_model_runner.pipeline = mock_pipeline
+        wrapper.worker.model_runner = mock_model_runner
+        wrapper.worker.init_lora_manager = Mock()
+        wrapper.worker.load_model = Mock()
+
+        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
+
+        # Call re_init_pipeline
+        wrapper.worker.re_init_pipeline(custom_args)
+
+        # Verify load_model was called with correct arguments
+        wrapper.worker.load_model.assert_called_once_with(
+            load_format="custom_pipeline",
+            custom_pipeline_name="tests.diffusion.test_worker_wrapper_base.MockCustomPipeline",
+        )
+        wrapper.worker.init_lora_manager.assert_called_once()
+
+    @patch("torch.cuda.empty_cache")
+    @patch("gc.collect")
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_re_init_pipeline_cleanup(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
+        """Test that re_init_pipeline properly cleans up old pipeline."""
+        wrapper = WorkerWrapperBase(
+            gpu_id=0,
+            od_config=mock_od_config,
+            base_worker_class=DiffusionWorker,
+            worker_extension_cls=CustomPipelineWorkerExtension,
+        )
+
+        # Setup mock model_runner with pipeline
+        mock_model_runner = Mock()
+        mock_pipeline = Mock()
+        mock_model_runner.pipeline = mock_pipeline
+        wrapper.worker.model_runner = mock_model_runner
+        wrapper.worker.init_lora_manager = Mock()
+        wrapper.worker.load_model = Mock()
+
+        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
+
+        # Call re_init_pipeline
+        wrapper.worker.re_init_pipeline(custom_args)
+
+        # Verify cleanup was performed
+        mock_gc_collect.assert_called_once()
+        mock_empty_cache.assert_called_once()
+
+    @patch("torch.cuda.empty_cache")
+    @patch("gc.collect")
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_re_init_pipeline_none_pipeline(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
+        """Test re_init_pipeline when pipeline is None."""
+        wrapper = WorkerWrapperBase(
+            gpu_id=0,
+            od_config=mock_od_config,
+            base_worker_class=DiffusionWorker,
+            worker_extension_cls=CustomPipelineWorkerExtension,
+        )
+
+        # Setup mock model_runner with None pipeline
+        mock_model_runner = Mock()
+        mock_model_runner.pipeline = None
+        wrapper.worker.model_runner = mock_model_runner
+        wrapper.worker.init_lora_manager = Mock()
+        wrapper.worker.load_model = Mock()
+
+        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
+
+        # Should not raise an error
+        wrapper.worker.re_init_pipeline(custom_args)
+
+        # Verify load_model was still called
+        wrapper.worker.load_model.assert_called_once()
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_custom_pipeline_args_initialization(self, mock_worker_init, mock_od_config):
+        """Test initialization with custom_pipeline_args calls re_init_pipeline."""
+        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
+
+        with patch.object(DiffusionWorker, "__init__", return_value=None):
+            with patch.object(WorkerWrapperBase, "_prepare_worker_class") as mock_prepare:
+                # Create a mock worker class with re_init_pipeline
+                mock_worker_class = Mock()
+                mock_worker_instance = Mock()
+                mock_worker_instance.re_init_pipeline = Mock()
+                mock_worker_class.return_value = mock_worker_instance
+                mock_prepare.return_value = mock_worker_class
+
+                _ = WorkerWrapperBase(
+                    gpu_id=0,
+                    od_config=mock_od_config,
+                    base_worker_class=DiffusionWorker,
+                    custom_pipeline_args=custom_args,
+                )
+
+                # Verify re_init_pipeline was called with custom_pipeline_args
+                mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args)
+
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_custom_pipeline_with_explicit_extension(self, mock_worker_init, mock_od_config):
+        """Test that explicit worker_extension_cls is preserved when custom_pipeline_args is provided."""
+
+        class CustomExtension:
+            def re_init_pipeline(self, custom_pipeline_args: dict[str, Any]):
+                return "custom_re_init_pipeline"
+
+            def custom_extension_method(self):
+                return "custom_extension_method"
+
+        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
+
+        with patch.object(DiffusionWorker, "__init__", return_value=None):
+            wrapper = WorkerWrapperBase(
+                gpu_id=0,
+                od_config=mock_od_config,
+                base_worker_class=DiffusionWorker,
+                worker_extension_cls=CustomExtension,
+                custom_pipeline_args=custom_args,
+            )
+
+            # Should still have the explicitly provided extension
+            assert CustomExtension in wrapper.worker.__class__.__bases__
+            assert hasattr(wrapper.worker, "custom_extension_method")
+
+    @patch("torch.cuda.empty_cache")
+    @patch("gc.collect")
+    @patch.object(DiffusionWorker, "__init__", return_value=None)
+    def test_re_init_pipeline_multiple_calls(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
+        """Test calling re_init_pipeline multiple times."""
+        wrapper = WorkerWrapperBase(
+            gpu_id=0,
+            od_config=mock_od_config,
+            base_worker_class=DiffusionWorker,
+            worker_extension_cls=CustomPipelineWorkerExtension,
+        )
+
+        # Setup mock model_runner
+        mock_model_runner = Mock()
+        mock_pipeline1 = Mock()
+        mock_pipeline2 = Mock()
+        mock_model_runner.pipeline = mock_pipeline1
+        wrapper.worker.model_runner = mock_model_runner
+        wrapper.worker.init_lora_manager = Mock()
+        wrapper.worker.load_model = Mock()
+
+        # First call
+        custom_args1 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
+        wrapper.worker.re_init_pipeline(custom_args1)
+
+        # Update pipeline for second call
+        mock_model_runner.pipeline = mock_pipeline2
+
+        # Second call
+        custom_args2 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
+        wrapper.worker.re_init_pipeline(custom_args2)
+
+        # Verify load_model was called twice with different pipelines
+        assert wrapper.worker.load_model.call_count == 2
+        assert wrapper.worker.init_lora_manager.call_count == 2