Update test-merge.yml

Signed-off-by: Kevin H. Luu <khluu000@gmail.com>
[ROCm] [CI] [Docker] Point to use the latest vLLM v0.16.0 stable version (#1500 )
2026-02-26 03:54:05 -08:00 · 2026-02-26 19:16:14 +08:00 · 2026-02-26 19:09:42 +08:00 · 2026-02-26 18:58:35 +08:00 · 2026-02-26 10:34:28 +00:00 · 2026-02-26 17:23:55 +08:00
188 changed files with 19842 additions and 7047 deletions
@@ -9,314 +9,32 @@ steps:
    agents:
      queue: "cpu_queue_premerge"

-  - label: "Upload Nightly Pipeline"
+  # L2 Test
+  - label: "Upload Ready Pipeline"
    depends_on: image-build
-    if: build.env("NIGHTLY") == "1"
+    key: upload-ready-pipeline
+    if: build.branch != "main"
    commands:
-      - buildkite-agent pipeline upload .buildkite/test-nightly.yaml
+      - buildkite-agent pipeline upload .buildkite/test-ready.yml
    agents:
      queue: "cpu_queue_premerge"

-  - label: "Simple Unit Test"
+  # L3 Test
+  - label: "Upload Merge Pipeline"
    depends_on: image-build
+    key: upload-merge-pipeline
+    if: build.branch == "main"
    commands:
-      - "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+      - buildkite-agent pipeline upload .buildkite/test-merge.yml
    agents:
-      queue: "gpu_1_queue"
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
+      queue: "cpu_queue_premerge"

-  - label: "Diffusion Model Test"
-    timeout_in_minutes: 20
+  # L4 Test
+  - label: "Upload Nightly Pipeline"
    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Diffusion Model CPU offloading Test"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
-    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Audio Generation Model Test"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
-    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Diffusion Cache Backend Test"
-    timeout_in_minutes: 15
-    depends_on: image-build
-    commands:
-      - pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
-    agents:
-      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Diffusion Sequence Parallelism Test"
-    timeout_in_minutes: 20
-    depends_on: image-build
-    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Diffusion GPU Worker Test"
-    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: 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
-    depends_on: image-build
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  # - label: "Bagel Text2Img Model Test with H100"
-  #   timeout_in_minutes: 30
-  #   depends_on: image-build
-  #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 1
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
+      queue: "cpu_queue_premerge"
@@ -1,16 +1,6 @@
 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
+    depends_on: upload-merge-pipeline
    commands:
      - "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
    agents:
@@ -27,7 +17,7 @@ steps:

  - label: "Diffusion Model Test"
    timeout_in_minutes: 20
-    depends_on: image-build
+    depends_on: upload-merge-pipeline
    commands:
      - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
    agents:
@@ -44,7 +34,7 @@ steps:

  - label: "Diffusion Images API LoRA E2E"
    timeout_in_minutes: 20
-    depends_on: image-build
+    depends_on: upload-merge-pipeline
    commands:
      - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
    agents:
@@ -61,7 +51,7 @@ steps:

  - label: "Diffusion Model CPU offloading Test"
    timeout_in_minutes: 20
-    depends_on: image-build
+    depends_on: upload-merge-pipeline
    commands:
      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
@@ -79,7 +69,7 @@ steps:

  - label: "Audio Generation Model Test"
    timeout_in_minutes: 20
-    depends_on: image-build
+    depends_on: upload-merge-pipeline
    commands:
      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
    agents:
@@ -96,9 +86,9 @@ steps:

  - label: "Diffusion Cache Backend Test"
    timeout_in_minutes: 15
-    depends_on: image-build
+    depends_on: upload-merge-pipeline
    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:
@@ -113,7 +103,7 @@ steps:

  - label: "Diffusion Sequence Parallelism Test"
    timeout_in_minutes: 20
-    depends_on: image-build
+    depends_on: upload-merge-pipeline
    commands:
      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
    agents:
@@ -131,9 +121,9 @@ steps:

  - label: "Diffusion Tensor Parallelism Test"
    timeout_in_minutes: 20
-    depends_on: image-build
+    depends_on: upload-merge-pipeline
    commands:
-      - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
+      - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
    agents:
      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
    plugins:
@@ -149,7 +139,7 @@ steps:

  - label: "Diffusion GPU Worker Test"
    timeout_in_minutes: 20
-    depends_on: image-build
+    depends_on: upload-merge-pipeline
    commands:
      - pytest -s -v tests/diffusion/test_diffusion_worker.py
    agents:
@@ -165,10 +155,9 @@ steps:
          volumes:
            - "/fsx/hf_cache:/fsx/hf_cache"

-
-  - label: "Benchmark&Engine Test"
+  - label: "Benchmark & Engine Test with H100"
    timeout_in_minutes: 15
-    depends_on: image-build
+    depends_on: upload-merge-pipeline
    commands:
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
      - pytest -s -v tests/benchmarks/test_serve_cli.py
@@ -179,7 +168,7 @@ steps:
      - kubernetes:
          podSpec:
            containers:
-              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                resources:
                  limits:
                    nvidia.com/gpu: 2
@@ -204,63 +193,63 @@ steps:

  - label: "Omni Model Test"
    timeout_in_minutes: 15
-    depends_on: image-build
+    depends_on: upload-merge-pipeline
    commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+      - 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
+    depends_on: upload-merge-pipeline
    commands:
-       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-       - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-       - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-       - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+      - 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"
+      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
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate

  - label: "Diffusion Image Edit Test with H100 (1 GPU)"
    timeout_in_minutes: 20
-    depends_on: image-build
+    depends_on: upload-merge-pipeline
    commands:
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
      - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
@@ -270,7 +259,7 @@ steps:
      - kubernetes:
          podSpec:
            containers:
-              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
                resources:
                  limits:
                    nvidia.com/gpu: 1
@@ -295,7 +284,7 @@ steps:

  # - label: "Bagel Text2Img Model Test with H100"
  #   timeout_in_minutes: 30
-  #   depends_on: image-build
+  #   depends_on: upload-merge-pipeline
  #   commands:
  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
@@ -305,7 +294,7 @@ steps:
  #     - kubernetes:
  #         podSpec:
  #           containers:
-  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
  #               resources:
  #                 limits:
  #                   nvidia.com/gpu: 1
@@ -1,12 +1,12 @@
 steps:
  - label: "Omni Model Test with H100"
-    timeout_in_minutes: 60
-    depends_on: image-build
+    timeout_in_minutes: 90
+    depends_on: upload-nightly-pipeline
    if: build.env("NIGHTLY") == "1"
    commands:
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py
-      - pytest -s -v tests/examples/online_serving/test_qwen3_omni.py
+      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
+      - pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
    agents:
      queue: "mithril-h100-pool"
    plugins:
@@ -38,28 +38,27 @@ steps:

  - label: "Omni Model Test"
    timeout_in_minutes: 60
-    depends_on: image-build
+    depends_on: upload-nightly-pipeline
    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
+      - pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
    agents:
      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
    plugins:
      - docker#v5.2.0:
          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
          always-pull: true
+          shm-size: "8gb"
          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
+    depends_on: upload-nightly-pipeline
    if: build.env("NIGHTLY") == "1"
    commands:
      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
@@ -0,0 +1,303 @@
+steps:
+  - label: "Simple Unit Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model Test"
+    timeout_in_minutes: 20
+    depends_on: upload-ready-pipeline
+    commands:
+      - 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:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model CPU offloading Test"
+    timeout_in_minutes: 20
+    depends_on: upload-ready-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Audio Generation Model Test"
+    timeout_in_minutes: 20
+    depends_on: upload-ready-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Cache Backend Test"
+    timeout_in_minutes: 15
+    depends_on: upload-ready-pipeline
+    commands:
+      - pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Sequence Parallelism Test"
+    timeout_in_minutes: 20
+    depends_on: upload-ready-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion GPU Worker Test"
+    timeout_in_minutes: 20
+    depends_on: upload-ready-pipeline
+    commands:
+      - pytest -s -v tests/diffusion/test_diffusion_worker.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  # - label: "Benchmark & Engine Test with H100"
+  #   timeout_in_minutes: 15
+  #   depends_on: upload-ready-pipeline
+  #   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: 17
+    depends_on: upload-ready-pipeline
+    commands:
+      - export VLLM_LOGGING_LEVEL=DEBUG
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  # - label: "Omni Model Test with H100"
+  #   timeout_in_minutes: 20
+  #   depends_on: upload-ready-pipeline
+  #   commands:
+  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+  #     # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+  #     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "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: upload-ready-pipeline
+    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
+  #   depends_on: upload-ready-pipeline
+  #   commands:
+  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #     - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: 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: upload-ready-pipeline
+  #   commands:
+  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: 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
@@ -14,7 +14,7 @@
        soft_fail: false
        commands:
          - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
-          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target final --progress plain ."
+          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
          - "docker push {{ docker_image_amd }}"
        key: "amd-build"
        env:
@@ -0,0 +1,629 @@
+#!/usr/bin/env python3
+"""
+PR Reviewer using GLM API for vllm-omni project.
+"""
+
+import json
+import logging
+import os
+import sys
+import time
+from dataclasses import dataclass
+from typing import Any, TypedDict
+
+import requests
+
+
+# Type definitions for API responses
+class PRDetails(TypedDict):
+    """Type definition for GitHub PR details response."""
+
+    title: str
+    body: str
+    number: int
+    state: str
+    user: dict[str, Any]
+
+
+class GLMMessage(TypedDict):
+    """Type definition for GLM API message."""
+
+    role: str
+    content: str
+
+
+class GLMChoice(TypedDict):
+    """Type definition for GLM API choice."""
+
+    message: GLMMessage
+    finish_reason: str
+
+
+class GLMResponse(TypedDict):
+    """Type definition for GLM API response."""
+
+    choices: list[GLMChoice]
+    usage: dict[str, int] | None
+
+
+class GitHubComment(TypedDict):
+    """Type definition for GitHub comment."""
+
+    id: int
+    body: str
+    created_at: str
+    user: dict[str, Any]
+
+
+# Configuration
+TRIGGER_PHRASE: str = "@vllm-omni-reviewer"
+DEFAULT_GLM_API_URL: str = "https://open.bigmodel.cn/api/paas/v4/chat/completions"  # noqa: E501
+DEFAULT_GLM_MODEL: str = "glm-5"
+DEFAULT_COOLDOWN_MINUTES: int = 5
+DEFAULT_MAX_RETRIES: int = 3
+DEFAULT_RETRY_DELAY: float = 1.0
+MAX_DIFF_SIZE: int = 100_000  # Maximum diff size in characters
+
+
+@dataclass
+class Config:
+    """Configuration for the PR reviewer."""
+
+    glm_api_url: str
+    glm_model: str
+    cooldown_minutes: int
+    max_retries: int
+    retry_delay: float
+    max_diff_size: int
+
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="[PR Reviewer] %(asctime)s - %(levelname)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def get_config() -> Config:
+    """Load configuration from environment variables with defaults."""
+    return Config(
+        glm_api_url=os.getenv("GLM_API_URL", DEFAULT_GLM_API_URL),
+        glm_model=os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL),
+        cooldown_minutes=int(
+            os.getenv(
+                "PR_REVIEWER_COOLDOWN_MINUTES",
+                str(DEFAULT_COOLDOWN_MINUTES),
+            )
+        ),
+        max_retries=int(
+            os.getenv(
+                "PR_REVIEWER_MAX_RETRIES",
+                str(DEFAULT_MAX_RETRIES),
+            )
+        ),
+        retry_delay=float(os.getenv("PR_REVIEWER_RETRY_DELAY", str(DEFAULT_RETRY_DELAY))),
+        max_diff_size=int(os.getenv("PR_REVIEWER_MAX_DIFF_SIZE", str(MAX_DIFF_SIZE))),  # noqa: E501
+    )
+
+
+def get_env_var(name: str) -> str:
+    """
+    Get an environment variable or raise an error.
+
+    Args:
+        name: Name of the environment variable.
+
+    Returns:
+        The value of the environment variable.
+
+    Raises:
+        SystemExit: If the environment variable is not set.
+    """
+    value = os.environ.get(name)
+    if not value:
+        logger.error(f"Environment variable {name} is not set")
+        sys.exit(1)
+    return value
+
+
+def check_trigger(comment_body: str) -> bool:
+    """
+    Check if the comment contains the trigger phrase.
+
+    Args:
+        comment_body: The body of the comment to check.
+
+    Returns:
+        True if the trigger phrase is found, False otherwise.
+    """
+    return TRIGGER_PHRASE in comment_body
+
+
+def fetch_pr_diff(
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    max_size: int = MAX_DIFF_SIZE,
+) -> str | None:
+    """
+    Fetch the diff for a pull request.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        max_size: Maximum diff size in characters.
+
+    Returns:
+        The diff content as a string, or None if fetching failed.
+        Returns empty string if diff is larger than max_size.
+    """
+    url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3.diff",
+    }
+
+    logger.info(f"Fetching PR diff from {url}")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code == 200:
+        diff: str = response.text
+        if len(diff) > max_size:
+            logger.warning(
+                f"Diff size ({len(diff)} bytes) exceeds maximum "
+                f"({max_size} bytes), truncating to first "
+                f"{max_size} bytes"
+            )
+            return diff[:max_size] + "\n\n... [Diff truncated due to size] ..."
+        logger.info(f"Successfully fetched diff ({len(diff)} bytes)")
+        return diff
+    else:
+        logger.error(f"Failed to fetch PR diff: {response.status_code}")
+        logger.error(f"Response: {response.text}")
+        return None
+
+
+def fetch_pr_details(
+    repo_name: str,
+    pr_number: int,
+    token: str,
+) -> PRDetails | None:
+    """
+    Fetch PR details including title and description.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+
+    Returns:
+        A dictionary containing PR details, or None if fetching failed.
+    """
+    url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    logger.info(f"Fetching PR details from {url}")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code == 200:
+        return response.json()
+    else:
+        logger.error(f"Failed to fetch PR details: {response.status_code}")
+        return None
+
+
+def build_review_prompt(pr_title: str, pr_description: str, diff: str) -> str:
+    """
+    Build the prompt for the GLM-4.7 API.
+
+    Args:
+        pr_title: The title of the pull request.
+        pr_description: The description/body of the pull request.
+        diff: The diff content of the pull request.
+
+    Returns:
+        The formatted prompt string for the API.
+    """
+    return f"""You are an expert code reviewer for the VLLM-Omni project. \
+Please review the following pull request:
+
+## Pull Request Details
+**Title:** {pr_title}
+
+**Description:**
+{pr_description if pr_description else "No description provided."}
+
+## Code Changes (Diff)
+{diff}
+
+## Review Guidelines
+
+Please provide a comprehensive code review with the following sections:
+
+### 1. Overview
+- Brief summary of the changes
+- Overall assessment (positive, neutral, or concerns)
+
+### 2. Code Quality
+- Code style and consistency
+- Potential bugs or edge cases
+- Performance considerations
+- Error handling
+
+### 3. Architecture & Design
+- Integration with existing codebase
+- Design patterns and best practices
+- Potential improvements
+
+### 4. Security & Safety
+- Security concerns (if any)
+- Resource management
+- Input validation
+
+### 5. Testing & Documentation
+- Test coverage considerations
+- Documentation completeness
+- Examples and usage clarity
+
+### 6. Specific Suggestions
+- Line-by-line specific feedback (use `file:line` format)
+- Concrete actionable suggestions
+- Code examples for improvements (if applicable)
+
+### 7. Approval Status
+- **LGTM** (Looks Good To Me) if the PR is ready to merge
+- **LGTM with suggestions** if the PR is good but has minor suggestions
+- **Changes requested** if significant changes are needed
+
+## Important Notes
+- Be constructive and helpful
+- Focus on objective technical feedback
+- Acknowledge good practices when you see them
+- Prioritize critical issues over nitpicks
+- If the diff is empty or minimal, acknowledge this and provide
+  any relevant context-specific guidance
+
+Please format your response in Markdown with clear section headers.
+"""
+
+
+def validate_glm_response(data: dict[str, Any]) -> str | None:
+    """
+    Validate and extract content from GLM API response.
+
+    Args:
+        data: The response data from GLM API.
+
+    Returns:
+        The review content string if valid, None otherwise.
+    """
+    # Check if choices exists and is a non-empty list
+    if "choices" not in data:
+        logger.error("GLM API response missing 'choices' field")
+        logger.error(f"Response structure: {json.dumps(data, indent=2)}")
+        return None
+
+    choices = data["choices"]
+    if not isinstance(choices, list):
+        logger.error(f"GLM API 'choices' is not a list: {type(choices)}")
+        return None
+
+    if len(choices) == 0:
+        logger.error("GLM API 'choices' is an empty list")
+        return None
+
+    # Check if first choice has message
+    try:
+        first_choice = choices[0]
+        if not isinstance(first_choice, dict):
+            logger.error(f"GLM API choice is not a dict: {type(first_choice)}")
+            return None
+
+        if "message" not in first_choice:
+            logger.error("GLM API choice missing 'message' field")
+            logger.error(f"Choice structure: {json.dumps(first_choice, indent=2)}")  # noqa: E501
+            return None
+
+        message = first_choice["message"]
+        if not isinstance(message, dict):
+            logger.error(f"GLM API message is not a dict: {type(message)}")
+            return None
+
+        if "content" not in message:
+            logger.error("GLM API message missing 'content' field")
+            logger.error(f"Message structure: {json.dumps(message, indent=2)}")
+            return None
+
+        content = message["content"]
+        if not isinstance(content, str):
+            logger.error(f"GLM API content is not a string: {type(content)}")
+            return None
+
+        return content
+
+    except (KeyError, IndexError, TypeError) as e:
+        logger.error(f"Failed to parse GLM API response: {e}")
+        logger.error(f"Response: {json.dumps(data, indent=2)}")
+        return None
+
+
+def call_glm_api(prompt: str, api_key: str, config: Config) -> str | None:
+    """
+    Call the GLM-4.7 API to get code review with retry logic.
+
+    Args:
+        prompt: The prompt to send to the API.
+        api_key: The GLM API key.
+        config: Configuration object.
+
+    Returns:
+        The review content as a string, or None if all retries failed.
+    """
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+    }
+
+    payload: dict[str, Any] = {
+        "model": config.glm_model,
+        "messages": [{"role": "user", "content": prompt}],
+        "temperature": 0.3,
+        "max_tokens": 32000,
+        "top_p": 0.9,
+    }
+
+    last_error: str | None = None
+
+    for attempt in range(config.max_retries):
+        try:
+            logger.info(f"Calling GLM API ({config.glm_model}) - Attempt {attempt + 1}/{config.max_retries}")
+            response = requests.post(
+                config.glm_api_url,
+                headers=headers,
+                json=payload,
+                timeout=120,
+            )
+
+            if response.status_code == 200:
+                data = response.json()
+                review = validate_glm_response(data)
+                if review:
+                    logger.info(f"Successfully received review ({len(review)} chars)")  # noqa: E501
+                    return review
+                else:
+                    last_error = "Failed to validate API response structure"
+                    logger.error(last_error)
+            else:
+                last_error = f"GLM API request failed: {response.status_code} - {response.text}"
+                logger.error(last_error)
+
+        except requests.exceptions.Timeout:
+            last_error = f"GLM API request timed out (attempt {attempt + 1})"
+            logger.error(last_error)
+        except requests.exceptions.RequestException as e:
+            last_error = f"GLM API request exception: {e}"
+            logger.error(last_error)
+        except json.JSONDecodeError as e:
+            last_error = f"Failed to decode GLM API response as JSON: {e}"
+            logger.error(last_error)
+
+        # Exponential backoff before retry
+        if attempt < config.max_retries - 1:
+            wait_time: float = config.retry_delay * (2**attempt)
+            logger.info(f"Waiting {wait_time}s before retry...")  # noqa: E501
+            time.sleep(wait_time)
+
+    logger.error(
+        f"All {config.max_retries} attempts failed. Last error: {last_error}"  # noqa: E501
+    )
+    return None
+
+
+def check_cooldown(  # noqa: E501
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    cooldown_minutes: int,
+) -> bool:
+    """
+    Check if the PR is within the cooldown period.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        cooldown_minutes: Cooldown period in minutes.
+
+    Returns:
+        True if within cooldown period (should skip), False otherwise.
+    """
+    from datetime import datetime, timedelta
+
+    url: str = (
+        f"https://api.github.com/repos/{repo_name}/issues/"
+        f"{pr_number}/comments"  # noqa: E501
+    )
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    logger.info(f"Checking cooldown period ({cooldown_minutes} minutes)")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code != 200:
+        logger.warning(f"Failed to check cooldown: {response.status_code}, proceeding with review")
+        return False
+
+    comments: list[dict[str, Any]] = response.json()
+    cutoff_time: datetime = datetime.utcnow() - timedelta(minutes=cooldown_minutes)  # noqa: E501
+
+    for comment in reversed(comments):
+        # Check if this is a bot comment
+        body: str = comment.get("body", "")
+        if "VLLM-Omni PR Review" in body or "PR Reviewer Bot" in body:
+            created_at_str: str = comment.get("created_at", "")
+            try:
+                # Parse GitHub timestamp format
+                created_at = datetime.fromisoformat(created_at_str.replace("Z", "+00:00"))
+                created_at = created_at.replace(tzinfo=None)
+                if created_at > cutoff_time:
+                    logger.info(f"PR is within cooldown period (last review: {created_at_str})")
+                    return True
+            except ValueError:
+                logger.warning(f"Failed to parse comment timestamp: {created_at_str}")  # noqa: E501
+                continue
+
+    logger.info("PR is outside cooldown period, proceeding with review")
+    return False
+
+
+def post_review_comment(  # noqa: E501
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    review: str,
+) -> bool:
+    """
+    Post the review as a comment on the PR.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        review: The review content to post.
+
+    Returns:
+        True if posting succeeded, False otherwise.
+    """
+    url: str = (
+        f"https://api.github.com/repos/{repo_name}/issues/"
+        f"{pr_number}/comments"  # noqa: E501
+    )
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    # Format the review comment
+    comment_body: str = f"""## 🤖 VLLM-Omni PR Review
+
+{review}
+
+---
+*This review was generated automatically by the VLLM-Omni PR Reviewer Bot
+using {os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL)}.*
+"""
+
+    payload: dict[str, str] = {"body": comment_body}
+
+    logger.info(f"Posting review comment to PR #{pr_number}")
+    response = requests.post(url, headers=headers, json=payload, timeout=30)
+
+    if response.status_code == 201:
+        logger.info("Successfully posted review comment")
+        return True
+    else:
+        logger.error(f"Failed to post comment: {response.status_code}")
+        logger.error(f"Response: {response.text}")
+        return False
+
+
+def main() -> int:
+    """
+    Main entry point for the PR reviewer bot.
+
+    Returns:
+        0 on success, 1 on error.
+    """
+    logger.info("VLLM-Omni PR Reviewer Bot starting...")
+
+    # Load configuration
+    config: Config = get_config()
+    logger.info(
+        f"Configuration: model={config.glm_model}, "
+        f"cooldown={config.cooldown_minutes}min, "
+        f"max_retries={config.max_retries}"
+    )
+
+    # Get environment variables
+    token: str = get_env_var("GITHUB_TOKEN")
+    api_key: str = get_env_var("GLM_API_KEY")
+    repo_name: str = get_env_var("REPO_NAME")
+    pr_number_str: str = get_env_var("PR_NUMBER")
+    comment_body: str = get_env_var("COMMENT_BODY")
+
+    try:
+        pr_number: int = int(pr_number_str)
+    except ValueError:
+        logger.error(f"Invalid PR number: {pr_number_str}")
+        return 1
+
+    logger.info(f"Repository: {repo_name}")
+    logger.info(f"PR Number: {pr_number}")
+
+    # Check if the comment contains the trigger phrase
+    if not check_trigger(comment_body):
+        logger.info(
+            f"Comment does not contain trigger phrase '{TRIGGER_PHRASE}', exiting"  # noqa: E501
+        )
+        return 0
+
+    logger.info("Trigger phrase detected! Starting review process...")
+
+    # Check cooldown period
+    if check_cooldown(repo_name, pr_number, token, config.cooldown_minutes):
+        logger.info("Skipping review due to cooldown period")
+        return 0
+
+    # Fetch PR details
+    logger.info("Step 1/4: Fetching PR details...")
+    pr_details: PRDetails | None = fetch_pr_details(repo_name, pr_number, token)  # noqa: E501
+    if not pr_details:
+        logger.error("Failed to fetch PR details")
+        return 1
+
+    pr_title: str = pr_details.get("title", "Unknown")
+    pr_description: str = pr_details.get("body", "")
+
+    logger.info(f"PR Title: {pr_title}")
+
+    # Fetch PR diff
+    logger.info("Step 2/4: Fetching PR diff...")
+    diff: str | None = fetch_pr_diff(repo_name, pr_number, token, config.max_diff_size)
+    if diff is None:
+        logger.error("Failed to fetch PR diff")
+        return 1
+
+    if not diff:
+        logger.warning("Warning: Empty diff - this might be a draft PR or no code changes")
+
+    # Build prompt
+    logger.info("Step 3/4: Building review prompt...")
+    prompt: str = build_review_prompt(pr_title, pr_description, diff)
+
+    # Call GLM API
+    logger.info("Step 4/4: Calling GLM API...")
+    review: str | None = call_glm_api(prompt, api_key, config)
+    if not review:
+        logger.error("Failed to get review from GLM API")
+        return 1
+
+    # Post review comment
+    logger.info("Posting review comment...")
+    if not post_review_comment(repo_name, pr_number, token, review):
+        logger.error("Failed to post review comment")
+        return 1
+
+    logger.info("PR review completed successfully!")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
@@ -0,0 +1,62 @@
+name: VLLM-Omni PR Reviewer
+
+on:
+  issue_comment:
+    types: [created]
+
+permissions:
+  contents: read
+  pull-requests: write
+  issues: write
+
+jobs:
+  pr-reviewer:
+    name: Review Pull Request
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    # Only run when the comment is from a collaborator/owner/member
+    if: |
+      github.event_name == 'issue_comment' &&
+      github.event.issue.pull_request != null &&
+      contains(github.event.comment.body, '@vllm-omni-reviewer') &&
+      (github.event.comment.author_association == 'MEMBER' ||
+       github.event.comment.author_association == 'COLLABORATOR' ||
+       github.event.comment.author_association == 'OWNER')
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4.2.2
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v5.3.0
+        with:
+          python-version: '3.11'
+          cache: 'pip'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip==24.0
+          pip install requests==2.31.0 pyyaml==6.0.1
+
+      - name: Run PR Reviewer
+        id: reviewer
+        env:
+          GLM_API_KEY: ${{ secrets.GLM_API_KEY }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}
+          COMMENT_BODY: ${{ github.event.comment.body }}
+          REPO_NAME: ${{ github.repository }}
+          PR_REVIEWER_COOLDOWN_MINUTES: 5
+          PR_REVIEWER_MAX_RETRIES: 3
+        run: |
+          python .github/scripts/pr_reviewer.py 2>&1 | tee "/tmp/pr_review_${PR_NUMBER}.log"
+
+      - name: Upload review logs
+        if: always()
+        uses: actions/upload-artifact@v4.4.3
+        with:
+          name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
+          path: /tmp/pr_review_*.log
+          retention-days: 7
+          if-no-files-found: ignore
@@ -0,0 +1,198 @@
+import base64
+import mimetypes
+import os
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any
+
+import aiohttp
+from tqdm import tqdm
+
+
+@dataclass
+class RequestFuncInput:
+    prompt: str
+    api_url: str
+    model: str
+    width: int | None = None
+    height: int | None = None
+    num_frames: int | None = None
+    num_inference_steps: int | None = None
+    seed: int | None = None
+    fps: int | None = None
+    timestamp: float | None = None
+    slo_ms: float | None = None
+    extra_body: dict[str, Any] = field(default_factory=dict)
+    image_paths: list[str] | None = None
+    request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+
+
+@dataclass
+class RequestFuncOutput:
+    success: bool = False
+    latency: float = 0.0
+    error: str = ""
+    start_time: float = 0.0
+    response_body: dict[str, Any] = field(default_factory=dict)
+    peak_memory_mb: float = 0.0
+    slo_achieved: bool | None = None
+
+
+def _guess_mime_type(path: str) -> str:
+    mime, _ = mimetypes.guess_type(path)
+    return mime or "application/octet-stream"
+
+
+def _encode_image_as_data_url(path: str) -> str:
+    with open(path, "rb") as f:
+        encoded = base64.b64encode(f.read()).decode("utf-8")
+    mime = _guess_mime_type(path)
+    return f"data:{mime};base64,{encoded}"
+
+
+async def async_request_chat_completions(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput:
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    extra_body = dict(input.extra_body)
+    if input.width and input.height:
+        extra_body.setdefault("height", input.height)
+        extra_body.setdefault("width", input.width)
+    if input.num_frames:
+        extra_body.setdefault("num_frames", input.num_frames)
+    if input.num_inference_steps:
+        extra_body.setdefault("num_inference_steps", input.num_inference_steps)
+    if input.seed is not None:
+        extra_body.setdefault("seed", input.seed)
+    if input.fps:
+        extra_body.setdefault("fps", input.fps)
+
+    if input.image_paths and len(input.image_paths) > 0:
+        content = []
+        if input.prompt:
+            content.append({"type": "text", "text": input.prompt})
+        for img_path in input.image_paths:
+            if not os.path.exists(img_path):
+                output.error = f"Image file not found: {img_path}"
+                output.success = False
+                if pbar:
+                    pbar.update(1)
+                return output
+            content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": _encode_image_as_data_url(img_path)},
+                }
+            )
+        messages = [{"role": "user", "content": content}]
+    else:
+        messages = [{"role": "user", "content": input.prompt}]
+
+    payload = {
+        "model": input.model,
+        "messages": messages,
+    }
+    if extra_body:
+        payload["extra_body"] = extra_body
+
+    try:
+        async with session.post(input.api_url, json=payload) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+                if "peak_memory_mb" in resp_json:
+                    output.peak_memory_mb = resp_json["peak_memory_mb"]
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+async def async_request_openai_images(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput:
+    """
+    Send request to OpenAI's /v1/images/generations endpoint.
+    """
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    # Build size string from width/height
+    width = input.width or 1024
+    height = input.height or 1024
+    size = f"{width}x{height}"
+
+    payload: dict[str, Any] = {
+        "model": input.model,
+        "prompt": input.prompt,
+        "n": 1,
+        "size": size,
+        "response_format": "b64_json",
+    }
+
+    # Add optional parameters
+    if input.seed is not None:
+        payload["seed"] = input.seed
+    if input.num_inference_steps is not None:
+        payload["num_inference_steps"] = input.num_inference_steps
+
+    # Add any extra body parameters
+    if input.extra_body:
+        for key, value in input.extra_body.items():
+            if key not in payload:
+                payload[key] = value
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": "Bearer EMPTY",
+    }
+
+    try:
+        async with session.post(input.api_url, json=payload, headers=headers) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+                # Check for usage/memory info if available
+                if "usage" in resp_json and "peak_memory_mb" in resp_json.get("usage", {}):
+                    output.peak_memory_mb = resp_json["usage"]["peak_memory_mb"]
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+backends_function_mapping = {
+    "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
+    "openai": (async_request_openai_images, "/v1/images/generations"),
+}
@@ -6,80 +6,60 @@
 Benchmark online serving for diffusion models (Image/Video Generation).
 If you want to use i2v, i2i dataset, you should `uv pip install gdown` first

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

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


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

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

 """

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

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


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


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

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

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

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

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

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

-# Install vLLM wheel for the specific commit
-# URL format: https://wheels.vllm.ai/<commit_hash>/
-# Current vLLM 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 []
@@ -7,8 +7,9 @@ WORKDIR ${APP_DIR}

 COPY . .

-# Install vllm-omni with dev dependencies
-RUN pip install --no-cache-dir -e . --no-build-isolation
+RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
+ && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
+RUN pip install -v -e .

 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

@@ -7,8 +7,9 @@ WORKDIR ${APP_DIR}

 COPY . .

-# Install vllm-omni with dev dependencies
-RUN pip install --no-cache-dir -e . --no-build-isolation
+RUN sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
+ && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
+RUN pip install -v -e .

 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn

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

 ARG COMMON_WORKDIR=/app
-
 WORKDIR ${COMMON_WORKDIR}

 # Step 1: Setup - Install system dependencies
@@ -11,18 +23,31 @@ RUN apt-get update && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

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

-# Step 2: Copy vllm-omni code and install without uv
+# Step 3: Copy vllm-omni code and install without uv
+RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
 COPY . ${COMMON_WORKDIR}/vllm-omni
 RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" --no-build-isolation

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

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

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

-#Set entrypoint for vllm-openai official images
-FROM final AS vllm-openai
+# Set entrypoint for vllm-openai official images
+FROM base AS vllm-openai
 ENTRYPOINT ["vllm", "serve", "--omni"]
@@ -0,0 +1,139 @@
+# Argument to configure vllm base image if pre-built
+ARG VLLM_BASE=vllm-base
+
+FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
+
+WORKDIR /workspace/
+
+ARG PYTHON_VERSION=3.12
+ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
+
+RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
+    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
+    add-apt-repository -y ppa:kobuk-team/intel-graphics
+
+RUN apt clean && apt-get update -y && \
+    apt-get install -y --no-install-recommends --fix-missing \
+    curl \
+    ffmpeg \
+    git \
+    libsndfile1 \
+    libsm6 \
+    libxext6 \
+    libgl1 \
+    lsb-release \
+    libaio-dev \
+    numactl \
+    wget \
+    vim \
+    python3.12 \
+    python3.12-dev \
+    python3-pip
+
+RUN apt update && apt upgrade -y && \
+    apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \
+    apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
+
+ENV PATH="/root/.local/bin:$PATH"
+ENV VIRTUAL_ENV="/opt/venv"
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh
+RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+
+# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
+ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
+RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
+    bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
+    rm "${ONECCL_INSTALLER}" && \
+    echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
+    echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc
+RUN rm -f /opt/intel/oneapi/ccl/latest && \
+    ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest
+
+SHELL ["bash", "-c"]
+CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
+
+WORKDIR /workspace/
+ENV UV_HTTP_TIMEOUT=500
+
+# Configure package index for XPU
+ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+ENV UV_INDEX_STRATEGY="unsafe-best-match"
+ENV UV_LINK_MODE="copy"
+
+ARG VLLM_VERSION=v0.16.0
+RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
+WORKDIR /workspace/vllm
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --upgrade pip && \
+    uv pip install -r requirements/xpu.txt
+
+ # used for suffix method speculative decoding
+ # build deps for proto + nanobind-based extensions to set up the build environment
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install grpcio-tools protobuf nanobind
+ # arctic-inference is built from source which needs torch-xpu properly installed first
+RUN --mount=type=cache,target=/root/.cache/uv \
+    source /opt/intel/oneapi/setvars.sh --force && \
+    source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
+    export CMAKE_PREFIX_PATH="$(python -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
+    uv pip install --no-build-isolation arctic-inference==0.1.1
+
+ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
+
+ENV VLLM_TARGET_DEVICE=xpu
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --no-build-isolation .
+
+CMD ["/bin/bash"]
+
+FROM vllm-base AS vllm-openai
+
+# install additional dependencies for openai api server
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope
+
+# install development dependencies (for testing)
+RUN uv pip install -e tests/vllm_test_utils
+
+# install nixl from source code
+ENV NIXL_VERSION=0.7.0
+RUN python /workspace/vllm/tools/install_nixl_from_source_ubuntu.py
+
+# ensure vllm is properly installed
+RUN python -c "import vllm, inspect; print(vllm.__file__)"
+RUN uv pip show vllm
+
+CMD ["/bin/bash"]
+
+ENTRYPOINT []
+
+FROM ${VLLM_BASE} AS vllm-omni
+
+WORKDIR /workspace/vllm-omni
+COPY . .
+
+ENV VLLM_OMNI_TARGET_DEVICE=xpu
+RUN uv pip install --no-cache-dir ".[dev]" --no-build-isolation
+
+# FIX triton
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip uninstall triton triton-xpu && \
+    uv pip install triton-xpu==3.6.0
+
+# remove torch bundled oneccl to avoid conflicts
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip uninstall oneccl oneccl-devel
+
+FROM vllm-omni AS vllm-omni-openai
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT ["vllm", "serve", "--omni"]
@@ -40,7 +40,7 @@ nav:
  - Features:
    - Sleep Mode: features/sleep_mode.md
    - Diffusion Features:
-      - Overview: user_guide/diffusion_acceleration.md
+      - Acceleration Overview: user_guide/diffusion_acceleration.md
      - TeaCache: user_guide/diffusion/teacache.md
      - Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
      - Quantization:
@@ -48,6 +48,7 @@ nav:
        - FP8: user_guide/diffusion/quantization/fp8.md
      - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
      - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
+      - LoRA: user_guide/diffusion/lora.md
    - ComfyUI: features/comfyui.md
 - Developer Guide:
  - General:
@@ -58,12 +59,6 @@ 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
@@ -72,6 +67,11 @@ nav:
      - design/feature/disaggregated_inference.md
      - design/feature/ray_based_execution.md
      - design/feature/omni_connectors/
+      - design/feature/cfg_parallel.md
+      - design/feature/sequence_parallel.md
+      - design/feature/tensor_parallel.md
+      - design/feature/cache_dit.md
+      - design/feature/teacache.md
    - Module Design:
      - design/module/ar_module.md
      - design/module/dit_module.md
@@ -0,0 +1,9 @@
+# vllm-omni serve
+
+## JSON CLI Arguments
+
+--8<-- "docs/cli/json_tip.inc.md"
+
+## Arguments
+
+--8<-- "docs/generated/argparse_omni/omni_serve.inc.md"
@@ -273,3 +273,15 @@ Default: `True`
 Penalty applied to tokens that have already appeared in the generated sequence. Values greater than 1.0 discourage repetition, while values less than 1.0 encourage it. A value of 1.0 applies no penalty.

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

-# Metrics vLLM-Omni:
+# Metrics

 You can use these metrics in production to monitor the health and performance of the vLLM-omni system. Typical scenarios include:
+
 - **Performance Monitoring**: Track throughput (e.g., `e2e_avg_tokens_per_s`), latency (e.g., `e2e_total_ms`), and resource utilization to verify that the system meets expected standards.
+
 - **Debugging and Troubleshooting**: Use detailed per-request metrics to diagnose issues, such as high transfer times or unexpected token counts.

 ## How to Enable and View Metrics

-### 1. Start the Service with Metrics Logging
+### 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
+### Send a Request

 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
 ```

-### 3. What You Will See
+### What You Will See

 With `--log-stats` enabled, the server will output detailed metrics logs after each request. Example output:

@@ -69,9 +71,13 @@ With `--log-stats` enabled, the server will output detailed metrics logs after e


 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.
@@ -87,6 +93,8 @@ For **online inference** (serving mode), the summary is always per-request. `e2e

 ## Parameter Details

+### Summary Metrics
+
 | Field                     | Meaning                                                                                       |
 |---------------------------|----------------------------------------------------------------------------------------------|
 | `e2e_requests`            | Number of completed requests.                                                                |
@@ -98,7 +106,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e

 ---

-## E2E Table (per request)
+### E2E Table (per request)

 | Field                     | Meaning                                                               |
 |---------------------------|-----------------------------------------------------------------------|
@@ -110,7 +118,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e

 ---

-## Stage Table (per stage event / request)
+### Stage Table (per stage event / request)

 | Field                     | Meaning                                                                                         |
 |---------------------------|-------------------------------------------------------------------------------------------------|
@@ -125,7 +133,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e

 ---

-## Transfer Table (per edge / request)
+### Transfer Table (per edge / request)

 | Field                | Meaning                                                                   |
 |----------------------|---------------------------------------------------------------------------|
@@ -135,31 +143,31 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
 | `in_flight_time_ms`  | In-flight time in ms.                                                     |


-## Expectation of the Numbers (Verification)
+### 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
+**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`.
+so `e2e_total_tokens = 4,860 + 67 + 275 + 0 = 5,202`, which matches the table value `e2e_total_tokens`.
+
+**transfers_total_time_ms**

-### 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).
+192.581 + 53.315 = **245.896** = transfers_total_time_ms, which matches the calculation (difference is due to rounding)
@@ -1,4 +1,4 @@
-# Adding a Diffusion Model to vLLM-Omni
+# Adding a Diffusion Model

 This guide walks you through adding a new diffusion model to vLLM-Omni. We use **Qwen-Image** as the primary example, with references to other models (LongCat, Flux, Wan2.2) to illustrate different patterns.

@@ -680,7 +680,7 @@ vLLM-Omni automatically compiles blocks in `_repeated_blocks` when `torch.compil

 ### Tensor Parallelism

-See detailed guide: [How to add Tensor Parallel support](../features/tensor_parallel.md)
+See detailed guide: [How to add Tensor Parallel support](../../design/feature/tensor_parallel.md)

 **Quick setup:**

@@ -694,7 +694,7 @@ omni = Omni(model="your-model", tensor_parallel_size=2)

 ### CFG Parallelism

-See detailed guide: [How to add CFG-Parallel support](../features/cfg_parallel.md)
+See detailed guide: [How to add CFG-Parallel support](../../design/feature/cfg_parallel.md)

 **Quick setup:**

@@ -708,7 +708,7 @@ omni = Omni(model="your-model", cfg_parallel_size=2)

 ### Sequence Parallelism

-See detailed guide: [How to add Sequence Parallel support](../features/sequence_parallel.md)
+See detailed guide: [How to add Sequence Parallel support](../../design/feature/sequence_parallel.md)

 **Quick setup:**

@@ -724,7 +724,7 @@ omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)

 #### TeaCache

-See detailed guide: [How to add TeaCache support](../features/teacache.md)
+See detailed guide: [How to add TeaCache support](../../design/feature/teacache.md)

 **Quick setup:**

@@ -744,7 +744,7 @@ omni = Omni(model="your-model",

 #### Cache-DiT

-See detailed guide: [How to add Cache-DiT support](../features/cache_dit.md)
+See detailed guide: [How to add Cache-DiT support](../../design/feature/cache_dit.md)

 **Quick setup:**

@@ -1,4 +1,4 @@
-# Support Cache-DiT
+# 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.

@@ -1,4 +1,4 @@
-# Support CFG-Parallel
+# 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.

@@ -6,7 +6,8 @@ This guide explains how to configure and use distributed connectors
 Backend-specific setup lives in separate docs:

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

 ## Overview
@@ -19,7 +20,8 @@ Current connectors operate in D2H2D (device to host to device) mode.
 | Use Case | Recommended Connector | Notes |
 | :--- | :--- | :--- |
 | Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
-| Multi node (Mooncake) | MooncakeConnector | Requires Mooncake Master + metadata server. |
+| Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
+| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
 | Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |

 ## Core API
@@ -1,4 +1,4 @@
-# MooncakeConnector
+# MooncakeStoreConnector

 ## When to Use

@@ -38,7 +38,7 @@ Define the connector in runtime:
 runtime:
  connectors:
    connector_of_mooncake:
-      name: MooncakeConnector
+      name: MooncakeStoreConnector
      extra:
        host: "127.0.0.1"
        metadata_server: "http://<MASTER_IP>:8080/metadata"
@@ -0,0 +1,207 @@
+# MooncakeTransferEngineConnector
+
+## When to Use
+
+Best for high-performance multi-node data transfer between stages using Mooncake
+Transfer Engine. Supports both RDMA and TCP protocols with a managed memory pool,
+zero-copy deserialization, and optional GPUDirect RDMA. Applicable to any
+inter-stage data (KV caches, request payloads, etc.), not limited to KV cache transfer.
+
+Compared to `MooncakeStoreConnector` (TCP key-value store), this connector
+provides **~60x faster** data transfer via RDMA direct memory access.
+
+## Installation
+
+```bash
+pip install mooncake-transfer-engine
+```
+
+Ensure RDMA drivers are installed on all nodes (e.g., Mellanox OFED for
+InfiniBand/RoCE NICs).
+
+## Configuration
+
+Define the connector in runtime:
+
+```yaml
+runtime:
+  connectors:
+    rdma_connector:
+      name: MooncakeTransferEngineConnector
+      extra:
+        host: "auto"                  # Auto-detect local RDMA IP
+        zmq_port: 50051               # ZMQ base port (see "Port Offset Scheme" below)
+        protocol: "rdma"              # "rdma" or "tcp"
+        device_name: ""               # RDMA device (e.g., "mlx5_0"), empty for auto-detect
+        memory_pool_size: 2147483648  # 2GB memory pool
+        memory_pool_device: "cpu"     # "cpu" for pinned memory, "cuda" for GPUDirect RDMA
+```
+
+Wire stages to the connector:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: rdma_connector
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: rdma_connector
+```
+
+## Parameters
+
+### Required
+
+| Parameter | Description |
+|---|---|
+| `role` | **Internal, do not set manually.** Auto-injected by the orchestration layer (`"sender"` for `output_connectors`, `"receiver"` for `input_connectors`). Defaults to `"sender"` if omitted. |
+| `host` | Local IP address for RDMA. `"auto"` detects from network interfaces. |
+| `protocol` | Transport protocol: `"rdma"` (InfiniBand/RoCE) or `"tcp"`. |
+
+### Memory Pool
+
+| Parameter | Default | Description |
+|---|---|---|
+| `memory_pool_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. |
+| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity). |
+
+### Networking
+
+| Parameter | Default | Description |
+|---|---|---|
+| `zmq_port` | 50051 | ZMQ **base** port. The orchestration layer computes the actual port as `base + purpose_offset + stage_offset` (see table below). Users only set this base value. |
+| `sender_host` | `None` | **Internal.** Receiver-side only — dynamically resolved via `update_sender_info()`. Not needed in YAML. |
+| `sender_zmq_port` | `None` | **Internal.** Receiver-side only — defaults to the sender's adjusted port. Not needed in YAML. |
+| `device_name` | `""` | RDMA device name (e.g., `"mlx5_0"`). Empty for auto-detect. Can also be set via `RDMA_DEVICE_NAME` env var. |
+
+#### ZMQ Port Offset Scheme
+
+To avoid port conflicts when multiple edges, purposes, DP replicas, or TP ranks share the same node, the actual ZMQ port is computed as:
+
+```
+side_channel_port = zmq_port + purpose_offset + stage_offset + dp_index * tp_size
+sender_listen     = side_channel_port + tp_rank
+receiver_connect  = remote_side_channel_port + tp_rank
+```
+
+| Component | Value | Description |
+|---|---|---|
+| `zmq_port` | 50051 (default) | Base port from YAML config |
+| `purpose_offset` | `request_forwarding` = 0, `kv_transfer` = 100 | Separates control-plane vs KV-cache connections |
+| `stage_offset` | `int(from_stage)` (0, 1, 2...) | Separates edges from different source stages |
+| `dp_index * tp_size` | e.g., DP1 × TP2 = 2 | Each DP replica reserves a port range of size `tp_size` (following vLLM convention: `VLLM_MOONCAKE_BOOTSTRAP_PORT + dp_index * tp_size`) |
+| `tp_rank` | 0, 1, 2... | Each TP rank within a DP replica uses its own port |
+| orchestrator | +200 | Extra offset when caller is the orchestrator (avoids collision with stage workers on the same node) |
+
+**Example** (base=50051, stage 0→1, DP=2, TP=2, kv_transfer):
+
+| Caller | DP | TP rank | Port |
+|---|---|---|---|
+| Stage worker | DP0 | rank 0 | `50051 + 100 + 0 + 0×2 + 0 = 50151` |
+| Stage worker | DP0 | rank 1 | `50051 + 100 + 0 + 0×2 + 1 = 50152` |
+| Stage worker | DP1 | rank 0 | `50051 + 100 + 0 + 1×2 + 0 = 50153` |
+| Stage worker | DP1 | rank 1 | `50051 + 100 + 0 + 1×2 + 1 = 50154` |
+| Orchestrator | — | — | `50051 + 200 + 0 = 50251` |
+
+## Memory Pool Modes
+
+| Mode | Config | Data Flow | Best For |
+|---|---|---|---|
+| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
+| GPUDirect | `memory_pool_device: "cuda"` | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
+
+> **Note**: GPUDirect RDMA requires the NIC and GPU to share a direct PCIe
+> switch (PIX topology). On systems where they are connected via PXB or NODE,
+> CPU pinned memory is faster due to GPU BAR1 bandwidth limitations.
+
+## Environment Variables
+
+| Variable | Description |
+|---|---|
+| `RDMA_DEVICE_NAME` | Override RDMA device name (e.g., `mlx5_0`). |
+| `MC_IB_PCI_RELAXED_ORDERING` | Set to `1` to enable PCIe relaxed ordering for GPUDirect. |
+
+## Docker / Container Setup
+
+RDMA requires host-level device access:
+
+```bash
+docker run -it \
+    --cap-add=SYS_PTRACE \
+    --cap-add=IPC_LOCK \
+    --security-opt seccomp=unconfined \
+    --network=host \
+    --device=/dev/infiniband \
+    -v /sys/class/infiniband:/sys/class/infiniband:ro \
+    your-image:tag
+```
+
+## Performance
+
+Benchmark results on H800 GPUs with mlx5_0 RDMA NIC (~186 MB KV cache):
+
+| Metric | MooncakeStoreConnector | MooncakeTransferEngineConnector (CPU) |
+|---|---|---|
+| KV transfer wall time | ~810 ms | **~14 ms** |
+| RDMA throughput | N/A (TCP) | ~22 GB/s |
+| Speedup | 1x | **58x** |
+
+## Troubleshooting
+
+### Quick Diagnostics
+
+```bash
+# 1. Check RDMA devices and link status
+ibdev2netdev
+# Expected: "mlx5_X port 1 ==> <iface> (Up)"
+# RoCE devices map to Ethernet interfaces (e.g., enp75s0f0)
+# IB devices map to ib0, ib1, etc.
+
+# 2. Check InfiniBand device details
+ibstat
+
+# 3. Verify /dev/infiniband is accessible (required in containers)
+ls /dev/infiniband/
+
+# 4. Check Mooncake installation
+python -c "from mooncake.engine import TransferEngine; print('OK')"
+
+# 5. Check environment variables
+echo "RDMA_DEVICE_NAME=${RDMA_DEVICE_NAME:-<not set>}"
+echo "MC_IB_PCI_RELAXED_ORDERING=${MC_IB_PCI_RELAXED_ORDERING:-<not set>}"
+```
+
+### Common Issues
+
+| Symptom | Cause | Fix |
+|---|---|---|
+| `Failed to modify QP to RTR` | Cross-NIC QP handshake failure (multi-NIC DGX) | Set `device_name` to a single RoCE NIC (e.g., `mlx5_2`) or set `RDMA_DEVICE_NAME` env var |
+| `transport retry counter exceeded` | RDMA path between incompatible NICs | Same as above — restrict to one NIC |
+| `zmq.error.Again: Resource temporarily unavailable` | ZMQ recv timeout (transfer took too long) | Check NIC selection; increase data may need longer timeout |
+| `Mooncake Engine initialization failed` | Missing RDMA drivers or `/dev/infiniband` | Install Mellanox OFED; in Docker add `--device=/dev/infiniband` |
+| `MemoryError` in allocator | Memory pool too small for payload | Increase `memory_pool_size` |
+| GPU transfer slower than CPU | GPU BAR1 bandwidth limitation (PXB/NODE topology) | Use `memory_pool_device: "cpu"` instead of `"cuda"` |
+
+### Multi-NIC Environments (DGX)
+
+On DGX machines with 12+ RDMA NICs, only RoCE NICs (with a bound network
+interface) reliably support loopback. IB-only NICs may fail cross-NIC QP
+handshakes. To identify RoCE NICs:
+
+```bash
+ibdev2netdev | grep -v "ib[0-9]"
+# RoCE devices show Ethernet interface names like enp75s0f0
+```
+
+Then configure the connector:
+```yaml
+device_name: "mlx5_2"  # or set RDMA_DEVICE_NAME=mlx5_2
+```
+
+See the RDMA Test README in tests/distributed/omni_connectors/README.md
+for test-specific setup instructions.
+
+For more details on the underlying engine, refer to the
+[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
@@ -45,7 +45,7 @@ ray start --address=<HEAD_NODE_IP>:6399

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

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

@@ -1,4 +1,4 @@
-# Support Sequence Parallel
+# 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.

@@ -333,13 +333,64 @@ _sp_plan = {

 **Problem:** `SequenceParallelInput(auto_pad=False)` - auto_pad should be True to enable automatic sequence padding.

-**Solution:** In `SequenceParallelInput`, set `auto_pad=True`:
+**Solution:** In `SequenceParallelInput`, set `auto_pad=True` and add attention mask support.
+
+> **Experimental Feature:** `auto_pad=True` is an experimental feature and may be changed in the future. We plan to improve this solution to involve minimal changes to model files. More details are [here](https://github.com/vllm-project/vllm-omni/issues/1324).
+
+**Constraints of auto_pad:**
+
+| Constraint | Description |
+|------------|-------------|
+| **Attention Backend Compatibility** | The attention backends must support `attention_mask`. Currently only `TORCH_SDPA` and `FLASH_ATTN` (default for diffusion models) are supported. |
+| **Ring Attention Limitation** | Ring attention does not support `attention_mask`. Therefore, when using `auto_pad=True`, the combination of Ulysses + Ring attention is not feasible. |
+
+1. Enable `auto_pad=True` for all sequence-dimension inputs in `_sp_plan`:
 ```python
-"blocks.0": {
-    "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True)
+_sp_plan = {
+    "rope": {
+        0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
+        1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
+    },
+    "blocks.0": {
+        "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True)
+    },
+    ...
 }
 ```

+2. Create attention mask dynamically when padding is applied:
+```python
+from vllm_omni.diffusion.forward_context import get_forward_context
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+# In model forward(), before transformer blocks:
+hidden_states_mask = None
+ctx = get_forward_context()
+if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
+    batch_size = hidden_states.shape[0]
+    padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
+    hidden_states_mask = torch.ones(batch_size, padded_seq_len, dtype=torch.bool, device=hidden_states.device)
+    hidden_states_mask[:, ctx.sp_original_seq_len:] = False
+
+# Pass mask to attention layers
+attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask) if hidden_states_mask is not None else None
+output = self.attn(query, key, value, attn_metadata)
+```
+
+**Important Quality Considerations:**
+
+While `auto_pad` enables generation for irregular resolutions, be aware of potential quality impacts:
+
+| Aspect | Impact |
+|--------|--------|
+| **Training Distribution** | Models perform best on aspect ratios seen during training (e.g., 1:1, 16:9, 4:3). Unusual ratios like 700x400 (1.75:1) may produce lower quality results. |
+| **Padding Overhead** | Padded positions consume compute even when masked. For best efficiency, prefer resolutions divisible by `sp_size`. |
+
+**Recommendations for users:**
+- Use standard aspect ratios when possible (e.g., 768x432 for 16:9 instead of 700x400)
+- Ensure post-patch dimensions are divisible by `sp_size` for optimal quality
+- Test generation quality when using unusual resolutions
+
 ### Issue: Inline operations not sharded

 **Symptoms:** Some tensors remain full-sized, not sharded.
@@ -1,4 +1,4 @@
-# Support TeaCache
+# TeaCache

 This section describes how to add TeaCache to a diffusion transformer model. We use the Qwen-Image transformer as the reference implementation.

@@ -1,4 +1,4 @@
-# Support Tensor Parallel
+# 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.

@@ -5,4 +5,5 @@ vLLM-Omni supports the following hardware platforms:
 - [GPU](gpu.md)
    - [NVIDIA CUDA](gpu.md)
    - [AMD ROCm](gpu.md)
+    - [Intel XPU](gpu.md)
 - [NPU](npu.md)
@@ -18,6 +18,10 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr

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

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

 ### Create a new Python environment
@@ -32,11 +36,14 @@ Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.1

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

-
 === "AMD ROCm"

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

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

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

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

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

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

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

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

 === "AMD ROCm"

    --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-docker"
+
+=== "Intel XPU"
+
+    --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-docker"
@@ -20,7 +20,20 @@ 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.16.0 --torch-backend=auto
+# vllm 0.16.0 is still under prerelease
+uv pip install --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
+
+# vllm 0.16.0 may have some bugs for cuda 12.9, here is how we solve them:
+export FLASHINFER_CUDA_TAG="$(python3 -c 'import torch; print((torch.version.cuda or "12.4").replace(".", ""))')"
+
+uv pip install --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}"
+
+uv pip install --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
+uv pip install --upgrade --force-reinstall "numpy==2.2.6"
 ```

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

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

 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
 # --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.14.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.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
+uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
 ```

-The release 0.14.0 of vLLM requires ROCm 7.0 environment.
+The release 0.16.0 of vLLM requires ROCm 7.0 environment.

 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.14.0
+git checkout v0.16.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -0,0 +1,48 @@
+# --8<-- [start:requirements]
+
+- GPU: Validated on Intel® Arc™ B-Series.
+
+# --8<-- [end:requirements]
+# --8<-- [start:set-up-using-python]
+
+vLLM-Omni currently recommends using the Docker image setup steps below.
+
+# --8<-- [start:pre-built-wheels]
+
+# --8<-- [end:pre-built-wheels]
+
+# --8<-- [start:build-wheel-from-source]
+
+# --8<-- [end:build-wheel-from-source]
+
+# --8<-- [start:build-docker]
+
+#### Build docker image
+
+```bash
+DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.xpu -t vllm-omni-xpu --shm-size=4g .
+```
+
+#### Launch the docker image
+
+##### Launch with OpenAI API Server
+
+```
+docker run -it -d --shm-size 10g \
+  --name {container_name} \
+  --net=host \
+  --ipc=host \
+  --privileged \
+  -v /dev/dri/by-path:/dev/dri/by-path \
+  --device /dev/dri:/dev/dri \
+  -v ~/.cache/huggingface:/root/.cache/huggingface \
+  --env "HF_TOKEN=$HF_TOKEN" \
+  vllm-omni-xpu \
+  --model Qwen/Qwen2.5-Omni-3B --port 8091
+```
+
+# --8<-- [end:build-docker]
+
+# --8<-- [start:pre-built-images]
+
+# --8<-- [end:pre-built-images]
@@ -24,6 +24,10 @@ vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a c

 ### Build wheel from source

-=== "NPU"
+=== "NPU release"

-    --8<-- "docs/getting_started/installation/npu/npu.inc.md:installation"
+    --8<-- "docs/getting_started/installation/npu/npu.inc.md:installation-release"
+
+=== "NPU from main"
+
+    --8<-- "docs/getting_started/installation/npu/npu.inc.md:installation-main"
@@ -3,14 +3,11 @@
 For detailed hardware and software requirements, please refer to the [vllm-ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).

 # --8<-- [end:requirements]
-# --8<-- [start:installation]
+# --8<-- [start:installation-release]

 The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built Docker images:

 ```bash
-# Update DEVICE according to your NPUs (/dev/davinci[0-7])
-export DEVICE0=/dev/davinci0
-export DEVICE1=/dev/davinci1
 # Update the vllm-ascend image
 # Atlas A2:
 # export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
@@ -20,8 +17,10 @@ export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
 docker run --rm \
    --name vllm-omni-npu \
    --shm-size=1g \
-    --device $DEVICE0 \
-    --device $DEVICE1 \
+    --device /dev/davinci0 \
+    --device /dev/davinci1 \
+    --device /dev/davinci2 \
+    --device /dev/davinci3 \
    --device /dev/davinci_manager \
    --device /dev/devmm_svm \
    --device /dev/hisi_hdc \
@@ -39,8 +38,9 @@ cd /vllm-workspace
 git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git

 cd vllm-omni
-VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
-# OR pip install -v -e . --no-build-isolation
+sed -i -E 's/^([[:space:]]*)"fa3-fwd==0\.0\.1",/\1# "fa3-fwd==0.0.1",/' pyproject.toml \
+ && sed -i -E 's/\bonnxruntime\b/onnxruntime-cann/g' pyproject.toml
+pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```

@@ -50,11 +50,38 @@ For other installation methods (pip installation, building from source, custom D

 We are keeping [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.

-# --8<-- [end:installation]
+# --8<-- [end:installation-release]
+
+# --8<-- [start:installation-main]
+
+You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) for the status of the latest commit of vLLM-Omni main branch on NPU.)
+
+```bash
+# Pin vLLM version to 0.16.0
+cd /vllm-workspace/vllm
+git pull origin main
+git fetch origin --tags
+git checkout v0.16.0
+
+# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
+cd ../vllm-ascend
+git pull origin main
+git fetch origin --tags
+git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+
+# Install vLLM-Omni from the latest main branch
+cd ../vllm-omni
+git clone https://github.com/vllm-project/vllm-omni.git
+pip install -v -e . --no-build-isolation
+# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
+```
+
+# --8<-- [end:installation-main]

 # --8<-- [start:pre-built-images]

-`vllm-ascend` offers Docker images for deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
+vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-omni](https://quay.io/repository/ascend/vllm-omni?tab=tags) and run it with bash.

 Supported images as following.

@@ -63,9 +90,7 @@ Supported images as following.
 | image-tag | Atlas A2 | Ubuntu |
 | image-tag-a3 | Atlas A3 | Ubuntu |

-vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
-
-Here's an example deployment command that has been verified on 6 x NPUs:
+Here's an example deployment command that has been verified on 4 x NPUs:

 ```bash
 # Atlas A2:
@@ -80,8 +105,6 @@ docker run --rm \
    --device /dev/davinci1 \
    --device /dev/davinci2 \
    --device /dev/davinci3 \
-    --device /dev/davinci4 \
-    --device /dev/davinci5 \
    --device /dev/davinci_manager \
    --device /dev/devmm_svm \
    --device /dev/hisi_hdc \
@@ -19,7 +19,18 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate

 # On CUDA
-uv pip install vllm==0.16.0 --torch-backend=auto
+# vllm 0.16.0 is still under prerelease
+uv pip install --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
+# vllm 0.16.0 may have some bugs for cuda 12.9, here is how we solve them:
+export FLASHINFER_CUDA_TAG="$(python3 -c 'import torch; print((torch.version.cuda or "12.4").replace(".", ""))')"
+uv pip install --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}"
+uv pip install --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
+uv pip install --upgrade --force-reinstall "numpy==2.2.6"
+

 # On ROCm
 uv pip install vllm==0.16.0 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import ast
+import logging
+import re
+import sys
+from argparse import SUPPRESS, Action, ArgumentParser, HelpFormatter, _ArgumentGroup
+from collections.abc import Iterable
+from importlib.machinery import ModuleSpec
+from pathlib import Path
+from typing import Literal
+from unittest.mock import MagicMock
+
+from pydantic_core import core_schema
+
+
+class _FlexibleArgumentParser(ArgumentParser):
+    """Fallback parser for docs when vllm is unavailable.
+
+    Accepts the 'deprecated' kwarg used by vllm CLI and emits warnings
+    if a deprecated argument is actually provided.
+    """
+
+    _deprecated: set[Action] = set()
+    _deprecated_warned: set[str] = set()
+
+    if sys.version_info < (3, 13):
+
+        def parse_known_args(self, args=None, namespace=None):
+            namespace, args = super().parse_known_args(args, namespace)
+            for action in _FlexibleArgumentParser._deprecated:
+                if (
+                    hasattr(namespace, dest := action.dest)
+                    and getattr(namespace, dest) != action.default
+                    and dest not in _FlexibleArgumentParser._deprecated_warned
+                ):
+                    _FlexibleArgumentParser._deprecated_warned.add(dest)
+                    logger.warning("argument '%s' is deprecated", dest)
+            return namespace, args
+
+    def add_argument(self, *args, **kwargs):
+        deprecated = kwargs.pop("deprecated", False)
+        action = super().add_argument(*args, **kwargs)
+        if deprecated:
+            _FlexibleArgumentParser._deprecated.add(action)
+        return action
+
+    class _FlexibleArgumentGroup(_ArgumentGroup):
+        def add_argument(self, *args, **kwargs):
+            deprecated = kwargs.pop("deprecated", False)
+            action = super().add_argument(*args, **kwargs)
+            if deprecated:
+                _FlexibleArgumentParser._deprecated.add(action)
+            return action
+
+    def add_argument_group(self, *args, **kwargs):
+        group = self._FlexibleArgumentGroup(self, *args, **kwargs)
+        self._action_groups.append(group)
+        return group
+
+
+logger = logging.getLogger("mkdocs")
+
+# Define root and doc output directories
+ROOT_DIR = Path(__file__).parent.parent.parent.parent
+ARGPARSE_DOC_DIR = ROOT_DIR / "docs/generated/argparse_omni"
+
+# Ensure the repo root is in sys.path for dynamic imports
+sys.path.insert(0, str(ROOT_DIR))
+
+
+class PydanticMagicMock(MagicMock):
+    """`MagicMock` that's able to generate pydantic-core schemas and avoids _mock_name=None errors."""
+
+    def __init__(self, *args, **kwargs):
+        name = kwargs.pop("name", None)
+        super().__init__(*args, **kwargs)
+        self.__spec__ = ModuleSpec(name, None)
+        self._mock_name = name or "mock"
+
+    def __get_pydantic_core_schema__(self, source_type, handler):
+        return core_schema.any_schema()
+
+
+# --- Static extraction for CLI argument docs ---
+
+
+def extract_omni_serve_subparser_init():
+    """
+    Statically parse vllm_omni/entrypoints/cli/serve.py to extract the subparser_init method
+    and return a callable that adds arguments to a parser. This avoids import and mock issues.
+    """
+    serve_path = ROOT_DIR / "vllm_omni" / "entrypoints" / "cli" / "serve.py"
+    with open(serve_path, encoding="utf-8") as f:
+        source = f.read()
+    tree = ast.parse(source, filename=str(serve_path))
+    # Find class OmniServeCommand
+    for node in tree.body:
+        if isinstance(node, ast.ClassDef) and node.name == "OmniServeCommand":
+            for item in node.body:
+                if isinstance(item, ast.FunctionDef) and item.name == "subparser_init":
+                    # We'll exec this function body in a dummy context
+                    func_code = ast.Module(body=[item], type_ignores=[])
+                    code = compile(func_code, filename=str(serve_path), mode="exec")
+                    # Prepare dummy context
+                    local_vars = {}
+
+                    # Provide a dummy subparsers with add_parser returning our parser
+                    class DummySubparsers:
+                        def add_parser(self, name, **kwargs):
+                            return _FlexibleArgumentParser(prog=name)
+
+                    dummy_subparsers = DummySubparsers()
+                    # Provide globals for exec
+                    # Extract DESCRIPTION from the source file (assume it's a triple-quoted string at the top)
+                    m = re.search(r'DESCRIPTION\s*=\s*([ru]?""".*?""")', source, re.DOTALL)
+                    DESCRIPTION = m.group(1)[3:-3] if m else ""
+                    exec_globals = {
+                        "_FlexibleArgumentParser": _FlexibleArgumentParser,
+                        "FlexibleArgumentParser": _FlexibleArgumentParser,
+                        "make_arg_parser": lambda parser: parser,  # no-op for doc
+                        "VLLM_SUBCMD_PARSER_EPILOG": "",
+                        "logger": logger,
+                        "DummySubparsers": DummySubparsers,
+                        "argparse": __import__("argparse"),
+                        "DESCRIPTION": DESCRIPTION,
+                    }
+                    exec(code, exec_globals, local_vars)
+                    # Get the function
+                    subparser_init = local_vars["subparser_init"]
+
+                    # Return a callable that mimics OmniServeCommand().subparser_init
+                    def parser_factory():
+                        class DummySelf:
+                            name = "serve"
+
+                        return subparser_init(DummySelf(), dummy_subparsers)
+
+                    return parser_factory
+    raise RuntimeError("Could not statically extract OmniServeCommand.subparser_init")
+
+
+OmniServeCommand = type("OmniServeCommand", (), {"subparser_init": staticmethod(extract_omni_serve_subparser_init())})
+
+
+class MarkdownFormatter(HelpFormatter):
+    """Custom formatter that generates markdown for argument groups."""
+
+    def __init__(self, prog: str, starting_heading_level: int = 3):
+        super().__init__(prog, max_help_position=sys.maxsize, width=sys.maxsize)
+
+        self._section_heading_prefix = "#" * starting_heading_level
+        self._argument_heading_prefix = "#" * (starting_heading_level + 1)
+        self._markdown_output = []
+
+    def start_section(self, heading: str):
+        if heading not in {"positional arguments", "options"}:
+            heading_md = f"\n{self._section_heading_prefix} {heading}\n\n"
+            self._markdown_output.append(heading_md)
+
+    def end_section(self):
+        pass
+
+    def add_text(self, text: str):
+        if text:
+            self._markdown_output.append(f"{text.strip()}\n\n")
+
+    def add_usage(self, usage, actions, groups, prefix=None):
+        pass
+
+    def add_arguments(self, actions: Iterable[Action]):
+        for action in actions:
+            if len(action.option_strings) == 0 or "--help" in action.option_strings:
+                continue
+
+            option_strings = f"`{'`, `'.join(action.option_strings)}`"
+            heading_md = f"{self._argument_heading_prefix} {option_strings}\n\n"
+            self._markdown_output.append(heading_md)
+
+            if choices := action.choices:
+                choices = f"`{'`, `'.join(str(c) for c in choices)}`"
+                self._markdown_output.append(f"Possible choices: {choices}\n\n")
+            elif (metavar := action.metavar) and isinstance(metavar, (list, tuple)):
+                metavar = f"`{'`, `'.join(str(m) for m in metavar)}`"
+                self._markdown_output.append(f"Possible choices: {metavar}\n\n")
+
+            if action.help:
+                self._markdown_output.append(f"{action.help}\n\n")
+
+            if (default := action.default) != SUPPRESS:
+                # Make empty string defaults visible
+                if default == "":
+                    default = '""'
+                self._markdown_output.append(f"Default: `{default}`\n\n")
+
+    def format_help(self):
+        """Return the formatted help as markdown."""
+        return "".join(self._markdown_output)
+
+
+# Function to create parser using subparser_init style CLI class
+
+
+def create_parser_subparser_init(subcmd_class):
+    """
+    Create an argparse parser using subparser_init style CLI class, with MarkdownFormatter.
+    Compatible with both AST-extracted static parser_factory and real class instance method.
+    """
+
+    class DummySubparsers:
+        def add_parser(self, name, **kwargs):
+            return _FlexibleArgumentParser(prog=name)
+
+    dummy_subparsers = DummySubparsers()
+    # If subparser_init is a zero-arg static function (AST-extracted), call it directly
+    subparser_init = getattr(subcmd_class, "subparser_init", None)
+    if subparser_init is not None:
+        import inspect
+
+        sig = inspect.signature(subparser_init)
+        if len(sig.parameters) == 0:
+            parser = subparser_init()
+        else:
+            parser = subcmd_class().subparser_init(dummy_subparsers)
+    else:
+        parser = subcmd_class().subparser_init(dummy_subparsers)
+    parser.formatter_class = MarkdownFormatter
+    return parser
+
+
+def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
+    """
+    Entry point for doc generation. Builds doc directory and outputs markdown for each CLI command.
+    """
+    logger.info("Generating vllm-omni argparse documentation")
+    logger.debug("Root directory: %s", ROOT_DIR.resolve())
+    logger.debug("Output directory: %s", ARGPARSE_DOC_DIR.resolve())
+    if not ARGPARSE_DOC_DIR.exists():
+        ARGPARSE_DOC_DIR.mkdir(parents=True)
+
+    # Register all CLI parsers; you can easily add more commands here
+    parsers = {
+        "omni_serve": create_parser_subparser_init(OmniServeCommand),
+        # "another_cmd": create_parser_subparser_init(AnotherCommandClass),
+    }
+
+    for stem, parser in parsers.items():
+        doc_path = ARGPARSE_DOC_DIR / f"{stem}.inc.md"
+        with open(doc_path, "w", encoding="utf-8") as f:
+            f.write(super(type(parser), parser).format_help())
+        logger.info("Argparse generated: %s", doc_path.relative_to(ROOT_DIR))
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    on_startup("build", False)
@@ -36,10 +36,12 @@ th {
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
 |`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
+|`OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` |
 |`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
 |`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
+|`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |


 ## List of Supported Models for NPU
@@ -0,0 +1,255 @@
+# VLLM-Omni PR Reviewer
+
+## Overview
+
+The VLLM-Omni PR Reviewer is an automated code review bot powered by GLM-4.7 AI model. It helps maintain code quality by providing intelligent feedback on pull requests.
+
+## Features
+
+- **Intelligent Code Analysis**: Leverages GLM-4.7 for understanding code context and providing meaningful feedback
+- **Comprehensive Reviews**: Covers code quality, architecture, security, testing, and documentation
+- **Structured Output**: Provides well-formatted reviews with clear sections and actionable suggestions
+- **Rate Limiting**: Built-in cooldown mechanism to prevent excessive API usage
+- **Retry Logic**: Automatic retries with exponential backoff for transient API failures
+- **Defensive Parsing**: Robust validation of API responses to handle malformed data
+- **Cost Control**: Only repository members/collaborators/owners can trigger reviews
+
+## How to Use
+
+### Triggering a Review
+
+To trigger an automated PR review, mention the bot in a PR comment:
+
+```
+@vllm-omni-reviewer please review
+```
+
+Or include in your PR description:
+
+```
+@vllm-omni-reviewer
+```
+
+The bot will automatically review your changes and post a detailed comment.
+
+## What Gets Reviewed
+
+- **vLLM Architecture Compatibility**: Ensures changes align with vLLM's design patterns
+- **Multi-modal Integration**: Reviews audio, vision, and text processing implementations
+- **Performance Implications**: Analyzes impact on inference latency and throughput
+- **Code Quality**: Checks Python best practices, type hints, and documentation
+- **Security Considerations**: Identifies potential security vulnerabilities
+- **Testing Coverage**: Recommends additional test cases when needed
+
+## Review Output
+
+The bot posts a structured review comment with:
+
+- **Overview**: Brief summary of the PR's purpose
+- **Critical Issues (Must Fix)**: Blocking issues that need to be addressed
+- **Important Issues (Should Fix)**: Significant concerns that should be resolved
+- **Minor Issues & Suggestions**: Small improvements and optional suggestions
+- **Positive Aspects**: Highlights well-implemented features
+- **Performance Considerations**: Analysis of performance impact
+- **Testing Recommendations**: Suggestions for additional tests
+- **Overall Assessment**: Final recommendation (Approve/Request Changes/Needs Major Work)
+
+## Rate Limiting and Cooldown
+
+The bot includes a cooldown mechanism to prevent excessive API usage:
+
+- **Default cooldown**: 5 minutes between reviews per PR
+- **Configurable**: Can be adjusted via `PR_REVIEWER_COOLDOWN_MINUTES` environment variable
+- **Smart detection**: Checks for previous bot comments before starting a review
+
+If you trigger a review within the cooldown period, the bot will log a message and skip the review.
+
+## Architecture
+
+```
+┌─────────────────┐
+│  PR Comment     │
+│  @vllm-omni-    │
+│  reviewer       │
+└────────┬────────┘
+         │
+         ▼
+┌─────────────────────────────────┐
+│  GitHub Actions Workflow        │
+│  (.github/workflows/            │
+│   pr-reviewer.yml)              │
+│                                 │
+│  - Python 3.11                  │
+│  - requests==2.31.0             │
+│  - pyyaml==6.0.1                │
+└────────┬────────────────────────┘
+         │
+         ▼
+┌─────────────────────────────────┐
+│  PR Reviewer Script             │
+│  (.github/scripts/              │
+│   pr_reviewer.py)               │
+│                                 │
+│  1. Check cooldown              │
+│  2. Fetch PR details & diff     │
+│  3. Build review prompt         │
+│  4. Call GLM-4.7 API            │
+│     (with retry logic)          │
+│  5. Validate response           │
+│  6. Post review comment         │
+└────────┬────────────────────────┘
+         │
+         ▼
+┌─────────────────────────────────┐
+│  GLM-4.7 API                    │
+│  (open.bigmodel.cn)             │
+└─────────────────────────────────┘
+```
+
+1. **GitHub Actions Workflow** (`.github/workflows/pr-reviewer.yml`): Triggers on @mention
+2. **Python Script** (`.github/scripts/pr_reviewer.py`): Fetches PR data and calls GLM-4.7 API
+3. **GLM-4.7 API**: Provides intelligent code analysis
+
+## Testing
+
+### Testing the PR Reviewer Bot
+
+To test the PR reviewer bot before deploying to production:
+
+1. **Create a test PR** - Make a small, safe change (e.g., documentation update)
+2. **Open the PR** - Create a pull request with a descriptive title
+3. **Trigger the review** - Comment `@vllm-omni-reviewer` on the PR
+4. **Monitor results** - Check the Actions tab for workflow execution logs
+
+### Running Unit Tests
+
+The bot includes comprehensive unit tests that can be run locally:
+
+```bash
+# Run all tests
+pytest .github/tests/test_pr_reviewer.py -v
+
+# Run specific test
+pytest .github/tests/test_pr_reviewer.py::TestCheckTrigger -v
+
+# Run with coverage
+pytest .github/tests/test_pr_reviewer.py --cov=.github/scripts/pr_reviewer.py --cov-report=term-missing
+```
+
+### What to Look For
+
+When testing, verify that:
+- [ ] The workflow triggers on the `@vllm-omni-reviewer` comment
+- [ ] The cooldown mechanism works correctly
+- [ ] The GLM API call completes without errors (with retry if needed)
+- [ ] A review comment is posted to the PR
+- [ ] The review content is meaningful and well-structured
+- [ ] The cost is within the expected range (0.50-5 CNY)
+
+### Safe Test Changes
+
+For testing, consider making these types of safe changes:
+- Documentation updates (like adding this Testing section)
+- Comment improvements
+- README enhancements
+- Non-functional file additions
+
+### Example Test PR
+
+A good test PR might:
+- Update a documentation file
+- Add explanatory comments
+- Improve code formatting
+- Fix a minor typo
+
+These changes are safe to merge if the test is successful and won't affect functionality.
+
+## Troubleshooting
+
+### Bot Doesn't Respond
+
+1. **Check permissions** - Verify you have Owner/Member/Collaborator access
+2. **Check Actions tab** - Look for workflow execution and view logs
+3. **Check cooldown** - If another review was posted recently, wait for the cooldown period
+4. **Check API key** - Ensure `GLM_API_KEY` is configured in repository secrets
+
+### API Errors
+
+If the GLM API call fails:
+- Check the Actions tab for detailed error logs
+- Verify the `GLM_API_KEY` secret is correctly configured
+- Ensure sufficient API quota is available
+- The bot will automatically retry up to 3 times with exponential backoff
+
+### Review Seems Truncated
+
+If the review appears incomplete:
+- Large diffs may be truncated at 100,000 characters
+- Check the logs for truncation warnings
+- Consider breaking large PRs into smaller chunks
+
+## Configuration
+
+### Required Secrets
+
+The following secret must be configured in the repository settings:
+
+- `GLM_API_KEY` - Your GLM (BigModel) API key for accessing the GLM-4.7 API
+
+To add the secret:
+1. Go to repository Settings → Secrets and variables → Actions
+2. Click "New repository secret"
+3. Name: `GLM_API_KEY`
+4. Value: Your GLM API key
+
+### Optional Configuration
+
+The following optional environment variables can be set in the workflow file:
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `GLM_API_URL` | `https://open.bigmodel.cn/api/paas/v4/chat/completions` | GLM API endpoint |
+| `GLM_MODEL` | `glm-4.7` | Model to use for reviews |
+| `PR_REVIEWER_COOLDOWN_MINUTES` | `5` | Cooldown period between reviews |
+| `PR_REVIEWER_MAX_RETRIES` | `3` | Maximum API retry attempts |
+| `PR_REVIEWER_RETRY_DELAY` | `1.0` | Base delay for retry backoff (seconds) |
+| `PR_REVIEWER_MAX_DIFF_SIZE` | `100000` | Maximum diff size before truncation |
+
+### Workflow Customization
+
+The workflow can be customized in `.github/workflows/pr-reviewer.yml`:
+- Change Python version (default: 3.11)
+- Adjust timeout value (default: 10 minutes)
+- Modify trigger conditions
+- Add additional dependencies
+
+## Code Quality
+
+The PR reviewer script follows vllm-omni coding standards:
+
+- **Type hints**: All functions have complete type hints following mypy strict mode
+- **Logging**: Uses Python's logging module for structured logging
+- **Testing**: Comprehensive unit tests with pytest
+- **Pre-commit**: Script is checked by pre-commit hooks (flake8)
+
+## Cost Estimate
+
+| Component | Cost |
+|-----------|------|
+| GitHub Actions (public repo) | Free |
+| GLM API (glm-4.7) | ~0.50-5 CNY per PR (varies by size) |
+| Total (20 PRs/month) | ~10-100 CNY/month (~$2-15 USD) |
+
+## Contributing
+
+To improve the PR reviewer bot:
+
+1. Edit `.github/scripts/pr_reviewer.py` for logic changes
+2. Edit `.github/workflows/pr-reviewer.yml` for workflow changes
+3. Add tests to `.github/tests/test_pr_reviewer.py`
+4. Run `pre-commit run --files .github/scripts/pr_reviewer.py` to check code quality
+5. Test thoroughly with a test PR before deploying to production
+
+## License
+
+This bot is part of the VLLM-Omni project and follows the same license terms.
@@ -0,0 +1,63 @@
+# LoRA (Low-Rank Adaptation) Guide
+
+LoRA (Low-Rank Adaptation) enables fine-tuning diffusion models by adding trainable low-rank matrices to existing model weights. vLLM-Omni currently supports PEFT-style LoRA adapters, allowing you to customize model behavior without modifying the base model weights.
+
+## Overview
+
+LoRA adapters are lightweight, model-specific fine-tuning weights that can be dynamically loaded and applied to diffusion models. vLLM-Omni uses a unified LoRA handling mechanism similar to vLLM with LRU cache management.
+
+## LoRA Adapter Format
+
+LoRA adapters must be in **PEFT (Parameter-Efficient Fine-Tuning)** format. A typical LoRA adapter directory structure:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
+
+The `adapter_config.json` file contains metadata about the LoRA adapter, including:
+- `r`: LoRA rank
+- `lora_alpha`: LoRA alpha scaling factor
+- `target_modules`: List of module names to apply LoRA to
+
+## Quick Start
+
+### Offline Inference
+
+#### Pre-loaded LoRA
+
+Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
+
+```python
+from vllm_omni import Omni
+from vllm_omni.lora.request import LoRARequest
+
+lora_path="/path/to/lora_adapter"
+
+omni = Omni(
+    model="stabilityai/stable-diffusion-3.5-medium",
+    lora_path=lora_path
+)
+
+lora_request = LoRARequest(
+    lora_name="preloaded",
+    lora_int_id=1,
+    lora_path=lora_path
+)
+
+outputs = omni.generate(
+    prompt="A piece of cheesecake",
+    lora_request=lora_request,
+    lora_scale=2.0, # optional arg, default 1.0
+)
+```
+
+!!! note "Server-side Path Requirement"
+    The LoRA adapter path (`local_path`) must be readable on the **server** machine. If your client and server are on different machines, ensure the LoRA adapter is accessible via a shared mount or copied to the server.
+
+
+## See Also
+
+- [Text-to-Image Offline Example](../examples/offline_inference/text_to_image.md#lora) - Complete offline LoRA example
+- [Text-to-Image Online Example](../examples/online_serving/text_to_image.md#lora) - Complete online LoRA example
@@ -65,6 +65,7 @@ The following table shows which models are currently supported by each accelerat
 | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
 | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
+| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ |

 ### VideoGen

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

 > ℹ️ 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.

+## LoRA
+
+This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. Pass `--lora-path` to use a LoRA adapter and optionally `--lora-scale` (default 1.0); omit it to use the base model only.
+
+### Basic usage with LoRA
+
+```bash
+python text_to_image.py \
+  --model Tongyi-MAI/Z-Image-Turbo \
+  --prompt "A piece of cheesecake" \
+  --lora-path /path/to/lora/ \
+  --lora-scale 1.0 \
+  --output output.png
+```
+
+### LoRA parameters
+
+- `--lora-path`: Path to LoRA adapter folder (PEFT format). Loaded at initialization and used for generation.
+- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
+
+### LoRA adapter format
+
+LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
+
 ## Web UI Demo

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

+## LoRA
+
+This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
+
+### Using Python Client with LoRA
+
+```bash
+python openai_chat_client.py \
+  --prompt "A piece of cheesecake" \
+  --lora-path /path/to/lora_adapter \
+  --lora-name my_lora \
+  --lora-scale 1.0 \
+  --output output.png
+```
+
+### Using curl with LoRA (Images API)
+
+The `/v1/images/generations` endpoint supports a `lora` field in the request body:
+
+```bash
+curl -X POST http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "prompt": "A piece of cheesecake",
+    "size": "1024x1024",
+    "seed": 42,
+    "lora": {
+      "name": "my_lora",
+      "local_path": "/path/to/lora_adapter",
+      "scale": 1.0
+    }
+  }' | jq -r '.data[0].b64_json' | base64 -d > output.png
+```
+
+### LoRA Parameters
+
+| Parameter     | Type   | Description                                                      |
+| ------------- | ------ | ---------------------------------------------------------------- |
+| `name`        | str    | LoRA adapter name (optional, defaults to path stem)              |
+| `local_path`  | str    | Server-local path to LoRA adapter folder (PEFT format, required) |
+| `scale`       | float  | LoRA scale factor (default: 1.0)                                 |
+| `int_id`     | int    | LoRA integer ID for caching (optional, derived from path if not provided) |
+
+### LoRA Adapter Format
+
+LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
+
 ## Request Format

 ### Simple Text Generation
@@ -177,6 +177,74 @@ Example configuration for TP=2 on GPUs 0 and 1:
 | `max_inflight`        | `1`     | Maximum inflight requests        |
 | `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB)   |

+## Using Mooncake Connector
+
+[Mooncake](https://github.com/kvcache-ai/Mooncake) is a high-performance distributed KV cache transfer engine that enables efficient cross-node data movement via TCP or RDMA, making it ideal for multi-node disaggregated inference.
+
+By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can switch to the Mooncake connector for better performance on multi-GPU setups and to enable multi-node deployment.
+
+### Prerequisites
+
+Install the Mooncake transfer engine:
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+### Step 1: Start the Mooncake Master
+
+On the **primary node**, start the Mooncake master service (run in a separate terminal or background with `&`):
+
+```bash
+# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
+# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+### Step 2: Run Offline Inference with Mooncake
+
+Use the provided Mooncake stage config [`bagel_multiconnector.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml). Before launching, update the `metadata_server` and `master` addresses in the YAML to match your Mooncake master node's IP (use `127.0.0.1` for single-node testing).
+
+```bash
+cd examples/offline_inference/bagel
+
+# Text to Image with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --prompts "A cute cat" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+
+# Image to Text with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2text \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Describe this image" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+
+# Text to Text with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --prompts "What is the capital of France?" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+```
+
+For more details on the Mooncake connector and multi-node setup, see the [Mooncake Store Connector documentation](../../../docs/design/feature/omni_connectors/mooncake_store_connector.md).
+
+------
+
 ## FAQ

 - If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
@@ -2,6 +2,23 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 """
+Example script for image editing with OmniGen2.
+
+    python image_edit.py \
+        --image input.png \
+        --model "OmniGen2/OmniGen2" \
+        --prompt "Change the background to classroom." \
+        --negative-prompt "(((deformed))), blurry, over saturation, bad anatomy, disfigured, poorly drawn face, mutation, mutated, (extra_limb), (ugly), (poorly drawn hands), fused fingers, messy drawing, broken legs censor, censored, censor_bar" \
+        --num-inference-steps 50 \
+        --seed 0 \
+        --guidance-scale 5.0 \
+        --guidance-scale-2 2.0 \
+        --output outputs/image_edit.png \
+        --num-outputs-per-prompt 2
+
+    Note: For OmniGen2, `guidance_scale` works as `text_guidance_scale`,
+    and `guidance_scale_2` works as `image_guidance_scale`.
+
 Example script for image editing with Qwen-Image-Edit.

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

+## Batched Decoding
+
+The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
+
+```
+python end2end.py --query-type CustomVoice \
+    --txt-prompts benchmark_prompts.txt \
+    --batch-size 4 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
+```
+
+**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
+
 ## Notes

 - The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs.
- Use `--output-dir` (preferred) or `--output-wav` to change the output folder.
+- Use `--output-dir` to change the output folder.
@@ -0,0 +1,12 @@
+Hello, welcome to the voice synthesis benchmark test.
+She said she would be here by noon, but nobody showed up.
+The quick brown fox jumps over the lazy dog near the riverbank.
+I can't believe how beautiful the sunset looks from up here on the mountain.
+Please remember to bring your identification documents to the appointment tomorrow morning.
+Have you ever wondered what it would be like to travel through time and visit ancient civilizations?
+The restaurant on the corner serves the best pasta I have ever tasted in my entire life.
+After the meeting, we should discuss the quarterly results and plan for the next phase.
+Learning a new language takes patience, practice, and a genuine curiosity about other cultures.
+The train leaves at half past seven, so we need to arrive at the station before then.
+Could you please turn down the music a little bit, I'm trying to concentrate on my work.
+It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.
@@ -4,18 +4,21 @@ Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base
 tasks, then runs Omni generation and saves output wav files.
 """

+import logging
 import os
-from typing import NamedTuple
+from typing import Any, NamedTuple

 import soundfile as sf
+import torch

 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

-from vllm import SamplingParams
 from vllm.utils.argparse_utils import FlexibleArgumentParser

 from vllm_omni import Omni

+logger = logging.getLogger(__name__)
+

 class QueryResult(NamedTuple):
    """Container for a prepared Omni request."""
@@ -24,6 +27,44 @@ class QueryResult(NamedTuple):
    model_name: str


+def _estimate_prompt_len(
+    additional_information: dict[str, Any],
+    model_name: str,
+    _cache: dict[str, Any] = {},
+) -> int:
+    """Estimate prompt_token_ids placeholder length for the Talker stage.
+
+    The AR Talker replaces all input embeddings via ``preprocess``, so the
+    placeholder values are irrelevant but the **length** must match the
+    embeddings that ``preprocess`` will produce.
+    """
+    try:
+        from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
+        from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+            Qwen3TTSTalkerForConditionalGeneration,
+        )
+
+        if model_name not in _cache:
+            from transformers import AutoTokenizer
+
+            tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
+            cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
+            _cache[model_name] = (tok, getattr(cfg, "talker_config", None))
+
+        tok, tcfg = _cache[model_name]
+        task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
+        return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
+            additional_information=additional_information,
+            task_type=task_type,
+            tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
+            codec_language_id=getattr(tcfg, "codec_language_id", None),
+            spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
+        )
+    except Exception as exc:
+        logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
+        return 2048
+
+
 def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
    """Build CustomVoice sample inputs.

@@ -34,6 +75,7 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
        QueryResult with Omni inputs and the CustomVoice model path.
    """
    task_type = "CustomVoice"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
    if use_batch_sample:
        texts = ["其实我真的有发现，我是一个特别善于观察别人情绪的人。", "She said she would be here by noon."]
        instructs = ["", "Very happy."]
@@ -41,18 +83,18 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
        speakers = ["Vivian", "Ryan"]
        inputs = []
        for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
-            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+            additional_information = {
+                "task_type": [task_type],
+                "text": [text],
+                "instruct": [instruct],
+                "language": [language],
+                "speaker": [speaker],
+                "max_new_tokens": [2048],
+            }
            inputs.append(
                {
-                    "prompt": prompt,
-                    "additional_information": {
-                        "task_type": [task_type],
-                        "text": [text],
-                        "instruct": [instruct],
-                        "language": [language],
-                        "speaker": [speaker],
-                        "max_new_tokens": [2048],
-                    },
+                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+                    "additional_information": additional_information,
                }
            )
    else:
@@ -60,21 +102,21 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
        language = "Chinese"
        speaker = "Vivian"
        instruct = "用特别愤怒的语气说"
-        prompts = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+        additional_information = {
+            "task_type": [task_type],
+            "text": [text],
+            "language": [language],
+            "speaker": [speaker],
+            "instruct": [instruct],
+            "max_new_tokens": [2048],
+        }
        inputs = {
-            "prompt": prompts,
-            "additional_information": {
-                "task_type": [task_type],
-                "text": [text],
-                "language": [language],
-                "speaker": [speaker],
-                "instruct": [instruct],
-                "max_new_tokens": [2048],
-            },
+            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+            "additional_information": additional_information,
        }
    return QueryResult(
        inputs=inputs,
-        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+        model_name=model_name,
    )


@@ -88,6 +130,7 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
        QueryResult with Omni inputs and the VoiceDesign model path.
    """
    task_type = "VoiceDesign"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign"
    if use_batch_sample:
        texts = [
            "哥哥，你回来啦，人家等了你好久好久了，要抱抱！",
@@ -100,39 +143,39 @@ def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
        languages = ["Chinese", "English"]
        inputs = []
        for text, instruct, language in zip(texts, instructs, languages):
-            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
-            inputs.append(
-                {
-                    "prompt": prompt,
-                    "additional_information": {
-                        "task_type": [task_type],
-                        "text": [text],
-                        "language": [language],
-                        "instruct": [instruct],
-                        "max_new_tokens": [2048],
-                        "non_streaming_mode": [True],
-                    },
-                }
-            )
-    else:
-        text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
-        instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
-        language = "Chinese"
-        prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
-        inputs = {
-            "prompt": prompt,
-            "additional_information": {
+            additional_information = {
                "task_type": [task_type],
                "text": [text],
                "language": [language],
                "instruct": [instruct],
                "max_new_tokens": [2048],
                "non_streaming_mode": [True],
-            },
+            }
+            inputs.append(
+                {
+                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+                    "additional_information": additional_information,
+                }
+            )
+    else:
+        text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
+        instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
+        language = "Chinese"
+        additional_information = {
+            "task_type": [task_type],
+            "text": [text],
+            "language": [language],
+            "instruct": [instruct],
+            "max_new_tokens": [2048],
+            "non_streaming_mode": [True],
+        }
+        inputs = {
+            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+            "additional_information": additional_information,
        }
    return QueryResult(
        inputs=inputs,
-        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
+        model_name=model_name,
    )


@@ -147,6 +190,7 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que
        QueryResult with Omni inputs and the Base model path.
    """
    task_type = "Base"
+    model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-Base"
    ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
    ref_audio_single = ref_audio_path_1
    ref_text_single = (
@@ -163,38 +207,38 @@ def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> Que
        syn_lang_batch = ["Chinese", "English"]
        inputs = []
        for text, language in zip(syn_text_batch, syn_lang_batch):
-            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
-            inputs.append(
-                {
-                    "prompt": prompt,
-                    "additional_information": {
-                        "task_type": [task_type],
-                        "ref_audio": [ref_audio_single],
-                        "ref_text": [ref_text_single],
-                        "text": [text],
-                        "language": [language],
-                        "x_vector_only_mode": [x_vector_only_mode],
-                        "max_new_tokens": [2048],
-                    },
-                }
-            )
-    else:
-        prompt = f"<|im_start|>assistant\n{syn_text_single}<|im_end|>\n<|im_start|>assistant\n"
-        inputs = {
-            "prompt": prompt,
-            "additional_information": {
+            additional_information = {
                "task_type": [task_type],
                "ref_audio": [ref_audio_single],
                "ref_text": [ref_text_single],
-                "text": [syn_text_single],
-                "language": [syn_lang_single],
+                "text": [text],
+                "language": [language],
                "x_vector_only_mode": [x_vector_only_mode],
                "max_new_tokens": [2048],
-            },
+            }
+            inputs.append(
+                {
+                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+                    "additional_information": additional_information,
+                }
+            )
+    else:
+        additional_information = {
+            "task_type": [task_type],
+            "ref_audio": [ref_audio_single],
+            "ref_text": [ref_text_single],
+            "text": [syn_text_single],
+            "language": [syn_lang_single],
+            "x_vector_only_mode": [x_vector_only_mode],
+            "max_new_tokens": [2048],
+        }
+        inputs = {
+            "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+            "additional_information": additional_information,
        }
    return QueryResult(
        inputs=inputs,
-        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
+        model_name=model_name,
    )


@@ -204,6 +248,13 @@ def main(args):
    Args:
        args: Parsed CLI args from parse_args().
    """
+    if args.batch_size < 1 or (args.batch_size & (args.batch_size - 1)) != 0:
+        raise ValueError(
+            f"--batch-size must be a power of two (got {args.batch_size}); "
+            "non-power-of-two values do not align with CUDA graph capture sizes "
+            "of Code2Wav."
+        )
+
    query_func = query_map[args.query_type]
    if args.query_type in {"CustomVoice", "VoiceDesign"}:
        query_result = query_func(use_batch_sample=args.use_batch_sample)
@@ -216,6 +267,33 @@ def main(args):
        query_result = query_func()

    model_name = query_result.model_name
+
+    # Load prompts from text file if provided.
+    # Use the default query as a template so task-specific fields
+    # (e.g. ref_audio for Base) are preserved; only override text.
+    if args.txt_prompts:
+        with open(args.txt_prompts) as f:
+            lines = [line.strip() for line in f if line.strip()]
+        if not lines:
+            raise ValueError(f"No valid prompts found in {args.txt_prompts}")
+        template = query_result.inputs
+        if isinstance(template, list):
+            template = template[0]
+        template_info = template["additional_information"]
+        inputs = []
+        for text in lines:
+            additional_information = {**template_info, "text": [text]}
+            inputs.append(
+                {
+                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+                    "additional_information": additional_information,
+                }
+            )
+    else:
+        inputs = query_result.inputs
+        if not isinstance(inputs, list):
+            inputs = [inputs]
+
    omni = Omni(
        model=model_name,
        stage_configs_path=args.stage_configs_path,
@@ -223,40 +301,35 @@ def main(args):
        stage_init_timeout=args.stage_init_timeout,
    )

-    sampling_params = SamplingParams(
-        temperature=0.9,
-        top_p=1.0,
-        top_k=50,
-        max_tokens=2048,
-        seed=42,
-        detokenize=False,
-        repetition_penalty=1.05,
-    )
-
-    sampling_params_list = [
-        sampling_params,
-    ]
-
-    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    output_dir = args.output_dir
    os.makedirs(output_dir, exist_ok=True)

-    omni_generator = omni.generate(query_result.inputs, sampling_params_list)
-    for stage_outputs in omni_generator:
-        for output in stage_outputs.request_output:
-            request_id = output.request_id
-            audio_tensor = output.outputs[0].multimodal_output["audio"]
-            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-            audio_samplerate = output.outputs[0].multimodal_output["sr"].item()
-            # Convert to numpy array and ensure correct format
-            audio_numpy = audio_tensor.float().detach().cpu().numpy()
+    batch_size = args.batch_size
+    for batch_start in range(0, len(inputs), batch_size):
+        batch = inputs[batch_start : batch_start + batch_size]
+        omni_generator = omni.generate(batch, sampling_params_list=None)
+        for stage_outputs in omni_generator:
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+                audio_data = output.outputs[0].multimodal_output["audio"]
+                # async_chunk mode returns a list of chunks; concatenate them.
+                if isinstance(audio_data, list):
+                    audio_tensor = torch.cat(audio_data, dim=-1)
+                else:
+                    audio_tensor = audio_data
+                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+                sr_val = output.outputs[0].multimodal_output["sr"]
+                audio_samplerate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val[-1])
+                # Convert to numpy array and ensure correct format
+                audio_numpy = audio_tensor.float().detach().cpu().numpy()

-            # Ensure audio is 1D (flatten if needed)
-            if audio_numpy.ndim > 1:
-                audio_numpy = audio_numpy.flatten()
+                # Ensure audio is 1D (flatten if needed)
+                if audio_numpy.ndim > 1:
+                    audio_numpy = audio_numpy.flatten()

-            # Save audio file with explicit WAV format
-            sf.write(output_wav, audio_numpy, samplerate=audio_samplerate, format="WAV")
-            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+                # Save audio file with explicit WAV format
+                sf.write(output_wav, audio_numpy, samplerate=audio_samplerate, format="WAV")
+                print(f"Request ID: {request_id}, Saved audio to {output_wav}")


 def parse_args():
@@ -305,9 +378,9 @@ def parse_args():
        help="Threshold for using shared memory in bytes (default: 65536)",
    )
    parser.add_argument(
-        "--output-wav",
+        "--output-dir",
        default="output_audio",
-        help="[Deprecated] Output wav directory (use --output-dir).",
+        help="Output directory for generated wav files (default: output_audio).",
    )
    parser.add_argument(
        "--num-prompts",
@@ -365,6 +438,12 @@ def parse_args():
        choices=["icl", "xvec_only"],
        help="Mode tag for Base query x_vector_only_mode (default: icl).",
    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1,
+        help="Number of prompts per batch (default: 1, sequential).",
+    )

    return parser.parse_args()

@@ -1,6 +1,6 @@
 # Text-To-Image

-This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` using vLLM-Omni:
+This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` `stepfun-ai/NextStep-1.1` using vLLM-Omni, note that NextStep-1.1 has different architecture so we treat it differently regarding running arguments and pipeline.

 - `text_to_image.py`: command-line script for single image generation with advanced options.
 - `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
@@ -74,6 +74,8 @@ if __name__ == "__main__":

 ## Local CLI Usage

+### Qwen/Tongyi Models
+
 ```bash
 python text_to_image.py \
  --model Tongyi-MAI/Z-Image-Turbo \
@@ -87,7 +89,26 @@ python text_to_image.py \
  --output outputs/coffee.png
 ```

-Key arguments:
+### NextStep Models
+
+NextStep-1.1 can have extra arguments
+```bash
+python text_to_image.py \
+  --model stepfun-ai/NextStep-1.1 \
+  --prompt "A baby panda wearing an Iron Man mask, holding a board with 'NextStep-1' written on it" \
+  --height 512 \
+  --width 512 \
+  --num-inference-steps 28 \
+  --guidance-scale 7.5 \
+  --guidance-scale-2 1.0 \
+  --cfg-schedule constant \
+  --output nextstep_output.png \
+  --seed 42
+```
+
+### Key Arguments
+
+**Common arguments:**

 - `--prompt`: text description (string).
 - `--seed`: integer seed for deterministic sampling.
@@ -98,13 +119,50 @@ Key arguments:
 - `--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).
+- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion_acceleration.md#using-cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+- `--guidance-scale`: classifier-free guidance scale.
+
+**NextStep-1.1 specific:**
+- `--guidance-scale-2`: secondary guidance scale, e.g. image-level CFG (default: 1.0).
+- `--timesteps-shift`: timesteps shift parameter for sampling (default: 1.0).
+- `--cfg-schedule`: CFG schedule type, "constant" or "linear" (default: "constant").
+- `--use-norm`: apply layer normalization to sampled tokens.

 > ℹ️ 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.

+## LoRA
+
+This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. Pass `--lora-path` to use a LoRA adapter and optionally `--lora-scale` (default 1.0); omit it to use the base model only.
+
+### Basic usage with LoRA
+
+```bash
+python text_to_image.py \
+  --model Tongyi-MAI/Z-Image-Turbo \
+  --prompt "A piece of cheesecake" \
+  --lora-path /path/to/lora/ \
+  --lora-scale 1.0 \
+  --output output.png
+```
+
+### LoRA parameters
+
+- `--lora-path`: Path to LoRA adapter folder (PEFT format). Loaded at initialization and used for generation.
+- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
+
+### LoRA adapter format
+
+LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
+
 ## Web UI Demo

 Launch the gradio demo:
@@ -12,17 +12,32 @@ import torch
 from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.lora.request import LoRARequest
+from vllm_omni.lora.utils import stable_lora_int_id
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform


+def is_nextstep_model(model_name: str) -> bool:
+    """Check if the model is a NextStep model by reading its config."""
+    from vllm.transformers_utils.config import get_hf_file_to_dict
+
+    try:
+        cfg = get_hf_file_to_dict("config.json", model_name)
+        if cfg and cfg.get("model_type") == "nextstep":
+            return True
+    except Exception:
+        pass
+    return False
+
+
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate an image with Qwen-Image.")
+    parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
    parser.add_argument(
        "--model",
        default="Qwen/Qwen-Image",
        help="Diffusion model name or local path. Supported models: "
-        "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512",
+        "Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1",
    )
    parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
    parser.add_argument(
@@ -147,22 +162,71 @@ def parse_args() -> argparse.Namespace:
        default=1,
        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
    )
+    parser.add_argument(
+        "--lora-path",
+        type=str,
+        default=None,
+        help="Path to LoRA adapter folder (PEFT format). Loaded at initialization and used for generation.",
+    )
+    parser.add_argument(
+        "--lora-scale",
+        type=float,
+        default=1.0,
+        help="Scale factor for LoRA weights (default: 1.0).",
+    )
+    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(
        "--vae-patch-parallel-size",
        type=int,
        default=1,
        help="Number of ranks used for VAE patch/tile parallelism (decode/encode).",
    )
+    # NextStep-1.1 specific arguments
+    parser.add_argument(
+        "--guidance-scale-2",
+        type=float,
+        default=1.0,
+        help="Secondary guidance scale (e.g. image-level CFG for NextStep-1.1).",
+    )
+    parser.add_argument(
+        "--timesteps-shift",
+        type=float,
+        default=1.0,
+        help="[NextStep-1.1 only] Timesteps shift parameter for sampling.",
+    )
+    parser.add_argument(
+        "--cfg-schedule",
+        type=str,
+        default="constant",
+        choices=["constant", "linear"],
+        help="[NextStep-1.1 only] CFG schedule type.",
+    )
+    parser.add_argument(
+        "--use-norm",
+        action="store_true",
+        help="[NextStep-1.1 only] Apply layer normalization to sampled tokens.",
+    )
    return parser.parse_args()


 def main():
    args = parse_args()
    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
+    use_nextstep = is_nextstep_model(args.model)

-    # Configure cache based on backend type
    cache_config = None
-    if args.cache_backend == "cache_dit":
+    cache_backend = args.cache_backend
+
+    if cache_backend == "cache_dit":
        # cache-dit configuration: Hybrid DBCache + SCM + TaylorSeer
        # All parameters marked with [cache-dit only] in DiffusionCacheConfig
        cache_config = {
@@ -179,7 +243,7 @@ def main():
            "scm_steps_mask_policy": None,  # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
            "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
        }
-    elif args.cache_backend == "tea_cache":
+    elif cache_backend == "tea_cache":
        # TeaCache configuration
        # All parameters marked with [tea_cache only] in DiffusionCacheConfig
        cache_config = {
@@ -201,6 +265,12 @@ def main():
    # Check if profiling is requested via environment variable
    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))

+    # Prepare LoRA kwargs for Omni initialization
+    lora_args: dict[str, Any] = {}
+    if args.lora_path:
+        lora_args["lora_path"] = args.lora_path
+        print(f"Using LoRA from: {args.lora_path}")
+
    # Build quantization kwargs: use quantization_config dict when
    # ignored_layers is specified so the list flows through OmniDiffusionConfig
    quant_kwargs: dict[str, Any] = {}
@@ -213,19 +283,25 @@ def main():
    elif args.quantization:
        quant_kwargs["quantization"] = args.quantization

-    omni = Omni(
-        model=args.model,
-        enable_layerwise_offload=args.enable_layerwise_offload,
-        vae_use_slicing=args.vae_use_slicing,
-        vae_use_tiling=args.vae_use_tiling,
-        cache_backend=args.cache_backend,
-        cache_config=cache_config,
-        enable_cache_dit_summary=args.enable_cache_dit_summary,
-        parallel_config=parallel_config,
-        enforce_eager=args.enforce_eager,
-        enable_cpu_offload=args.enable_cpu_offload,
+    omni_kwargs = {
+        "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,
+        "cache_config": cache_config,
+        "enable_cache_dit_summary": args.enable_cache_dit_summary,
+        "parallel_config": parallel_config,
+        "enforce_eager": args.enforce_eager,
+        "enable_cpu_offload": args.enable_cpu_offload,
+        **lora_args,
        **quant_kwargs,
-    )
+    }
+    if use_nextstep:
+        # NextStep-1.1 requires explicit pipeline class
+        omni_kwargs["model_class_name"] = "NextStep11Pipeline"
+    omni = Omni(**omni_kwargs)

    if profiler_enabled:
        print("[Profiler] Starting profiling...")
@@ -236,7 +312,7 @@ def main():
    print("Generation Configuration:")
    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"  Cache backend: {cache_backend if 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}")
@@ -247,9 +323,32 @@ def main():
    )
    print(f"  CPU offload: {args.enable_cpu_offload}")
    print(f"  Image size: {args.width}x{args.height}")
+    if args.lora_path:
+        print(f"  LoRA: scale={args.lora_scale}")
    print(f"{'=' * 60}\n")

+    # Build LoRA request when --lora-path is set
+    lora_request = None
+    if args.lora_path:
+        lora_request_id = stable_lora_int_id(args.lora_path)
+        lora_request = LoRARequest(
+            lora_name=Path(args.lora_path).stem,
+            lora_int_id=lora_request_id,
+            lora_path=args.lora_path,
+        )
+
    generation_start = time.perf_counter()
+
+    extra_args = {
+        "timesteps_shift": args.timesteps_shift,
+        "cfg_schedule": args.cfg_schedule,
+        "use_norm": args.use_norm,
+    }
+
+    if lora_request:
+        extra_args["lora_request"] = lora_request
+        extra_args["lora_scale"] = args.lora_scale
+
    outputs = omni.generate(
        {
            "prompt": args.prompt,
@@ -261,10 +360,13 @@ def main():
            generator=generator,
            true_cfg_scale=args.cfg_scale,
            guidance_scale=args.guidance_scale,
+            guidance_scale_2=args.guidance_scale_2,
            num_inference_steps=args.num_inference_steps,
            num_outputs_per_prompt=args.num_images_per_prompt,
+            extra_args=extra_args,
        ),
    )
+
    generation_end = time.perf_counter()
    generation_time = generation_end - generation_start

@@ -45,6 +45,129 @@ For larger models or multi-GPU environments, you can enable Tensor Parallelism (
 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/your/custom_bagel.yaml
 ```

+#### Using Mooncake Connector
+
+By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can use the [Mooncake](https://github.com/kvcache-ai/Mooncake) connector to transfer KV cache between stages, which also enables multi-node deployment.
+
+**1. Install Mooncake**
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+**2. Start Mooncake Master** on the primary node:
+
+```bash
+# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
+# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+**3. Launch the server** with the Mooncake stage config:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+```
+
+> **Note**: Before launching, edit [`bagel_multiconnector.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. For single-node testing, `127.0.0.1` works.
+
+The client-side usage is identical to the default setup -- the Mooncake connector is transparent to the API. See the requests section below.
+
+For more details on the Mooncake connector configuration, see the [Mooncake Store Connector documentation](../../../docs/design/feature/omni_connectors/mooncake_store_connector.md).
+
+#### Multi-Node Deployment
+
+You can deploy each stage on a **separate node** for better resource utilization. In this example, the orchestrator (Stage 0 / Thinker) and Stage 1 (DiT) run on different machines, connected via Mooncake.
+
+Replace `<ORCHESTRATOR_IP>` below with the actual IP address of your orchestrator node (e.g., `10.244.227.244`).
+
+> [!WARNING]
+> **Before launching**, edit [`bagel_multiconnector.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. Mismatched addresses will cause silent connection failures.
+
+**1. Start Mooncake Master** (on the orchestrator node):
+
+```bash
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=<ORCHESTRATOR_IP> \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003
+```
+
+**2. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
+    --port 8000 \ # API server port for client requests
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
+    --stage-id 0 \
+    -oma <ORCHESTRATOR_IP> \
+    -omp 8091
+```
+
+**3. Launch Stage 1 (DiT)** on the remote node in headless mode:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
+    --stage-id 1 \
+    --headless \
+    -oma <ORCHESTRATOR_IP> \
+    -omp 8091
+```
+
+**Mooncake Master arguments:**
+
+| Argument | Description |
+| :------- | :---------- |
+| `--rpc_port` | Mooncake RPC port for control-plane coordination between stages |
+| `--enable_http_metadata_server` | Enable the HTTP metadata server for service discovery |
+| `--http_metadata_server_host` | IP address to bind the metadata server (use the orchestrator node's IP) |
+| `--http_metadata_server_port` | Port for the HTTP metadata server |
+| `--metrics_port` | Port for Prometheus-compatible metrics endpoint |
+
+**vllm serve arguments:**
+
+| Argument | Description |
+| :------- | :---------- |
+| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) |
+| `--headless` | Run without the API server (worker-only mode) |
+| `-oma` | Orchestrator master address |
+| `-omp` | Orchestrator master port for Stage 1 to connect to Stage 0 for task coordination |
+
+> [!IMPORTANT]
+> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless).
+> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior.
+
+**Network Requirements**
+
+All nodes must have network connectivity to each other. Ensure the following ports are open **between all participating nodes**:
+
+| Port | Protocol | Service | Direction |
+| :--- | :------- | :------ | :-------- |
+| 50051 | TCP | Mooncake Master RPC | Worker → Orchestrator |
+| 8080 | TCP | Mooncake HTTP Metadata Server | Worker → Orchestrator |
+| 8091 | TCP | Orchestrator Master (`-omp`) | Worker → Orchestrator |
+| 8000 | TCP | API Server (`--port`) | Client → Orchestrator |
+| 9003 | TCP | Metrics (optional) | Monitoring → Orchestrator |
+
+> **Tip**: If nodes are behind a firewall or in different VPCs/security groups, make sure the above ports are allowed in ingress/egress rules. All nodes should be reachable via their IP addresses (no NAT). Using nodes on the same subnet or VPC is recommended to minimize latency for Mooncake KV cache transfers.
+
 ### Send Multi-modal Request

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

 SEED = 42

+DEFAULT_SAMPLING_PARAMS = {
+    "thinker": {
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "top_k": -1,
+        "max_tokens": 2048,
+        "seed": SEED,
+        "detokenize": True,
+        "repetition_penalty": 1.1,
+    },
+    "talker": {
+        "temperature": 0.9,
+        "top_p": 0.8,
+        "top_k": 40,
+        "max_tokens": 2048,
+        "seed": SEED,
+        "detokenize": True,
+        "repetition_penalty": 1.05,
+        "stop_token_ids": [8294],
+    },
+    "code2wav": {
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "top_k": -1,
+        "max_tokens": 2048,
+        "seed": SEED,
+        "detokenize": True,
+        "repetition_penalty": 1.1,
+    },
+}
+
 SUPPORTED_MODELS: dict[str, dict[str, Any]] = {
+    "Qwen/Qwen2.5-Omni-3B": {
+        "sampling_params": DEFAULT_SAMPLING_PARAMS,
+    },
    "Qwen/Qwen2.5-Omni-7B": {
-        "sampling_params": {
-            "thinker": {
-                "temperature": 0.0,
-                "top_p": 1.0,
-                "top_k": -1,
-                "max_tokens": 2048,
-                "seed": SEED,
-                "detokenize": True,
-                "repetition_penalty": 1.1,
-            },
-            "talker": {
-                "temperature": 0.9,
-                "top_p": 0.8,
-                "top_k": 40,
-                "max_tokens": 2048,
-                "seed": SEED,
-                "detokenize": True,
-                "repetition_penalty": 1.05,
-                "stop_token_ids": [8294],
-            },
-            "code2wav": {
-                "temperature": 0.0,
-                "top_p": 1.0,
-                "top_k": -1,
-                "max_tokens": 2048,
-                "seed": SEED,
-                "detokenize": True,
-                "repetition_penalty": 1.1,
-            },
-        },
+        "sampling_params": DEFAULT_SAMPLING_PARAMS,
    },
 }
 # Ensure deterministic behavior across runs.
@@ -90,6 +90,25 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 curl http://localhost:8091/v1/audio/voices
 ```

+## Server Configuration
+
+### TTS-Specific CLI Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `--tts-max-instructions-length` | int | 500 | Maximum character length for voice style instructions |
+
+**Example with custom instructions length:**
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
+    --omni \
+    --port 8000 \
+    --tts-max-instructions-length 1000
+```
+
+This is useful when you need longer voice style descriptions, especially for the VoiceDesign task type.
+
 ## API Reference

 ### Endpoint
@@ -58,6 +58,59 @@ python gradio_demo.py
 # Visit http://localhost:7860
 ```

+## LoRA
+
+This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
+
+### Using Python Client with LoRA
+
+```bash
+python openai_chat_client.py \
+  --prompt "A piece of cheesecake" \
+  --lora-path /path/to/lora_adapter \
+  --lora-name my_lora \
+  --lora-scale 1.0 \
+  --output output.png
+```
+
+### Using curl with LoRA (Images API)
+
+The `/v1/images/generations` endpoint supports a `lora` field in the request body:
+
+```bash
+curl -X POST http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "prompt": "A piece of cheesecake",
+    "size": "1024x1024",
+    "seed": 42,
+    "lora": {
+      "name": "my_lora",
+      "local_path": "/path/to/lora_adapter",
+      "scale": 1.0
+    }
+  }' | jq -r '.data[0].b64_json' | base64 -d > output.png
+```
+
+### LoRA Parameters
+
+| Parameter     | Type   | Description                                                      |
+| ------------- | ------ | ---------------------------------------------------------------- |
+| `name`        | str    | LoRA adapter name (optional, defaults to path stem)              |
+| `local_path`  | str    | Server-local path to LoRA adapter folder (PEFT format, required) |
+| `scale`       | float  | LoRA scale factor (default: 1.0)                                 |
+| `int_id`     | int    | LoRA integer ID for caching (optional, derived from path if not provided) |
+
+### LoRA Adapter Format
+
+LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
+
 ## Request Format

 ### Simple Text Generation
@@ -24,6 +24,10 @@ def generate_image(
    seed: int | None = None,
    negative_prompt: str | None = None,
    num_outputs_per_prompt: int = 1,
+    lora_path: str | None = None,
+    lora_name: str | None = None,
+    lora_scale: float | None = None,
+    lora_int_id: int | None = None,
 ) -> bytes | None:
    """Generate an image using the images generation API.

@@ -37,6 +41,10 @@ def generate_image(
        seed: Random seed
        negative_prompt: Negative prompt
        num_outputs_per_prompt: Number of images to generate
+        lora_path: Server-local LoRA adapter folder path (PEFT format)
+        lora_name: LoRA name (optional, defaults to path stem)
+        lora_scale: LoRA scale factor (default: 1.0)
+        lora_int_id: LoRA integer ID (optional, derived from path if not provided)

    Returns:
        Image bytes or None if failed
@@ -63,6 +71,18 @@ def generate_image(
    if seed is not None:
        payload["seed"] = seed

+    # Add LoRA if provided
+    if lora_path:
+        lora_body: dict = {
+            "local_path": lora_path,
+            "name": lora_name or Path(lora_path).stem,
+        }
+        if lora_scale is not None:
+            lora_body["scale"] = float(lora_scale)
+        if lora_int_id is not None:
+            lora_body["int_id"] = int(lora_int_id)
+        payload["lora"] = lora_body
+
    try:
        response = requests.post(
            f"{server_url}/v1/images/generations",
@@ -99,6 +119,16 @@ def main():
    parser.add_argument("--seed", type=int, default=0, help="Random seed")
    parser.add_argument("--negative", help="Negative prompt")

+    parser.add_argument("--lora-path", default=None, help="Server-local LoRA adapter folder (PEFT format)")
+    parser.add_argument("--lora-name", default=None, help="LoRA name (optional)")
+    parser.add_argument("--lora-scale", type=float, default=1.0, help="LoRA scale")
+    parser.add_argument(
+        "--lora-int-id",
+        type=int,
+        default=None,
+        help="LoRA integer id (cache key). If omitted, the server derives a stable id from lora_path.",
+    )
+
    args = parser.parse_args()

    print(f"Generating image for: {args.prompt}")
@@ -112,6 +142,10 @@ def main():
        true_cfg_scale=args.cfg_scale,
        seed=args.seed,
        negative_prompt=args.negative,
+        lora_path=args.lora_path,
+        lora_name=args.lora_name,
+        lora_scale=args.lora_scale if args.lora_path else None,
+        lora_int_id=args.lora_int_id if args.lora_path else None,
    )

    if image_bytes:
@@ -59,6 +59,7 @@ hooks:
  - docs/mkdocs/hooks/generate_api_readme.py
  - docs/mkdocs/hooks/url_schemes.py
  - docs/mkdocs/hooks/generate_examples.py
+  - docs/mkdocs/hooks/generate_argparse.py

 # Exclude include files from navigation warnings
 exclude_docs: |
@@ -41,6 +41,7 @@ dev = [
    "pytest>=7.0.0",
    "pytest-asyncio>=0.21.0",
    "pytest-cov>=4.0.0",
+    "pytest-mock>=3.10.0",
    "mypy==1.11.1",
    "pre-commit==4.0.1",
    "openai-whisper>=20250625",
@@ -14,3 +14,4 @@ imageio[ffmpeg]>=2.37.2
 sox>=1.5.0
 prettytable>=3.8.0
 aenum==3.1.16
+pyzmq>=25.0.0
@@ -7,9 +7,9 @@ Unit tests for patch.py

 import asyncio
 import json
-from unittest.mock import AsyncMock, MagicMock

 import pytest
+from pytest_mock import MockerFixture
 from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput

 from vllm_omni.benchmarks.patch.patch import MixRequestFuncOutput, async_request_openai_chat_omni_completions
@@ -51,7 +51,7 @@ def create_sse_chunk(data_dict):


@pytest.mark.asyncio
-async def test_output_tokens_assigned_with_metrics():
+async def test_output_tokens_assigned_with_metrics(mocker: MockerFixture):
    """Test that output.output_tokens is assigned when metrics are present"""
    # Arrange
    request_input = RequestFuncInput(
@@ -82,8 +82,8 @@ async def test_output_tokens_assigned_with_metrics():
    ]

    mock_response = MockResponse(200, chunks)
-    mock_session = AsyncMock()
-    mock_session.post = MagicMock(return_value=mock_response)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)

    # Act
    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -95,7 +95,7 @@ async def test_output_tokens_assigned_with_metrics():


@pytest.mark.asyncio
-async def test_output_tokens_not_assigned_without_metrics():
+async def test_output_tokens_not_assigned_without_metrics(mocker: MockerFixture):
    """Test that output.output_tokens defaults to 0 when no metrics present"""
    # Arrange
    request_input = RequestFuncInput(
@@ -125,8 +125,8 @@ async def test_output_tokens_not_assigned_without_metrics():
    ]

    mock_response = MockResponse(200, chunks)
-    mock_session = AsyncMock()
-    mock_session.post = MagicMock(return_value=mock_response)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)

    # Act
    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -138,7 +138,7 @@ async def test_output_tokens_not_assigned_without_metrics():


@pytest.mark.asyncio
-async def test_output_tokens_assigned_multiple_metrics():
+async def test_output_tokens_assigned_multiple_metrics(mocker: MockerFixture):
    """Test that output.output_tokens is updated with the latest metrics value"""
    # Arrange
    request_input = RequestFuncInput(
@@ -177,8 +177,8 @@ async def test_output_tokens_assigned_multiple_metrics():
    ]

    mock_response = MockResponse(200, chunks)
-    mock_session = AsyncMock()
-    mock_session.post = MagicMock(return_value=mock_response)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)

    # Act
    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -190,7 +190,7 @@ async def test_output_tokens_assigned_multiple_metrics():


@pytest.mark.asyncio
-async def test_output_tokens_with_audio_and_text():
+async def test_output_tokens_with_audio_and_text(mocker: MockerFixture):
    """Test output_tokens assignment in mixed audio and text response"""
    # Arrange
    request_input = RequestFuncInput(
@@ -226,8 +226,8 @@ async def test_output_tokens_with_audio_and_text():
    ]

    mock_response = MockResponse(200, chunks)
-    mock_session = AsyncMock()
-    mock_session.post = MagicMock(return_value=mock_response)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)

    # Act
    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -238,7 +238,7 @@ async def test_output_tokens_with_audio_and_text():


@pytest.mark.asyncio
-async def test_output_tokens_with_missing_num_tokens_out():
+async def test_output_tokens_with_missing_num_tokens_out(mocker: MockerFixture):
    """Test that output_tokens defaults to 0 when num_tokens_out is missing in metrics"""
    # Arrange
    request_input = RequestFuncInput(
@@ -269,8 +269,8 @@ async def test_output_tokens_with_missing_num_tokens_out():
    ]

    mock_response = MockResponse(200, chunks)
-    mock_session = AsyncMock()
-    mock_session.post = MagicMock(return_value=mock_response)
+    mock_session = mocker.AsyncMock()
+    mock_session.post = mocker.MagicMock(return_value=mock_response)

    # Act
    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -312,7 +312,7 @@ class TestTextLatencyAttribute:
        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):
+    async def test_text_latency_assigned_with_text_response(self, mocker: MockerFixture):
        """Test that text_latency is assigned when text response is received"""
        request_input = RequestFuncInput(
            model="test-model",
@@ -341,8 +341,8 @@ class TestTextLatencyAttribute:
        ]

        mock_response = MockResponse(200, chunks)
-        mock_session = AsyncMock()
-        mock_session.post = MagicMock(return_value=mock_response)
+        mock_session = mocker.AsyncMock()
+        mock_session.post = mocker.MagicMock(return_value=mock_response)

        # Act
        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -353,7 +353,7 @@ class TestTextLatencyAttribute:
        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):
+    async def test_text_latency_updated_with_multiple_text_chunks(self, mocker: MockerFixture):
        """Test that text_latency is updated as more text chunks arrive"""
        request_input = RequestFuncInput(
            model="test-model",
@@ -388,8 +388,8 @@ class TestTextLatencyAttribute:
        ]

        mock_response = MockResponse(200, chunks)
-        mock_session = AsyncMock()
-        mock_session.post = MagicMock(return_value=mock_response)
+        mock_session = mocker.AsyncMock()
+        mock_session.post = mocker.MagicMock(return_value=mock_response)

        # Act
        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -401,7 +401,7 @@ class TestTextLatencyAttribute:
        assert output.generated_text == "First second third"

    @pytest.mark.asyncio
-    async def test_text_latency_with_only_audio_response(self):
+    async def test_text_latency_with_only_audio_response(self, mocker: MockerFixture):
        """Test text_latency behavior when only audio is received"""
        request_input = RequestFuncInput(
            model="test-model",
@@ -424,8 +424,8 @@ class TestTextLatencyAttribute:
        ]

        mock_response = MockResponse(200, chunks)
-        mock_session = AsyncMock()
-        mock_session.post = MagicMock(return_value=mock_response)
+        mock_session = mocker.AsyncMock()
+        mock_session.post = mocker.MagicMock(return_value=mock_response)

        # Act
        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -437,7 +437,7 @@ class TestTextLatencyAttribute:
        assert output.text_latency >= 0, "text_latency should be non-negative"

    @pytest.mark.asyncio
-    async def test_text_latency_not_affected_by_metrics(self):
+    async def test_text_latency_not_affected_by_metrics(self, mocker: MockerFixture):
        """Test that text_latency is independent of metrics data"""
        request_input = RequestFuncInput(
            model="test-model",
@@ -467,8 +467,8 @@ class TestTextLatencyAttribute:
        ]

        mock_response = MockResponse(200, chunks)
-        mock_session = AsyncMock()
-        mock_session.post = MagicMock(return_value=mock_response)
+        mock_session = mocker.AsyncMock()
+        mock_session.post = mocker.MagicMock(return_value=mock_response)

        # Act
        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -481,7 +481,7 @@ class TestTextLatencyAttribute:
        assert output.output_tokens == 100, "metrics should not affect text_latency"

    @pytest.mark.asyncio
-    async def test_text_latency_mixed_modalities(self):
+    async def test_text_latency_mixed_modalities(self, mocker: MockerFixture):
        """Test text_latency with mixed text and audio modalities"""
        request_input = RequestFuncInput(
            model="test-model",
@@ -516,8 +516,8 @@ class TestTextLatencyAttribute:
        ]

        mock_response = MockResponse(200, chunks)
-        mock_session = AsyncMock()
-        mock_session.post = MagicMock(return_value=mock_response)
+        mock_session = mocker.AsyncMock()
+        mock_session.post = mocker.MagicMock(return_value=mock_response)

        # Act
        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -529,7 +529,7 @@ class TestTextLatencyAttribute:
        assert output.generated_text == "Text more text"

    @pytest.mark.asyncio
-    async def test_text_latency_value_consistency(self):
+    async def test_text_latency_value_consistency(self, mocker: MockerFixture):
        """Test that text_latency matches latency minus ttft relationship"""
        request_input = RequestFuncInput(
            model="test-model",
@@ -557,8 +557,8 @@ class TestTextLatencyAttribute:
        ]

        mock_response = MockResponse(200, chunks)
-        mock_session = AsyncMock()
-        mock_session.post = MagicMock(return_value=mock_response)
+        mock_session = mocker.AsyncMock()
+        mock_session.post = mocker.MagicMock(return_value=mock_response)

        # Act
        output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -848,41 +848,42 @@ def modify_stage_config(
                del config[key]

    # Apply updates
-    for key, value in updates.items():
-        if key == "stage_args":
-            if value and isinstance(value, dict):
-                stage_args = config.get("stage_args", [])
-                if not stage_args:
-                    raise ValueError("stage_args does not exist in config")
+    if updates:
+        for key, value in updates.items():
+            if key == "stage_args":
+                if value and isinstance(value, dict):
+                    stage_args = config.get("stage_args", [])
+                    if not stage_args:
+                        raise ValueError("stage_args does not exist in config")

-                for stage_id, stage_updates in value.items():
-                    # Find stage by ID
-                    target_stage = None
-                    for stage in stage_args:
-                        if stage.get("stage_id") == int(stage_id):
-                            target_stage = stage
-                            break
+                    for stage_id, stage_updates in value.items():
+                        # Find stage by ID
+                        target_stage = None
+                        for stage in stage_args:
+                            if stage.get("stage_id") == int(stage_id):
+                                target_stage = stage
+                                break

-                    if target_stage is None:
-                        available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
-                        raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}")
+                        if target_stage is None:
+                            available_ids = [s.get("stage_id") for s in stage_args if "stage_id" in s]
+                            raise KeyError(f"Stage ID {stage_id} not found, available: {available_ids}")

-                    # Apply updates to this stage
-                    for path, val in stage_updates.items():
-                        # Check if this is a simple key (not dot-separated)
-                        # Example: 'engine_input_source' vs 'engine_args.max_model_len'
-                        if "." not in path:
-                            # Direct key assignment (e.g., updating a list value)
-                            target_stage[path] = val
-                        else:
-                            # Dot-separated path (e.g., nested dict access)
-                            apply_update(target_stage, path, val)
-        elif "." in key:
-            # Apply using dot-separated path
-            apply_update(config, key, value)
-        else:
-            # Direct top-level key
-            config[key] = value
+                        # Apply updates to this stage
+                        for path, val in stage_updates.items():
+                            # Check if this is a simple key (not dot-separated)
+                            # Example: 'engine_input_source' vs 'engine_args.max_model_len'
+                            if "." not in path:
+                                # Direct key assignment (e.g., updating a list value)
+                                target_stage[path] = val
+                            else:
+                                # Dot-separated path (e.g., nested dict access)
+                                apply_update(target_stage, path, val)
+            elif "." in key:
+                # Apply using dot-separated path
+                apply_update(config, key, value)
+            else:
+                # Direct top-level key
+                config[key] = value

    # Save to new file with timestamp
    timestamp = int(time.time())
@@ -1057,7 +1058,8 @@ def omni_server(request, run_level):
    Multi-stage initialization can take 10-20+ minutes.
    """
    with _omni_server_lock:
-        model, stage_config_path = request.param
+        *port, model, stage_config_path = request.param
+        port = port[0] if port else None
        if run_level == "advanced_model":
            stage_config_path = modify_stage_config(
                stage_config_path,
@@ -1070,7 +1072,8 @@ def omni_server(request, run_level):
                },
            )

-        with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
+        server_args = ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]
+        with OmniServer(model, server_args, port=port) if port else OmniServer(model, server_args) as server:
            print("OmniServer started successfully")
            yield server
            print("OmniServer stopping...")
@@ -1288,11 +1291,15 @@ class OpenAIClientHandler:

        responses = []
        stream = request_config.get("stream", False)
+        modalities = request_config.get("modalities", ["text", "audio"])

        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
+                model=request_config.get("model"),
+                messages=request_config.get("messages"),
+                stream=stream,
+                modalities=modalities,
            )

            if stream:
@@ -1314,6 +1321,7 @@ class OpenAIClientHandler:
                        self.client.chat.completions.create,
                        model=request_config.get("model"),
                        messages=request_config.get("messages"),
+                        modalities=modalities,
                        stream=stream,
                    )
                    futures.append(future)
@@ -38,21 +38,77 @@ class _FakeLinearBase(LinearBase):
        torch.nn.Module.__init__(self)


+class _DummyBaseLayerWithLoRA(torch.nn.Module):
+    def __init__(self, base_layer: torch.nn.Module):
+        super().__init__()
+        self.base_layer = base_layer
+
+        self.set_calls: list[
+            tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
+        ] = []
+        self.reset_calls: int = 0
+        self.create_calls: int = 0
+
+    def set_lora(self, index: int, lora_a, lora_b):
+        assert index == 0
+        self.set_calls.append((lora_a, lora_b))
+
+    def reset_lora(self, index: int):
+        assert index == 0
+        self.reset_calls += 1
+
+    def create_lora_weights(self, max_loras, lora_config, model_config):
+        # Needs to be callable for scale test when rank changes, but not
+        # actually used since we mock everything and check everything based
+        # on set calls.
+        self.create_calls += 1
+
+
+class _DummyPipeline(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.transformer = torch.nn.Module()
+        self.transformer.foo = _FakeLinearBase()
+
+
+class _DummyLM(torch.nn.Module):
+    """LoRA enabled wrapper for _DummyPipeline."""
+
+    def __init__(self, rank: int):
+        super().__init__()
+        self.transformer = torch.nn.Module()
+        self.transformer.foo = _DummyBaseLayerWithLoRA(_FakeLinearBase())
+        self.rank = rank
+        self.loras = self.get_lora_modules()
+
+    def get_lora_modules(self):
+        return {"transformer.foo": self._get_initial_lora(self.rank)}
+
+    def get_lora(self, k: str) -> LoRALayerWeights:
+        """Get the unscaled LoRA weights for transformer.foo"""
+        return self.loras[k]
+
+    def _get_initial_lora(self, rank: int) -> LoRALayerWeights:
+        """Initializes a dummy LoRA for the current rank."""
+        A = torch.ones((rank, 4))
+        B = torch.ones((4, rank))
+        return LoRALayerWeights(
+            module_name="foo",
+            rank=rank,
+            lora_alpha=rank,
+            lora_a=A,
+            lora_b=B,
+        )
+
+
 def test_lora_manager_supported_modules_are_stable_with_wrapped_layers(monkeypatch):
    # Simulate a pipeline that already contains LoRA wrappers where the original
    # LinearBase is nested under ".base_layer".
    import vllm_omni.diffusion.lora.manager as manager_mod

-    class _DummyBaseLayerWithLoRA(torch.nn.Module):
-        def __init__(self, base_layer: torch.nn.Module):
-            super().__init__()
-            self.base_layer = base_layer
-
    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)

-    pipeline = torch.nn.Module()
-    pipeline.transformer = torch.nn.Module()
-    pipeline.transformer.foo = _DummyBaseLayerWithLoRA(_FakeLinearBase())
+    pipeline = _DummyLM(rank=2)

    # vLLM helper would see only the nested LinearBase and yield "base_layer".
    assert get_supported_lora_modules(pipeline) == ["base_layer"]
@@ -71,11 +127,6 @@ def test_lora_manager_supported_modules_are_stable_with_wrapped_layers(monkeypat
 def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
    import vllm_omni.diffusion.lora.manager as manager_mod

-    class _DummyBaseLayerWithLoRA(torch.nn.Module):
-        def __init__(self, base_layer: torch.nn.Module):
-            super().__init__()
-            self.base_layer = base_layer
-
    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)

    def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
@@ -116,11 +167,6 @@ def test_lora_manager_replace_layers_does_not_rewrap_base_layer(monkeypatch):
 def test_lora_manager_replaces_packed_layer_when_targeting_sublayers(monkeypatch):
    import vllm_omni.diffusion.lora.manager as manager_mod

-    class _DummyBaseLayerWithLoRA(torch.nn.Module):
-        def __init__(self, base_layer: torch.nn.Module):
-            super().__init__()
-            self.base_layer = base_layer
-
    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)

    def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs):
@@ -189,9 +235,8 @@ def test_lora_manager_activates_fused_lora_on_packed_layer():
            },
        )()
    }
-    manager._adapter_scales = {7: 0.5}

-    manager._activate_adapter(7)
+    manager._activate_adapter(7, 0.5)

    assert packed_layer.reset_calls == 0
    assert len(packed_layer.set_calls) == 1
@@ -234,9 +279,8 @@ def test_lora_manager_activates_packed_lora_from_sublayers():
    manager._registered_adapters = {
        1: type("LM", (), {"id": 1, "loras": loras, "get_lora": lambda self, k: self.loras.get(k)})()
    }
-    manager._adapter_scales = {1: 2.0}

-    manager._activate_adapter(1)
+    manager._activate_adapter(1, scale=2.0)

    assert packed_layer.reset_calls == 0
    assert len(packed_layer.set_calls) == 1
@@ -274,7 +318,7 @@ def test_lora_manager_evicts_lru_adapter_when_cache_full(monkeypatch):

    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
    monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
-    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id, scale: None)

    req1 = _dummy_lora_request(1)
    req2 = _dummy_lora_request(2)
@@ -306,7 +350,7 @@ def test_lora_manager_does_not_evict_pinned_adapter(monkeypatch):

    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
    monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
-    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id, scale: None)

    manager.set_active_adapter(_dummy_lora_request(1), lora_scale=1.0)
    assert manager.pin_adapter(1)
@@ -332,7 +376,7 @@ def test_lora_manager_warns_when_all_adapters_pinned(monkeypatch):

    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
    monkeypatch.setattr(manager, "_replace_layers_with_lora", lambda _peft: None)
-    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id: None)
+    monkeypatch.setattr(manager, "_activate_adapter", lambda _adapter_id, scale: None)

    manager.set_active_adapter(_dummy_lora_request(1), lora_scale=1.0)
    manager.set_active_adapter(_dummy_lora_request(2), lora_scale=1.0)
@@ -341,6 +385,118 @@ def test_lora_manager_warns_when_all_adapters_pinned(monkeypatch):
    assert manager.pin_adapter(2)

    manager.max_cached_adapters = 1
-    manager._evict_if_needed()
+    manager._evict_for_new_adapter()

    assert set(manager.list_adapters()) == {1, 2}
+
+
+def test_lora_manager_applies_multiple_scales_correctly(monkeypatch):
+    """Ensure that the LoRA manager applies scales correctly when the
+    active adapter receives a different scale, i.e., the rank is unchanged.
+    """
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyLoRALayer)
+
+    rank = 2
+    adapter_id = 7
+    req1 = _dummy_lora_request(adapter_id)
+    scale_1 = 0.25
+    scale_2 = 0.5
+
+    lora_model = _DummyLM(rank=rank)
+    manager = DiffusionLoRAManager(
+        pipeline=_DummyPipeline(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+
+    def _fake_load(_req: LoRARequest):
+        peft_helper = type("PH", (), {"r": rank})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    manager._registered_adapters = {
+        adapter_id: lora_model,
+    }
+    manager._lora_modules = {"transformer.foo": lora_model.transformer.foo}
+
+    # After the first scale, all B values should go from 1 -> scale_1
+    manager.set_active_adapter(req1, lora_scale=scale_1)
+    assert len(lora_model.transformer.foo.set_calls) == 1
+    lora_a, lora_b = lora_model.transformer.foo.set_calls[0]
+    assert torch.all(lora_a == 1)
+    assert torch.all(lora_b == scale_1)
+
+    # After the second scale, all B values should go from 1 -> scale_2
+    manager.set_active_adapter(req1, lora_scale=scale_2)
+    assert len(lora_model.transformer.foo.set_calls) == 2
+
+    lora_a, lora_b = lora_model.transformer.foo.set_calls[1]
+    assert torch.all(lora_a == 1)
+    assert torch.all(lora_b == scale_2)
+
+
+def test_lora_manager_scales_correctly_with_rank_changes(monkeypatch):
+    """Ensure that the LoRA manager correctly handles scaling when the rank
+    is changed and the buffers are reset + we reactivate.
+    """
+    import vllm_omni.diffusion.lora.manager as manager_mod
+
+    monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
+
+    rank = 2
+    adapter_id = 7
+    req1 = _dummy_lora_request(adapter_id)
+    initial_scale = 0.5
+
+    lora_model = _DummyLM(rank=rank)
+    manager = DiffusionLoRAManager(
+        pipeline=_DummyPipeline(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+
+    def _fake_load(_req: LoRARequest):
+        peft_helper = type("PH", (), {"r": rank})()
+        return lora_model, peft_helper
+
+    monkeypatch.setattr(manager, "_load_adapter", _fake_load)
+    manager._registered_adapters = {
+        adapter_id: lora_model,
+    }
+    manager._lora_modules = {"transformer.foo": lora_model.transformer.foo}
+
+    # Activate adapter with initial scale
+    manager.set_active_adapter(req1, lora_scale=initial_scale)
+    assert lora_model.transformer.foo.create_calls == 0
+    assert len(lora_model.transformer.foo.set_calls) == 1
+    lora_a, lora_b = lora_model.transformer.foo.set_calls[0]
+    assert torch.all(lora_a == 1)
+    assert torch.all(lora_b == initial_scale)
+
+    # Increase the rank; this resets the buffers, so the adapter is activated again
+    manager._ensure_max_lora_rank(8)
+
+    # Ensure we actually took the rank expansion path, which recreates
+    # and sets the weight buffets, but that the scale didn't change
+    assert lora_model.transformer.foo.create_calls == 1
+    assert len(lora_model.transformer.foo.set_calls) == 2
+    lora_a, lora_b = lora_model.transformer.foo.set_calls[1]
+    assert torch.all(lora_a == 1)
+    assert torch.all(lora_b == initial_scale)
+
+
+def test_scale_keys_are_rounded():
+    """Ensure that added adapter scales are rounded to avoid lookup
+    issues due to precision differences, e.g., computed scales.
+    """
+    manager = DiffusionLoRAManager(
+        pipeline=_DummyPipeline(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+    )
+    adapter_id = 1
+    # Currently we round keys to 3 decimal places
+    manager._update_adapter_scale(adapter_id, 0.0031)
+    assert manager._adapter_scales[adapter_id] == 0.003
@@ -0,0 +1,217 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+from PIL import Image
+
+import vllm_omni.diffusion.models.nextstep_1_1.pipeline_nextstep_1_1 as nextstep_pipeline_module
+from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep_heads import FlowMatchingHead
+from vllm_omni.diffusion.models.nextstep_1_1.pipeline_nextstep_1_1 import NextStep11Pipeline
+
+
+class _DummyImageHead:
+    def __init__(self, token_dim: int):
+        self.token_dim = token_dim
+        self.calls = []
+
+    def sample(
+        self,
+        c: torch.Tensor,
+        cfg: float,
+        cfg_img: float,
+        cfg_mult: int,
+        timesteps_shift: float,
+        num_sampling_steps: int,
+        noise_repeat: int,
+    ) -> torch.Tensor:
+        self.calls.append(
+            {
+                "batch": c.shape[0],
+                "cfg": cfg,
+                "cfg_img": cfg_img,
+                "cfg_mult": cfg_mult,
+                "noise_repeat": noise_repeat,
+            }
+        )
+        batch_per_prompt = c.shape[0] // cfg_mult
+        return torch.ones(batch_per_prompt, self.token_dim, dtype=c.dtype, device=c.device)
+
+
+class _DummyModel:
+    def __init__(self, hidden_dim: int, token_dim: int):
+        self.hidden_dim = hidden_dim
+        self.image_head = _DummyImageHead(token_dim)
+        self.forward_batches = []
+
+    def image_out_projector(self, c: torch.Tensor) -> torch.Tensor:
+        return c
+
+    def image_in_projector(self, sampled_tokens: torch.Tensor) -> torch.Tensor:
+        bsz = sampled_tokens.shape[0]
+        return torch.zeros(bsz, 1, self.hidden_dim, dtype=sampled_tokens.dtype, device=sampled_tokens.device)
+
+    def forward_model(self, inputs_embeds: torch.Tensor, attention_mask, past_key_values, use_cache: bool):
+        del attention_mask, use_cache
+        self.forward_batches.append(inputs_embeds.shape[0])
+        return SimpleNamespace(
+            last_hidden_state=torch.zeros(
+                inputs_embeds.shape[0],
+                1,
+                self.hidden_dim,
+                dtype=inputs_embeds.dtype,
+                device=inputs_embeds.device,
+            ),
+            past_key_values=past_key_values,
+        )
+
+
+def _make_minimal_pipeline_for_decoding(hidden_dim: int = 8, token_dim: int = 4):
+    pipeline = object.__new__(NextStep11Pipeline)
+    pipeline.config = SimpleNamespace(latent_channels=token_dim, latent_patch_size=1, use_gen_pos_embed=False)
+    pipeline.model = _DummyModel(hidden_dim=hidden_dim, token_dim=token_dim)
+    return pipeline
+
+
+@pytest.mark.parametrize(
+    ("cfg", "cfg_img", "has_image_conditions", "expected_cfg_mult", "expected_cfg_img"),
+    [
+        (1.0, 1.0, False, 1, 1.0),
+        (7.5, 1.0, False, 2, 1.0),
+        (7.5, 8.0, False, 2, 1.0),
+        (7.5, 1.5, True, 3, 1.5),
+    ],
+)
+def test_resolve_cfg_layout(cfg, cfg_img, has_image_conditions, expected_cfg_mult, expected_cfg_img):
+    cfg_mult, effective_cfg_img = NextStep11Pipeline._resolve_cfg_layout(cfg, cfg_img, has_image_conditions)
+    assert cfg_mult == expected_cfg_mult
+    assert effective_cfg_img == expected_cfg_img
+
+
+def test_build_captions_ignores_image_cfg_without_image_conditions():
+    pipeline = object.__new__(NextStep11Pipeline)
+    pipeline._image_str = lambda hw: f"<image:{hw}>"
+
+    captions, images, cfg_mult, effective_cfg_img = pipeline._build_captions(
+        captions=["a prompt"],
+        images=None,
+        num_images_per_caption=1,
+        positive_prompt=None,
+        negative_prompt="bad quality",
+        cfg=7.5,
+        cfg_img=8.0,
+    )
+
+    assert cfg_mult == 2
+    assert effective_cfg_img == 1.0
+    assert images is None
+    assert captions == ["a prompt", "bad quality"]
+
+
+def test_build_captions_enables_three_way_cfg_when_image_conditions_exist():
+    pipeline = object.__new__(NextStep11Pipeline)
+    pipeline._image_str = lambda hw: f"<image:{hw}>"
+
+    image = Image.new("RGB", (64, 32))
+    captions, images, cfg_mult, effective_cfg_img = pipeline._build_captions(
+        captions=["a prompt"],
+        images=[image],
+        num_images_per_caption=1,
+        positive_prompt=None,
+        negative_prompt="bad quality",
+        cfg=7.5,
+        cfg_img=1.5,
+    )
+
+    assert cfg_mult == 3
+    assert effective_cfg_img == 1.5
+    assert len(captions) == 3
+    assert captions[1].startswith("<image:")
+    assert captions[2] == "bad quality"
+    assert len(images) == 2
+
+
+def test_decoding_non_parallel_uses_cfg_mult_for_sampling_and_duplication(monkeypatch):
+    pipeline = _make_minimal_pipeline_for_decoding()
+    monkeypatch.setattr(nextstep_pipeline_module, "get_classifier_free_guidance_world_size", lambda: 1)
+
+    c = torch.zeros(2, 1, 8)
+    attention_mask = torch.ones(2, 3, dtype=torch.long)
+
+    tokens = pipeline.decoding(
+        c=c,
+        attention_mask=attention_mask,
+        past_key_values=None,
+        max_new_len=1,
+        num_images_per_caption=1,
+        cfg=7.5,
+        cfg_img=1.0,
+        cfg_mult=2,
+        progress=False,
+    )
+
+    assert tokens.shape == (1, 1, 4)
+    assert pipeline.model.image_head.calls[0]["cfg_mult"] == 2
+    assert pipeline.model.forward_batches == [2]
+
+
+def test_decoding_cfg_parallel_mismatch_falls_back_to_non_parallel(monkeypatch):
+    pipeline = _make_minimal_pipeline_for_decoding()
+
+    monkeypatch.setattr(nextstep_pipeline_module, "get_classifier_free_guidance_world_size", lambda: 2)
+
+    def _unexpected(*args, **kwargs):
+        del args, kwargs
+        raise AssertionError("CFG rank/group should not be queried on mismatch fallback.")
+
+    monkeypatch.setattr(nextstep_pipeline_module, "get_classifier_free_guidance_rank", _unexpected)
+    monkeypatch.setattr(nextstep_pipeline_module, "get_cfg_group", _unexpected)
+
+    c = torch.zeros(3, 1, 8)
+    attention_mask = torch.ones(3, 3, dtype=torch.long)
+
+    tokens = pipeline.decoding(
+        c=c,
+        attention_mask=attention_mask,
+        past_key_values=None,
+        max_new_len=1,
+        num_images_per_caption=1,
+        cfg=7.5,
+        cfg_img=1.5,
+        cfg_mult=3,
+        progress=False,
+    )
+
+    assert tokens.shape == (1, 1, 4)
+    assert pipeline.model.image_head.calls[0]["cfg_mult"] == 3
+    assert pipeline.model.forward_batches == [3]
+
+
+def test_decoding_rejects_incompatible_batch_and_cfg_mult(monkeypatch):
+    pipeline = _make_minimal_pipeline_for_decoding()
+    monkeypatch.setattr(nextstep_pipeline_module, "get_classifier_free_guidance_world_size", lambda: 1)
+
+    with pytest.raises(ValueError, match="not divisible"):
+        pipeline.decoding(
+            c=torch.zeros(5, 1, 8),
+            attention_mask=torch.ones(5, 3, dtype=torch.long),
+            past_key_values=None,
+            max_new_len=1,
+            num_images_per_caption=1,
+            cfg=7.5,
+            cfg_img=1.0,
+            cfg_mult=2,
+            progress=False,
+        )
+
+
+def test_flow_matching_head_sample_validates_cfg_mult_divisibility():
+    head = FlowMatchingHead(input_dim=4, cond_dim=8, dim=8, layers=1)
+
+    with pytest.raises(ValueError, match="not divisible"):
+        head.sample(
+            c=torch.zeros(5, 8),
+            cfg=2.0,
+            cfg_img=1.0,
+            cfg_mult=2,
+            num_sampling_steps=1,
+        )
@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
+from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@contextmanager
+def _noop_forward_context(*args, **kwargs):
+    del args, kwargs
+    yield
+
+
+class _DummyPipeline:
+    def __init__(self, output):
+        self._output = output
+        self.forward_calls = 0
+
+    def forward(self, req):
+        del req
+        self.forward_calls += 1
+        return self._output
+
+
+def _make_request(skip_cache_refresh: bool = True):
+    sampling_params = SimpleNamespace(
+        generator=None,
+        seed=None,
+        generator_device=None,
+        num_inference_steps=4,
+    )
+    return SimpleNamespace(
+        prompts=["a prompt"],
+        sampling_params=sampling_params,
+        skip_cache_refresh=skip_cache_refresh,
+    )
+
+
+def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
+    runner = object.__new__(DiffusionModelRunner)
+    runner.vllm_config = object()
+    runner.device = torch.device("cpu")
+    runner.pipeline = _DummyPipeline(output=SimpleNamespace(output="ok"))
+    runner.cache_backend = cache_backend
+    runner.offload_backend = None
+    runner.od_config = SimpleNamespace(
+        cache_backend=cache_backend_name,
+        enable_cache_dit_summary=enable_cache_dit_summary,
+    )
+    runner.kv_transfer_manager = SimpleNamespace(receive_kv_cache=lambda req, target_device: None)
+    return runner
+
+
+def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch):
+    """Guard cache diagnostics with runtime backend state to avoid stale-config crashes."""
+    runner = _make_runner(cache_backend=None, cache_backend_name="cache_dit")
+    req = _make_request(skip_cache_refresh=True)
+
+    cache_summary_calls = []
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(
+        model_runner_module,
+        "cache_summary",
+        lambda pipeline, details: cache_summary_calls.append((pipeline, details)),
+    )
+
+    output = DiffusionModelRunner.execute_model(runner, req)
+
+    assert output.output == "ok"
+    assert cache_summary_calls == []
+
+
+def test_execute_model_emits_cache_summary_with_active_cache_dit_backend(monkeypatch):
+    class _EnabledCacheBackend:
+        def is_enabled(self):
+            return True
+
+    runner = _make_runner(cache_backend=_EnabledCacheBackend(), cache_backend_name="cache_dit")
+    req = _make_request(skip_cache_refresh=True)
+
+    cache_summary_calls = []
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(
+        model_runner_module,
+        "cache_summary",
+        lambda pipeline, details: cache_summary_calls.append((pipeline, details)),
+    )
+
+    output = DiffusionModelRunner.execute_model(runner, req)
+
+    assert output.output == "ok"
+    assert cache_summary_calls == [(runner.pipeline, True)]
+
+
+def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch):
+    class _DummyLoader:
+        def __init__(self, load_config):
+            del load_config
+
+        def load_model(self, **kwargs):
+            del kwargs
+            return SimpleNamespace(transformer=torch.nn.Identity())
+
+    class _DummyMemoryProfiler:
+        consumed_memory = 0
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, tb):
+            del exc_type, exc, tb
+            return False
+
+    class _DummyCacheBackend:
+        def __init__(self):
+            self.enabled = False
+
+        def enable(self, pipeline):
+            del pipeline
+            self.enabled = True
+
+    dummy_cache_backend = _DummyCacheBackend()
+
+    runner = object.__new__(DiffusionModelRunner)
+    runner.vllm_config = object()
+    runner.device = torch.device("cpu")
+    runner.pipeline = None
+    runner.cache_backend = None
+    runner.offload_backend = None
+    runner.od_config = SimpleNamespace(
+        enable_cpu_offload=False,
+        enable_layerwise_offload=False,
+        cache_backend="cache_dit",
+        cache_config={},
+        model_class_name="NextStep11Pipeline",
+        enforce_eager=True,
+    )
+
+    monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object())
+    monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _DummyLoader)
+    monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _DummyMemoryProfiler)
+    monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda od_config, device: None)
+    monkeypatch.setattr(
+        model_runner_module, "get_cache_backend", lambda cache_backend, cache_config: dummy_cache_backend
+    )
+
+    DiffusionModelRunner.load_model(runner)
+
+    assert runner.cache_backend is None
+    assert runner.od_config.cache_backend is None
+    assert dummy_cache_backend.enabled is False
@@ -10,10 +10,9 @@ This module tests the DiffusionWorker implementation:
 - wake_up: Waking worker from sleep mode
 """

-from unittest.mock import Mock, patch
-
 import pytest
 import torch
+from pytest_mock import MockerFixture

 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker

@@ -21,9 +20,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]


@pytest.fixture
-def mock_od_config():
+def mock_od_config(mocker: MockerFixture):
    """Create a mock OmniDiffusionConfig."""
-    config = Mock()
+    config = mocker.Mock()
    config.num_gpus = 1
    config.master_port = 12345
    config.enable_sleep_mode = False
@@ -34,22 +33,23 @@ def mock_od_config():


@pytest.fixture
-def mock_gpu_worker(mock_od_config):
+def mock_gpu_worker(mocker: MockerFixture, mock_od_config):
    """Create a DiffusionWorker with mocked initialization."""
-    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()
-        worker.model_runner.pipeline = Mock()
-        worker.device = torch.device("cuda", 0)
-        worker._sleep_saved_buffers = {}
-        return worker
+    mocker.patch.object(DiffusionWorker, "init_device")
+    mocker.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 = mocker.Mock()
+    worker.model_runner.pipeline = mocker.Mock()
+    worker.device = torch.device("cuda", 0)
+    worker._sleep_saved_buffers = {}
+    return worker


 class TestDiffusionWorkerLoadWeights:
    """Test DiffusionWorker.load_weights method."""

-    def test_load_weights_calls_pipeline(self, mock_gpu_worker):
+    def test_load_weights_calls_pipeline(self, mocker: MockerFixture, mock_gpu_worker):
        """Test that load_weights delegates to model_runner.load_weights."""
        # Setup mock weights
        mock_weights = [
@@ -59,7 +59,7 @@ class TestDiffusionWorkerLoadWeights:
        expected_loaded = {"layer1.weight", "layer2.weight"}

        # Configure model_runner mock
-        mock_gpu_worker.model_runner.load_weights = Mock(return_value=expected_loaded)
+        mock_gpu_worker.model_runner.load_weights = mocker.Mock(return_value=expected_loaded)

        # Call load_weights
        result = mock_gpu_worker.load_weights(mock_weights)
@@ -68,9 +68,9 @@ class TestDiffusionWorkerLoadWeights:
        mock_gpu_worker.model_runner.load_weights.assert_called_once_with(mock_weights)
        assert result == expected_loaded

-    def test_load_weights_empty_iterable(self, mock_gpu_worker):
+    def test_load_weights_empty_iterable(self, mocker: MockerFixture, mock_gpu_worker):
        """Test load_weights with empty weights iterable."""
-        mock_gpu_worker.model_runner.load_weights = Mock(return_value=set())
+        mock_gpu_worker.model_runner.load_weights = mocker.Mock(return_value=set())

        result = mock_gpu_worker.load_weights([])

@@ -81,11 +81,12 @@ 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_get_process_memory, mock_gpu_worker):
+    def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
        """Test sleep mode level 1 (offload weights only)."""
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
+
        # Setup process-scoped memory mocks
        # Before sleep: 3GB used
        # After sleep: 1GB used (freed 2GB)
@@ -95,9 +96,9 @@ class TestDiffusionWorkerSleep:
        ]

        # Setup allocator mock
-        mock_allocator = Mock()
-        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
-        mock_allocator.sleep = Mock()
+        mock_allocator = mocker.Mock()
+        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
+        mock_allocator.sleep = mocker.Mock()

        # Call sleep with level 1
        result = mock_gpu_worker.sleep(level=1)
@@ -108,11 +109,12 @@ 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_get_process_memory, mock_gpu_worker):
+    def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker):
        """Test sleep mode level 2 (offload all, save buffers)."""
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
+
        # Setup process-scoped memory mocks
        mock_get_process_memory.side_effect = [
            5 * 1024**3,  # Before sleep
@@ -120,14 +122,14 @@ class TestDiffusionWorkerSleep:
        ]

        # Setup allocator mock
-        mock_allocator = Mock()
-        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
-        mock_allocator.sleep = Mock()
+        mock_allocator = mocker.Mock()
+        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
+        mock_allocator.sleep = mocker.Mock()

        # Mock pipeline buffers
        mock_buffer1 = torch.randn(10, 10)
        mock_buffer2 = torch.randn(20, 20)
-        mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
+        mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
            return_value=[
                ("buffer1", mock_buffer1),
                ("buffer2", mock_buffer2),
@@ -146,42 +148,32 @@ 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_get_process_memory,
-        mock_gpu_worker,
-    ):
+    def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker):
        """Test that sleep validates memory was actually freed."""
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
+
        # 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_allocator = Mock()
-        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
-        mock_allocator.sleep = Mock()
+        mock_allocator = mocker.Mock()
+        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
+        mock_allocator.sleep = mocker.Mock()

        # This should raise an assertion error
        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,
-    ):
+    def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker):
        """Test sleep uses device-scoped fallback when NVML is unavailable."""
+
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+        mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
+        mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
        mock_get_process_memory.side_effect = [None, None]
        mock_platform.get_free_memory.side_effect = [
            1 * 1024**3,  # Before sleep
@@ -189,9 +181,9 @@ class TestDiffusionWorkerSleep:
        ]
        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()
+        mock_allocator = mocker.Mock()
+        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
+        mock_allocator.sleep = mocker.Mock()

        result = mock_gpu_worker.sleep(level=1)

@@ -202,13 +194,14 @@ class TestDiffusionWorkerSleep:
 class TestDiffusionWorkerWakeUp:
    """Test DiffusionWorker.wake_up method."""

-    @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker):
+    def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker):
        """Test wake_up without saved buffers (level 1 sleep)."""
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+
        # Setup allocator mock
-        mock_allocator = Mock()
-        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
-        mock_allocator.wake_up = Mock()
+        mock_allocator = mocker.Mock()
+        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
+        mock_allocator.wake_up = mocker.Mock()

        # Ensure no saved buffers
        mock_gpu_worker._sleep_saved_buffers = {}
@@ -220,13 +213,14 @@ class TestDiffusionWorkerWakeUp:
        mock_allocator.wake_up.assert_called_once_with(["weights"])
        assert result is True

-    @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_wake_up_with_buffers(self, mock_allocator_class, mock_gpu_worker):
+    def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker):
        """Test wake_up with saved buffers (level 2 sleep)."""
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+
        # Setup allocator mock
-        mock_allocator = Mock()
-        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
-        mock_allocator.wake_up = Mock()
+        mock_allocator = mocker.Mock()
+        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
+        mock_allocator.wake_up = mocker.Mock()

        # Create saved buffers
        saved_buffer1 = torch.randn(10, 10)
@@ -237,12 +231,12 @@ class TestDiffusionWorkerWakeUp:
        }

        # Mock pipeline buffers (these will be restored)
-        mock_buffer1 = Mock()
-        mock_buffer1.data = Mock()
-        mock_buffer2 = Mock()
-        mock_buffer2.data = Mock()
+        mock_buffer1 = mocker.Mock()
+        mock_buffer1.data = mocker.Mock()
+        mock_buffer2 = mocker.Mock()
+        mock_buffer2.data = mocker.Mock()

-        mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
+        mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
            return_value=[
                ("buffer1", mock_buffer1),
                ("buffer2", mock_buffer2),
@@ -263,13 +257,14 @@ class TestDiffusionWorkerWakeUp:
        assert len(mock_gpu_worker._sleep_saved_buffers) == 0
        assert result is True

-    @patch("vllm.device_allocator.cumem.CuMemAllocator")
-    def test_wake_up_partial_buffer_restore(self, mock_allocator_class, mock_gpu_worker):
+    def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker):
        """Test wake_up only restores buffers that were saved."""
+        mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
+
        # Setup allocator mock
-        mock_allocator = Mock()
-        mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
-        mock_allocator.wake_up = Mock()
+        mock_allocator = mocker.Mock()
+        mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
+        mock_allocator.wake_up = mocker.Mock()

        # Only save buffer1, not buffer2
        saved_buffer1 = torch.randn(10, 10)
@@ -278,12 +273,12 @@ class TestDiffusionWorkerWakeUp:
        }

        # Mock pipeline has both buffers
-        mock_buffer1 = Mock()
-        mock_buffer1.data = Mock()
-        mock_buffer2 = Mock()
-        mock_buffer2.data = Mock()
+        mock_buffer1 = mocker.Mock()
+        mock_buffer1.data = mocker.Mock()
+        mock_buffer2 = mocker.Mock()
+        mock_buffer2.data = mocker.Mock()

-        mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
+        mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
            return_value=[
                ("buffer1", mock_buffer1),
                ("buffer2", mock_buffer2),
@@ -13,9 +13,9 @@ This module tests the WorkerWrapperBase implementation:
 """

 from typing import Any
-from unittest.mock import Mock, patch

 import pytest
+from pytest_mock import MockerFixture

 from vllm_omni.diffusion.worker.diffusion_worker import (
    CustomPipelineWorkerExtension,
@@ -23,15 +23,18 @@ from vllm_omni.diffusion.worker.diffusion_worker import (
    WorkerWrapperBase,
 )

+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 # -------------------------------------------------------------------------
 # Fixtures
 # -------------------------------------------------------------------------


@pytest.fixture
-def mock_od_config():
+def mock_od_config(mocker: MockerFixture):
    """Create a mock OmniDiffusionConfig for use in tests."""
-    config = Mock()
+    config = mocker.Mock()
    config.num_gpus = 1
    config.master_port = 12345
    config.enable_sleep_mode = False
@@ -72,9 +75,10 @@ class MockCustomPipeline:
 class TestWorkerWrapperBaseInitialization:
    """Test WorkerWrapperBase initialization behavior."""

-    @patch.object(DiffusionWorker, "__init__", return_value=None)
-    def test_basic_initialization(self, mock_worker_init, mock_od_config):
+    def test_basic_initialization(self, mocker: MockerFixture, mock_od_config):
        """Test basic initialization without extensions."""
+        mock_worker_init = mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
+
        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
@@ -103,49 +107,49 @@ class TestWorkerWrapperBaseInitialization:
 class TestWorkerWrapperBaseExtension:
    """Test WorkerWrapperBase worker extension functionality."""

-    def test_prepare_worker_class_without_extension(self, mock_od_config):
+    def test_prepare_worker_class_without_extension(self, mocker: MockerFixture, 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
+        mocker.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):
+    def test_prepare_worker_class_with_extension_class(self, mocker: MockerFixture, 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,
-            )
+        mocker.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__
+        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):
+    def test_prepare_worker_class_with_extension_string(self, mocker: MockerFixture, mock_od_config):
        """Test _prepare_worker_class with worker extension as string."""
+        mock_resolve = mocker.patch("vllm.utils.import_utils.resolve_obj_by_qualname")
        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",
-            )
+        mocker.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")
+        assert hasattr(wrapper.worker.__class__, "custom_method")


 # -------------------------------------------------------------------------
@@ -156,70 +160,70 @@ class TestWorkerWrapperBaseExtension:
 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):
+    def test_generate_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test that generate() delegates to worker.generate()."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        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_output = mocker.Mock()
+        wrapper.worker.generate = mocker.Mock(return_value=mock_output)

-        mock_requests = [Mock()]
+        mock_requests = [mocker.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):
+    def test_execute_model_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test that execute_model() delegates to worker.execute_model()."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        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_output = mocker.Mock()
+        wrapper.worker.execute_model = mocker.Mock(return_value=mock_output)

-        mock_reqs = [Mock()]
+        mock_reqs = [mocker.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):
+    def test_load_weights_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test that load_weights() delegates to worker.load_weights()."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        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)
+        wrapper.worker.load_weights = mocker.Mock(return_value=expected_result)

-        mock_weights = [("weight1", Mock()), ("weight2", Mock())]
+        mock_weights = [("weight1", mocker.Mock()), ("weight2", mocker.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):
+    def test_sleep_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test that sleep() delegates to worker.sleep()."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
-        wrapper.worker.sleep = Mock(return_value=True)
+        wrapper.worker.sleep = mocker.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):
+    def test_wake_up_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test that wake_up() delegates to worker.wake_up()."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
-        wrapper.worker.wake_up = Mock(return_value=True)
+        wrapper.worker.wake_up = mocker.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):
+    def test_shutdown_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test that shutdown() delegates to worker.shutdown()."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
-        wrapper.worker.shutdown = Mock(return_value=None)
+        wrapper.worker.shutdown = mocker.Mock(return_value=None)

        result = wrapper.shutdown()
        wrapper.worker.shutdown.assert_called_once()
@@ -234,39 +238,39 @@ class TestWorkerWrapperBaseDelegation:
 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):
+    def test_execute_method_success(self, mocker: MockerFixture, mock_od_config):
        """Test execute_method successfully calls worker method."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
-        wrapper.worker.test_method = Mock(return_value="method_result")
+        wrapper.worker.test_method = mocker.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):
+    def test_execute_method_with_no_args(self, mocker: MockerFixture, mock_od_config):
        """Test execute_method with no arguments."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        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")
+        wrapper.worker.no_args_method = mocker.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):
+    def test_execute_method_error(self, mocker: MockerFixture, mock_od_config):
        """Test execute_method raises exception on error."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
-        wrapper.worker.error_method = Mock(side_effect=RuntimeError("Test error"))
+        wrapper.worker.error_method = mocker.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):
+    def test_execute_method_invalid_type(self, mocker: MockerFixture, mock_od_config):
        """Test execute_method with invalid method type."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)

        with pytest.raises(AssertionError, match="Method must be str"):
@@ -281,26 +285,26 @@ class TestWorkerWrapperBaseExecuteMethod:
 class TestWorkerWrapperBaseGetAttr:
    """Test WorkerWrapperBase.__getattr__ delegation."""

-    @patch.object(DiffusionWorker, "__init__", return_value=None)
-    def test_getattr_delegation(self, mock_worker_init, mock_od_config):
+    def test_getattr_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test __getattr__ delegates to worker attributes."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        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):
+    def test_getattr_method_access(self, mocker: MockerFixture, mock_od_config):
        """Test __getattr__ delegates to worker methods."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
-        wrapper.worker.custom_method = Mock(return_value="method_result")
+        wrapper.worker.custom_method = mocker.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):
+    def test_getattr_missing_attribute(self, mocker: MockerFixture, mock_od_config):
        """Test __getattr__ raises AttributeError for missing attributes."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
        with pytest.raises(AttributeError):
            _ = wrapper.nonexistent_attribute
@@ -314,47 +318,46 @@ class TestWorkerWrapperBaseGetAttr:
 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):
+    def test_extension_conflict_warning(self, mocker: MockerFixture, mock_od_config, caplog):
        """Test a warning is logged when an extension conflicts with worker."""
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)

        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
+        mocker.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):
+    def test_multiple_extensions_same_class(self, mocker: MockerFixture, 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,
-            )
+        mocker.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")
+        assert hasattr(wrapper1.worker, "custom_method")
+        assert hasattr(wrapper2.worker, "custom_method")


 # -------------------------------------------------------------------------
@@ -365,11 +368,12 @@ class TestWorkerWrapperBaseEdgeCases:
 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):
+    def test_re_init_pipeline_basic(self, mocker: MockerFixture, mock_od_config):
        """Test basic re_init_pipeline functionality."""
+        mocker.patch("torch.cuda.empty_cache")
+        mocker.patch("gc.collect")
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
+
        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
@@ -378,12 +382,12 @@ class TestCustomPipelineWorkerExtension:
        )

        # Setup mock model_runner and pipeline
-        mock_model_runner = Mock()
-        mock_pipeline = Mock()
+        mock_model_runner = mocker.Mock()
+        mock_pipeline = mocker.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()
+        wrapper.worker.init_lora_manager = mocker.Mock()
+        wrapper.worker.load_model = mocker.Mock()

        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}

@@ -397,11 +401,12 @@ class TestCustomPipelineWorkerExtension:
        )
        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):
+    def test_re_init_pipeline_cleanup(self, mocker: MockerFixture, mock_od_config):
        """Test that re_init_pipeline properly cleans up old pipeline."""
+        mock_gc_collect = mocker.patch("gc.collect")
+        mock_empty_cache = mocker.patch("torch.cuda.empty_cache")
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
+
        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
@@ -410,12 +415,12 @@ class TestCustomPipelineWorkerExtension:
        )

        # Setup mock model_runner with pipeline
-        mock_model_runner = Mock()
-        mock_pipeline = Mock()
+        mock_model_runner = mocker.Mock()
+        mock_pipeline = mocker.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()
+        wrapper.worker.init_lora_manager = mocker.Mock()
+        wrapper.worker.load_model = mocker.Mock()

        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}

@@ -426,11 +431,12 @@ class TestCustomPipelineWorkerExtension:
        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):
+    def test_re_init_pipeline_none_pipeline(self, mocker: MockerFixture, mock_od_config):
        """Test re_init_pipeline when pipeline is None."""
+        mocker.patch("torch.cuda.empty_cache")
+        mocker.patch("gc.collect")
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
+
        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
@@ -439,11 +445,11 @@ class TestCustomPipelineWorkerExtension:
        )

        # Setup mock model_runner with None pipeline
-        mock_model_runner = Mock()
+        mock_model_runner = mocker.Mock()
        mock_model_runner.pipeline = None
        wrapper.worker.model_runner = mock_model_runner
-        wrapper.worker.init_lora_manager = Mock()
-        wrapper.worker.load_model = Mock()
+        wrapper.worker.init_lora_manager = mocker.Mock()
+        wrapper.worker.load_model = mocker.Mock()

        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}

@@ -453,32 +459,30 @@ class TestCustomPipelineWorkerExtension:
        # 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):
+    def test_custom_pipeline_args_initialization(self, mocker: MockerFixture, 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
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
+        mock_prepare = mocker.patch.object(WorkerWrapperBase, "_prepare_worker_class")
+        # Create a mock worker class with re_init_pipeline
+        mock_worker_class = mocker.Mock()
+        mock_worker_instance = mocker.Mock()
+        mock_worker_instance.re_init_pipeline = mocker.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,
-                )
+        _ = 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)
+        # 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):
+    def test_custom_pipeline_with_explicit_extension(self, mocker: MockerFixture, mock_od_config):
        """Test that explicit worker_extension_cls is preserved when custom_pipeline_args is provided."""

        class CustomExtension:
@@ -490,24 +494,25 @@ class TestCustomPipelineWorkerExtension:

        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,
-            )
+        mocker.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")
+        # 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):
+    def test_re_init_pipeline_multiple_calls(self, mocker: MockerFixture, mock_od_config):
        """Test calling re_init_pipeline multiple times."""
+        mocker.patch("torch.cuda.empty_cache")
+        mocker.patch("gc.collect")
+        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
+
        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
@@ -516,13 +521,13 @@ class TestCustomPipelineWorkerExtension:
        )

        # Setup mock model_runner
-        mock_model_runner = Mock()
-        mock_pipeline1 = Mock()
-        mock_pipeline2 = Mock()
+        mock_model_runner = mocker.Mock()
+        mock_pipeline1 = mocker.Mock()
+        mock_pipeline2 = mocker.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()
+        wrapper.worker.init_lora_manager = mocker.Mock()
+        wrapper.worker.load_model = mocker.Mock()

        # First call
        custom_args1 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
@@ -1,19 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-from unittest.mock import MagicMock
-
 import pytest
+from pytest_mock import MockerFixture

 from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector, try_send_via_connector
 from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
 from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec, OmniTransferConfig
 from vllm_omni.distributed.omni_connectors.utils.initialization import get_connectors_config_for_stage

+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+

@pytest.fixture
-def mock_objects():
-    return {"connector": MagicMock(), "metrics": MagicMock(), "queue_fn": MagicMock()}
+def mock_objects(mocker: MockerFixture):
+    return {"connector": mocker.MagicMock(), "metrics": mocker.MagicMock(), "queue_fn": mocker.MagicMock()}


 def test_send_success(mock_objects):
@@ -138,7 +139,7 @@ def test_recv_no_connector():
    assert inputs is None


-def test_shm_connector_flow():
+def test_shm_connector_flow(mocker: MockerFixture):
    """
    Verify the full flow: Send -> Adapter -> Connector -> Adapter -> Recv.
    Using real SharedMemoryConnector (inline mode for simplicity).
@@ -161,7 +162,7 @@ def test_shm_connector_flow():
    def mock_submit(payload):
        queue_capture.append(payload)

-    mock_metrics = MagicMock()
+    mock_metrics = mocker.MagicMock()

    # 3. Send
    success = try_send_via_connector(
@@ -217,7 +218,7 @@ def test_get_connectors_for_stage():
    assert stage_2_config["from_stage_1"]["spec"]["name"] == "C2"


-def test_recv_with_missing_metadata():
+def test_recv_with_missing_metadata(mocker: MockerFixture):
    """Test recv when queue payload is malformed (missing metadata)."""
    # Connector expects metadata but task doesn't have it
    task = {
@@ -226,7 +227,7 @@ def test_recv_with_missing_metadata():
        "from_stage": "0",
        # Missing "connector_metadata"
    }
-    mock_conn = MagicMock()
+    mock_conn = mocker.MagicMock()
    # If get is called with None metadata, connector usually handles it or adapter handles exception
    mock_conn.get.side_effect = Exception("Get failed")

@@ -1,14 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from unittest.mock import MagicMock

 import pytest
+from pytest_mock import MockerFixture

 from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
 from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
 from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
 from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer

+# pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+

 def test_basic_serialization():
    """Test basic msgpack serialization."""
@@ -79,19 +81,19 @@ def test_put_get_inline(shm_connector):
    assert size == ret_size


-def test_put_get_shm(shm_connector, monkeypatch):
+def test_put_get_shm(mocker: MockerFixture, shm_connector, monkeypatch: pytest.MonkeyPatch):
    """Test SHM transfer logic for large data (Mocked)."""
    # Create data larger than 100 bytes
    data = {"large": "x" * 200}

    # Mock SHM return values
    mock_handle = {"name": "test_shm", "size": 200}
-    mock_write = MagicMock(return_value=mock_handle)
+    mock_write = mocker.MagicMock(return_value=mock_handle)
    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_write_bytes", mock_write)

    # When reading, return the serialized bytes of the data
    serialized_data = shm_connector.serialize_obj(data)
-    mock_read = MagicMock(return_value=serialized_data)
+    mock_read = mocker.MagicMock(return_value=serialized_data)
    monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_read_bytes", mock_read)

    # Put
@@ -4,10 +4,10 @@
 import threading
 from collections import deque
 from types import SimpleNamespace
-from unittest.mock import MagicMock

 import pytest
 import torch
+from pytest_mock import MockerFixture
 from vllm.v1.request import RequestStatus

 from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
@@ -39,9 +39,9 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)


@pytest.fixture
-def build_adapter(monkeypatch):
+def build_adapter(monkeypatch, mocker: MockerFixture):
    def _build(*, stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2):
-        connector = MagicMock()
+        connector = mocker.MagicMock()
        connector.stage_id = stage_id
        connector.get.return_value = None
        connector.put.return_value = (True, 1, {})
@@ -105,6 +105,98 @@ def test_manager_extraction(kv_config, mock_connector, common_constants):
    assert data["layer_blocks"]["key_cache"][0].shape == expected_shape


+def test_manager_extraction_tuple_layout(kv_config, mock_connector, common_constants):
+    """Test extraction with tuple layout."""
+    num_layers = common_constants["num_layers"]
+    block_size = common_constants["block_size"]
+    num_heads = common_constants["num_heads"]
+    head_dim = common_constants["head_dim"]
+    seq_len = common_constants["seq_len"]
+    req_id = common_constants["req_id"]
+
+    num_blocks = 10
+    kv_caches = []
+    for _ in range(num_layers):
+        k_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
+        v_cache = torch.randn(num_blocks, block_size, num_heads, head_dim)
+        kv_caches.append((k_cache, v_cache))
+
+    block_ids = [1, 3, 5]
+    finished_reqs = {req_id: {"block_ids": block_ids, "seq_len": seq_len}}
+
+    manager = OmniKVTransferManager(kv_config)
+    manager._connector = mock_connector
+
+    processed = manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32")
+    assert req_id in processed
+
+    full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}"
+    expected_key = f"stage1->stage2:{full_request_id}"
+    assert expected_key in mock_connector.store
+
+    data = mock_connector.store[expected_key]
+    expected_shape = (seq_len, num_heads, head_dim)
+    for idx in range(len(kv_caches)):
+        assert data["layer_blocks"]["key_cache"][idx].shape == expected_shape
+        assert data["layer_blocks"]["value_cache"][idx].shape == expected_shape
+
+
+def test_manager_extraction_mismatched_kv_block_counts(kv_config, mock_connector, common_constants):
+    """Mismatched key/value block counts should not crash extraction."""
+    block_size = common_constants["block_size"]
+    num_heads = common_constants["num_heads"]
+    head_dim = common_constants["head_dim"]
+    req_id = common_constants["req_id"]
+
+    key_blocks = torch.randn(3, block_size, num_heads, head_dim)
+    value_blocks = torch.randn(2, block_size, num_heads, head_dim)
+    kv_caches = [(key_blocks, value_blocks)]
+
+    finished_reqs = {req_id: {"block_ids": [0, 1, 2], "seq_len": 32}}
+
+    manager = OmniKVTransferManager(kv_config)
+    manager._connector = mock_connector
+
+    processed = manager.handle_finished_requests_kv_transfer(finished_reqs, kv_caches, block_size, "float32")
+    assert req_id in processed
+
+    full_request_id = f"omni_stage1_to_stage2_kv_cache_{req_id}"
+    expected_key = f"stage1->stage2:{full_request_id}"
+    assert expected_key in mock_connector.store
+
+    data = mock_connector.store[expected_key]
+    expected_shape = (2 * block_size, num_heads, head_dim)
+    assert data["layer_blocks"]["key_cache"][0].shape == expected_shape
+    assert data["layer_blocks"]["value_cache"][0].shape == expected_shape
+
+
+@pytest.mark.parametrize(
+    "invalid_case",
+    ["invalid_stacked_shape", "invalid_tuple_length", "non_tensor_entries"],
+)
+def test_normalize_layer_kv_rejects_invalid_inputs(kv_config, common_constants, invalid_case):
+    """_normalize_layer_kv should reject malformed KV representations."""
+    block_size = common_constants["block_size"]
+    num_heads = common_constants["num_heads"]
+    head_dim = common_constants["head_dim"]
+    req_id = common_constants["req_id"]
+
+    if invalid_case == "invalid_stacked_shape":
+        layer_kv = torch.randn(3, block_size, num_heads, head_dim)
+    elif invalid_case == "invalid_tuple_length":
+        layer_kv = (
+            torch.randn(2, block_size, num_heads, head_dim),
+            torch.randn(2, block_size, num_heads, head_dim),
+            torch.randn(2, block_size, num_heads, head_dim),
+        )
+    else:
+        layer_kv = (torch.randn(2, block_size, num_heads, head_dim), "not-a-tensor")
+
+    manager = OmniKVTransferManager(kv_config)
+    normalized = manager._normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0)
+    assert normalized is None
+
+
 def test_manager_reception(kv_config, mock_connector, common_constants):
    """Test reception and injection logic in OmniKVTransferManager."""
    num_layers = common_constants["num_layers"]
@@ -0,0 +1,227 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for BufferAllocator and ManagedBuffer.
+These tests do NOT require Mooncake or RDMA environment.
+"""
+
+import threading
+
+import pytest
+import torch
+
+from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
+    BufferAllocator,
+    ManagedBuffer,
+)
+
+# All tests in this file are pure-CPU unit tests for the memory allocator.
+pytestmark = [pytest.mark.cpu, pytest.mark.parallel]
+
+
+@pytest.mark.core_model
+class TestBufferAllocator:
+    """Unit tests for BufferAllocator."""
+
+    def test_basic_alloc_free(self):
+        """Verify alloc, free, and reuse of freed space."""
+        allocator = BufferAllocator(total_size=4096, alignment=64)
+
+        offset1 = allocator.alloc(512)
+        assert offset1 == 0
+
+        offset2 = allocator.alloc(512)
+        assert offset2 > 0
+
+        # Free first block, should be reusable
+        allocator.free(offset1, 512)
+        offset3 = allocator.alloc(512)
+        assert offset3 == 0
+
+    def test_alignment(self):
+        """Verify allocation respects alignment."""
+        allocator = BufferAllocator(total_size=4096, alignment=128)
+
+        _offset1 = allocator.alloc(100)
+        offset2 = allocator.alloc(100)
+
+        assert offset2 % 128 == 0
+        assert offset2 == 128
+
+    def test_exhaustion_and_recovery(self):
+        """Test that full allocation fails, then succeeds after free."""
+        allocator = BufferAllocator(total_size=1024, alignment=64)
+
+        offset = allocator.alloc(1024)
+        assert offset == 0
+
+        with pytest.raises(MemoryError):
+            allocator.alloc(64)
+
+        allocator.free(offset, 1024)
+        offset2 = allocator.alloc(1024)
+        assert offset2 == 0
+
+    def test_thread_safety(self):
+        """Verify allocator is thread-safe under concurrent access."""
+        allocator = BufferAllocator(total_size=1024 * 1024, alignment=64)
+        errors = []
+
+        def worker(worker_id):
+            try:
+                for i in range(100):
+                    size = 1024 + (i % 10) * 64
+                    offset = allocator.alloc(size)
+                    allocator.free(offset, size)
+            except Exception as e:
+                errors.append(f"Worker {worker_id}: {e}")
+
+        threads = [threading.Thread(target=worker, args=(i,)) for i in range(10)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        assert len(errors) == 0, f"Thread safety errors: {errors}"
+
+
+class TestAllocatorInvariants:
+    """
+    Defensive invariant tests for BufferAllocator: double-free, partial
+    overlap corruption, and adjacent-block merging.
+
+    Marked @pytest.mark.slow so they are skipped in quick CI but retained
+    as regression safety-net.
+    """
+
+    @pytest.mark.slow
+    def test_double_free_exact_is_safe(self):
+        """Double-free of exact same block should warn but NOT crash."""
+        allocator = BufferAllocator(total_size=4096, alignment=64)
+        offset = allocator.alloc(256)
+        allocator.free(offset, 256)
+        # Second free of the same block — should be silently ignored
+        allocator.free(offset, 256)
+        # Pool should still be consistent: allocate full size back
+        offset2 = allocator.alloc(4096)
+        assert offset2 == 0
+
+    @pytest.mark.slow
+    def test_double_free_after_merge_is_safe(self):
+        """
+        Free A then B (adjacent → merged), then free A again.
+        The allocator must detect A is already within the merged block.
+        """
+        allocator = BufferAllocator(total_size=4096, alignment=64)
+        a = allocator.alloc(64)
+        b = allocator.alloc(64)
+        allocator.free(a, 64)
+        allocator.free(b, 64)  # triggers merge with A
+        # Now free A again — contained within the merged block
+        allocator.free(a, 64)  # should not raise
+        # Pool should still be fully usable
+        offset = allocator.alloc(4096)
+        assert offset == 0
+
+    @pytest.mark.slow
+    def test_partial_overlap_raises_corruption(self):
+        """Freeing a region that partially overlaps a free block must raise RuntimeError."""
+        allocator = BufferAllocator(total_size=4096, alignment=64)
+        a = allocator.alloc(128)
+        b = allocator.alloc(128)
+        allocator.free(a, 128)  # [0, 128) is now free
+        # Try to free a region that starts inside [0,128) but extends beyond
+        with pytest.raises(RuntimeError):
+            allocator.free(64, 128)  # [64, 192) overlaps with free [0, 128)
+
+        # b is still allocated; freeing b should be fine
+        allocator.free(b, 128)
+
+    @pytest.mark.slow
+    def test_merge_adjacent_blocks(self):
+        """Free three adjacent blocks; they should merge into one contiguous region."""
+        allocator = BufferAllocator(total_size=4096, alignment=64)
+        a = allocator.alloc(64)
+        b = allocator.alloc(64)
+        c = allocator.alloc(64)
+        # Free in non-sequential order to exercise sorting + merging
+        allocator.free(b, 64)
+        allocator.free(a, 64)
+        allocator.free(c, 64)
+        # After merge, free_blocks should contain one block starting at 0
+        # covering at least 192 bytes (3 * 64).
+        # Verify by allocating a contiguous block of 192 bytes.
+        offset = allocator.alloc(192)
+        assert offset == 0, "Adjacent blocks were not merged properly"
+
+    @pytest.mark.slow
+    def test_fragmentation_and_defrag(self):
+        """
+        Allocate A B C D that exactly fill the pool, free B and D to
+        create fragmentation, verify a large contiguous alloc fails,
+        then free A and C — should result in full defrag.
+        """
+        # Total pool = 4 * 64 = 256 bytes, so 4 allocs exhaust it completely
+        allocator = BufferAllocator(total_size=256, alignment=64)
+        a = allocator.alloc(64)
+        b = allocator.alloc(64)
+        c = allocator.alloc(64)
+        d = allocator.alloc(64)
+
+        allocator.free(b, 64)  # free blocks: [64, 128)
+        allocator.free(d, 64)  # free blocks: [64, 128) and [192, 256)
+
+        # Pool has two 64-byte holes; contiguous 128 is unavailable
+        with pytest.raises(MemoryError):
+            allocator.alloc(128)
+
+        allocator.free(a, 64)
+        allocator.free(c, 64)
+
+        # After freeing everything, full pool should be available
+        offset = allocator.alloc(256)
+        assert offset == 0
+
+
+@pytest.mark.core_model
+class TestManagedBuffer:
+    """Unit tests for ManagedBuffer."""
+
+    @pytest.fixture(autouse=True)
+    def _setup(self):
+        # automatically invoked for every test method in the class
+        self.allocator = BufferAllocator(total_size=4096, alignment=64)
+        self.pool = torch.zeros(4096, dtype=torch.uint8)
+
+    def test_tensor_view(self):
+        """Verify tensor property and as_tensor return correct views."""
+        offset = self.allocator.alloc(64)
+        buf = ManagedBuffer(self.allocator, offset, 64, self.pool)
+
+        # Write float32 data via pool
+        src = torch.arange(16, dtype=torch.float32)
+        self.pool[offset : offset + 64] = src.view(torch.uint8)
+
+        # Raw uint8 view
+        assert buf.tensor.shape[0] == 64
+
+        # Typed view
+        typed = buf.as_tensor(dtype=torch.float32, shape=(4, 4))
+        assert typed.shape == (4, 4)
+        assert torch.equal(typed.flatten(), src)
+
+        buf.release()
+
+    def test_context_manager_releases_buffer(self):
+        """Verify context manager releases buffer and space is reusable."""
+        offset = self.allocator.alloc(128)
+
+        with ManagedBuffer(self.allocator, offset, 128, self.pool) as buf:
+            assert not buf._released
+
+        assert buf._released
+
+        # Space should be reusable
+        new_offset = self.allocator.alloc(128)
+        assert new_offset == offset
@@ -0,0 +1,754 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Integration tests for MooncakeTransferEngineConnector.
+
+These tests require Mooncake TransferEngine and an RDMA environment.
+Set RDMA_DEVICE_NAME to force a specific device for single-node testing.
+"""
+
+import hashlib
+import os
+import subprocess
+import threading
+import time
+
+import pytest
+import torch
+
+from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
+    ManagedBuffer,
+    MooncakeTransferEngineConnector,
+    TransferEngine,
+)
+
+# All tests in this file require Mooncake TransferEngine and an RDMA environment.
+pytestmark = [pytest.mark.parallel, pytest.mark.gpu]
+
+# ---------------------------------------------------------------------------
+# Shared helpers
+# ---------------------------------------------------------------------------
+
+
+def get_rdma_host() -> str:
+    """Resolve RDMA-capable host IP (env > auto-detect IB > hostname > 127.0.0.1)."""
+    import socket
+
+    env_host = os.environ.get("RDMA_TEST_HOST")
+    if env_host:
+        return env_host
+
+    for ip_cmd in ("ip", "/sbin/ip", "/usr/sbin/ip"):
+        try:
+            result = subprocess.run([ip_cmd, "addr", "show"], capture_output=True, text=True, timeout=5)
+            if result.returncode != 0:
+                continue
+            lines = result.stdout.split("\n")
+            rdma_patterns = ("ibp", "ib0", "ib1", "ib2", "mlx", "roce")
+            for line in lines:
+                if "inet " in line:
+                    for pat in rdma_patterns:
+                        if pat in line.lower():
+                            ip = line.strip().split()[1].split("/")[0]
+                            if not ip.startswith("127.") and not ip.startswith("169.254."):
+                                return ip
+            for line in lines:
+                if "inet " in line and "scope global" in line:
+                    ip = line.strip().split()[1].split("/")[0]
+                    if not ip.startswith("127.") and not ip.startswith("172.17."):
+                        return ip
+            break
+        except (FileNotFoundError, Exception):
+            continue
+
+    try:
+        host_ip = socket.gethostbyname(socket.gethostname())
+        if host_ip and not host_ip.startswith("127."):
+            return host_ip
+    except Exception:
+        pass  # Hostname resolution failed; fall back to loopback below
+    return "127.0.0.1"
+
+
+def _detect_rdma_device() -> str:
+    """Auto-detect a usable RDMA device for single-node loopback testing.
+
+    On DGX machines with 12+ RDMA NICs, only RoCE NICs (with a bound network
+    interface) can reliably do loopback.  IB-only NICs often fail cross-NIC
+    QP handshakes, causing intermittent test failures.
+
+    Priority:
+      1. RDMA_DEVICE_NAME / RDMA_DEVICE environment variable
+      2. ibdev2netdev: pick first RoCE device whose netdev is UP
+      3. Mooncake TransferEngine topology (if available)
+      4. Empty string (Mooncake uses all devices — risky on multi-NIC)
+    """
+    # 1. Explicit env override
+    for env_key in ("RDMA_DEVICE_NAME", "RDMA_DEVICE"):
+        env_val = os.environ.get(env_key, "").strip()
+        if env_val:
+            print(f"[_detect_rdma_device] Using {env_key}={env_val}")
+            return env_val
+
+    # 2. ibdev2netdev: most reliable, no Mooncake dependency
+    #    Prefer RoCE (Ethernet-backed) over IB for loopback reliability.
+    try:
+        result = subprocess.run(["ibdev2netdev"], capture_output=True, text=True, timeout=5)
+        if result.returncode == 0:
+            up_lines = [ln.strip() for ln in result.stdout.strip().splitlines() if "(Up)" in ln]
+            # Prefer Ethernet-backed (RoCE) — interface name is NOT ib*
+            for line in up_lines:
+                # Format: "mlx5_2 port 1 ==> enp75s0f0 (Up)"
+                parts = line.split("==>")
+                if len(parts) == 2:
+                    iface = parts[1].strip().split()[0]  # e.g. "enp75s0f0"
+                    if not iface.startswith("ib"):
+                        dev = line.split()[0]
+                        print(f"[_detect_rdma_device] ibdev2netdev selected RoCE: {dev} ({line})")
+                        return dev
+            # Fallback: any Up device (IB)
+            if up_lines:
+                dev = up_lines[0].split()[0]
+                print(f"[_detect_rdma_device] ibdev2netdev selected IB: {dev} ({up_lines[0]})")
+                return dev
+            print("[_detect_rdma_device] ibdev2netdev: no device with (Up) link found")
+    except FileNotFoundError:
+        print("[_detect_rdma_device] ibdev2netdev not found, skipping")
+    except Exception as exc:
+        print(f"[_detect_rdma_device] ibdev2netdev failed: {exc}")
+
+    # 3. Mooncake topology query
+    if TransferEngine is not None:
+        try:
+            import json
+
+            host = get_rdma_host()
+            eng = TransferEngine()
+            if eng.initialize(host, "P2PHANDSHAKE", "rdma", "") == 0:
+                topo_str = getattr(eng, "get_local_topology", lambda: None)()
+                if topo_str:
+                    topo = json.loads(topo_str)
+                    # Prefer RoCE (IPv4-mapped GID)
+                    for name, info in topo.items():
+                        if isinstance(info, dict):
+                            gid = info.get("gid", "")
+                            if gid.startswith("00:00:00:00:00:00:00:00:00:00:ff:ff"):
+                                print(f"[_detect_rdma_device] Mooncake topology selected RoCE: {name}")
+                                return name
+                    # Prefer bonded device
+                    for name in topo:
+                        if "bond" in name:
+                            print(f"[_detect_rdma_device] Mooncake topology selected bond: {name}")
+                            return name
+                    # First available
+                    for name, info in topo.items():
+                        if isinstance(info, dict):
+                            print(f"[_detect_rdma_device] Mooncake topology fallback: {name}")
+                            return name
+                else:
+                    print("[_detect_rdma_device] get_local_topology() returned None/empty")
+        except Exception as exc:
+            print(f"[_detect_rdma_device] Mooncake topology query failed: {exc}")
+
+    print("[_detect_rdma_device] WARNING: no device detected, Mooncake will use ALL NICs")
+    return ""
+
+
+RDMA_HOST = get_rdma_host()
+RDMA_DEVICE = _detect_rdma_device()
+
+
+def _free_port() -> int:
+    import socket
+
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind((RDMA_HOST, 0))
+        return s.getsockname()[1]
+
+
+def _connector_config(
+    zmq_port: int,
+    pool_size: int = 16 * 1024 * 1024,
+    pool_device: str = "cpu",
+) -> dict:
+    cfg = {
+        "host": RDMA_HOST,
+        "zmq_port": zmq_port,
+        "protocol": "rdma",
+        "memory_pool_size": pool_size,
+        "memory_pool_device": pool_device,
+    }
+    if RDMA_DEVICE:
+        cfg["device_name"] = RDMA_DEVICE
+    return cfg
+
+
+def _md5(tensor: torch.Tensor) -> str:
+    t = tensor.cpu() if tensor.is_cuda else tensor
+    return hashlib.md5(t.contiguous().view(torch.uint8).numpy().tobytes()).hexdigest()
+
+
+# ---------------------------------------------------------------------------
+# 1. Basic functionality (single connector, no RDMA transfer)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+class TestBasicConnector:
+    """Verify connector initialization, put, cleanup, and health check."""
+
+    def test_initialization(self):
+        port = _free_port()
+        with MooncakeTransferEngineConnector(_connector_config(port, pool_size=1024 * 1024)) as c:
+            assert c.rpc_port != 0
+            assert c.pool_size == 1024 * 1024
+            assert c.pool.is_pinned()
+            health = c.health()
+            assert health["status"] == "healthy"
+
+    def test_put_tensor_bytes_object(self):
+        """Put tensor / bytes / dict and verify metadata."""
+        port = _free_port()
+        with MooncakeTransferEngineConnector(_connector_config(port)) as c:
+            ok, sz, meta = c.put("s0", "s1", "t", torch.randn(100))
+            assert ok
+            assert meta["is_fast_path"]
+
+            ok, sz, meta = c.put("s0", "s1", "b", b"hello" * 100)
+            assert ok
+            assert meta["is_fast_path"]
+
+            ok, sz, meta = c.put("s0", "s1", "d", {"k": [1, 2, 3]})
+            assert ok
+            assert not meta["is_fast_path"]
+
+    def test_cleanup_releases_buffer(self):
+        port = _free_port()
+        with MooncakeTransferEngineConnector(_connector_config(port)) as c:
+            c.put("s0", "s1", "r1", torch.randn(100))
+            key = MooncakeTransferEngineConnector._make_key("r1", "s0", "s1")
+            assert key in c._local_buffers
+            c.cleanup("r1", from_stage="s0", to_stage="s1")
+            assert key not in c._local_buffers
+
+    def test_pool_exhaustion_and_recovery(self):
+        """Fill pool, verify failure, free, verify recovery."""
+        port = _free_port()
+        with MooncakeTransferEngineConnector(_connector_config(port, pool_size=64 * 1024)) as c:
+            ids = []
+            for i in range(10):
+                ok, _, _ = c.put("s", "s", f"f{i}", torch.randn(1000))
+                if ok:
+                    ids.append(f"f{i}")
+                else:
+                    break
+            for rid in ids:
+                c.cleanup(rid, from_stage="s", to_stage="s")
+            ok, _, _ = c.put("s", "s", "recovery", torch.randn(1000))
+            assert ok, "Pool recovery failed after cleanup"
+
+
+# ---------------------------------------------------------------------------
+# 2. End-to-end RDMA transfer (producer + consumer, single node)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+class TestEndToEnd:
+    """E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload, mixed types."""
+
+    def _pair(self, pool_size=16 * 1024 * 1024):
+        p = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
+        c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
+        return p, c
+
+    def test_tensor_e2e(self):
+        p, c = self._pair()
+        try:
+            orig = torch.randn(1024, 1024, dtype=torch.float32)
+            ok, sz, meta = p.put("s0", "s1", "t1", orig)
+            assert ok
+            time.sleep(0.5)
+            result = c.get("s0", "s1", "t1", meta)
+            assert result is not None
+            buf, _ = result
+            assert isinstance(buf, ManagedBuffer)
+            recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
+            assert torch.equal(recon, orig)
+            buf.release()
+        finally:
+            p.close()
+            c.close()
+
+    def test_bytes_e2e(self):
+        p, c = self._pair()
+        try:
+            orig = b"Hello RDMA! " * 1000
+            ok, _, meta = p.put("s0", "s1", "b1", orig)
+            assert ok
+            time.sleep(0.5)
+            result = c.get("s0", "s1", "b1", meta)
+            assert result is not None
+            buf, _ = result
+            recv = buf.to_bytes() if hasattr(buf, "to_bytes") else buf
+            if hasattr(buf, "release"):
+                buf.release()
+            assert recv == orig
+        finally:
+            p.close()
+            c.close()
+
+    def test_object_e2e(self):
+        p, c = self._pair()
+        try:
+            orig = {"msg": "hello", "nums": [1, 2, 3], "nested": {"k": "v"}}
+            ok, _, meta = p.put("s0", "s1", "o1", orig)
+            assert ok
+            assert not meta["is_fast_path"]
+            time.sleep(1.0)
+            result = c.get("s0", "s1", "o1", meta)
+            assert result is not None
+            recv, _ = result
+            assert recv == orig
+        finally:
+            p.close()
+            c.close()
+
+    def test_zero_copy_e2e(self):
+        p, c = self._pair()
+        try:
+            shape = (1024, 1024)
+            nbytes = 1024 * 1024 * 4
+            offset = p.allocator.alloc(nbytes)
+            sbuf = ManagedBuffer(p.allocator, offset, nbytes, p.pool)
+            ref = torch.randn(*shape, dtype=torch.float32)
+            sbuf.as_tensor(dtype=torch.float32, shape=shape).copy_(ref)
+
+            ok, sz, meta = p.put("s0", "s1", "zc", sbuf)
+            assert ok
+            assert meta.get("is_fast_path")
+            time.sleep(0.5)
+
+            result = c.get("s0", "s1", "zc", meta)
+            assert result is not None
+            rbuf, _ = result
+            recon = rbuf.as_tensor(dtype=torch.float32, shape=shape)
+            assert torch.equal(recon, ref)
+            rbuf.release()
+        finally:
+            p.close()
+            c.close()
+
+    def test_large_tensor_100mb(self):
+        """Transfer ~100 MB tensor and verify MD5 integrity."""
+        p, c = self._pair(pool_size=128 * 1024 * 1024)
+        try:
+            orig = torch.randn(5000, 5000, dtype=torch.float32)
+            md5_send = _md5(orig)
+            ok, _, meta = p.put("s0", "s1", "lg", orig)
+            assert ok
+            time.sleep(1.0)
+            result = c.get("s0", "s1", "lg", meta)
+            assert result is not None
+            buf, _ = result
+            recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
+            assert _md5(recon) == md5_send
+            buf.release()
+        finally:
+            p.close()
+            c.close()
+
+    def test_mixed_types_sequential(self):
+        """Transfer different dtypes and data types in sequence."""
+        p, c = self._pair(pool_size=32 * 1024 * 1024)
+        try:
+            cases = [
+                ("f32", torch.randn(500, 500, dtype=torch.float32)),
+                ("f16", torch.randn(500, 500, dtype=torch.float16)),
+                ("i64", torch.randint(0, 100, (200, 200), dtype=torch.int64)),
+                ("byte", b"X" * (1024 * 1024)),
+                ("obj", {"key": "value", "list": [1, 2, 3]}),
+            ]
+            time.sleep(0.5)
+            for name, data in cases:
+                ok, _, meta = p.put("s0", "s1", name, data)
+                assert ok, f"put failed: {name}"
+                time.sleep(0.3)
+                result = c.get("s0", "s1", name, meta)
+                assert result is not None, f"get failed: {name}"
+                recv, _ = result
+                if isinstance(data, torch.Tensor):
+                    assert isinstance(recv, ManagedBuffer)
+                    recon = recv.as_tensor(dtype=data.dtype, shape=data.shape)
+                    assert torch.equal(recon, data), f"mismatch: {name}"
+                    recv.release()
+                elif isinstance(data, bytes):
+                    got = recv.to_bytes() if hasattr(recv, "to_bytes") else recv
+                    if hasattr(recv, "release"):
+                        recv.release()
+                    assert got == data
+                else:
+                    assert recv == data
+        finally:
+            p.close()
+            c.close()
+
+    def test_concurrent_put(self):
+        """10 concurrent puts should all succeed."""
+        port = _free_port()
+        conn = MooncakeTransferEngineConnector(_connector_config(port, pool_size=64 * 1024 * 1024))
+        errors: list[str] = []
+        lock = threading.Lock()
+
+        def worker(rid, data):
+            try:
+                ok, _, _ = conn.put("s0", "s1", rid, data)
+                if not ok:
+                    with lock:
+                        errors.append(f"{rid}: put failed")
+            except Exception as e:
+                with lock:
+                    errors.append(f"{rid}: {e}")
+
+        with conn:
+            threads = []
+            for i in range(10):
+                t = threading.Thread(target=worker, args=(f"r{i}", torch.randn(256, 256)))
+                threads.append(t)
+                t.start()
+            for t in threads:
+                t.join()
+            assert len(errors) == 0, f"errors: {errors}"
+
+    def test_auto_cleanup(self):
+        """Producer buffer should be released after consumer get."""
+        p, c = self._pair()
+        try:
+            ok, _, meta = p.put("s0", "s1", "ac", torch.randn(100))
+            assert ok
+            key = MooncakeTransferEngineConnector._make_key("ac", "s0", "s1")
+            assert key in p._local_buffers
+            time.sleep(0.5)
+            result = c.get("s0", "s1", "ac", meta)
+            assert result is not None
+            buf, _ = result
+            if isinstance(buf, ManagedBuffer):
+                buf.release()
+            time.sleep(0.3)
+            assert key not in p._local_buffers
+        finally:
+            p.close()
+            c.close()
+
+
+# ---------------------------------------------------------------------------
+# 3. Lifecycle & resource management
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+class TestLifecycle:
+    """Close, context manager, double-close safety."""
+
+    def test_close_releases_resources(self):
+        c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size=1024 * 1024))
+        c.put("s0", "s1", "x", torch.randn(100))
+        c.close()
+        assert c._stop_event.is_set()
+        assert len(c._local_buffers) == 0
+
+    def test_context_manager(self):
+        with MooncakeTransferEngineConnector(_connector_config(_free_port())) as c:
+            ok, _, _ = c.put("s0", "s1", "ctx", torch.randn(50))
+            assert ok
+        assert c._stop_event.is_set()
+
+    def test_double_close_safe(self):
+        c = MooncakeTransferEngineConnector(_connector_config(_free_port()))
+        c.close()
+        c.close()
+
+
+# ---------------------------------------------------------------------------
+# 4. GPU memory pool (requires CUDA)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.cuda
+@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+class TestGPUPool:
+    """GPU memory pool: initialization, put (CPU/GPU tensor), E2E transfer."""
+
+    @staticmethod
+    def _gpu_cfg(port, pool_size=32 * 1024 * 1024):
+        return _connector_config(port, pool_size, pool_device="cuda:0")
+
+    def test_gpu_pool_init(self):
+        with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c:
+            assert c.pool_device == "cuda:0"
+            assert c.pool.is_cuda
+
+    def test_gpu_pool_put_cpu_and_gpu_tensor(self):
+        with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c:
+            ok, _, meta = c.put("s0", "s1", "h2d", torch.randn(256, 256))
+            assert ok
+            assert meta["is_fast_path"]
+
+            ok, _, meta = c.put("s0", "s1", "d2d", torch.randn(256, 256, device="cuda:0"))
+            assert ok
+            assert meta["is_fast_path"]
+
+    def test_gpu_e2e_transfer(self):
+        p = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port()))
+        c = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port()))
+        try:
+            orig = torch.randn(512, 512, dtype=torch.float32, device="cuda:0")
+            ok, _, meta = p.put("s0", "s1", "ge", orig)
+            assert ok
+            time.sleep(0.5)
+            result = c.get("s0", "s1", "ge", meta)
+            assert result is not None
+            buf, _ = result
+            recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
+            assert recon.is_cuda
+            assert torch.equal(recon.cpu(), orig.cpu())
+            buf.release()
+        finally:
+            p.close()
+            c.close()
+
+
+# ---------------------------------------------------------------------------
+# 5. Stress / Correctness tests (marked slow, skipped in quick CI)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+class TestStressCorrectness:
+    """
+    Slow but high-value regression tests: concurrent put+get with data
+    integrity, large payload, edge cases, and sustained stress.
+    """
+
+    def _pair(self, pool_size=64 * 1024 * 1024):
+        p = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
+        c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size))
+        return p, c
+
+    # -- Concurrent put + get with data integrity --
+
+    def test_concurrent_put_get_integrity(self):
+        """
+        10 threads each put a unique tensor; a single consumer thread gets all
+        of them and verifies MD5 matches.
+        """
+        p, c = self._pair(pool_size=128 * 1024 * 1024)
+        num_workers = 10
+        errors: list[str] = []
+        lock = threading.Lock()
+        md5_sent: dict[str, str] = {}
+        meta_map: dict[str, dict] = {}
+
+        def producer_worker(rid: str, tensor: torch.Tensor):
+            try:
+                md5_val = _md5(tensor)
+                with lock:
+                    md5_sent[rid] = md5_val
+                ok, _, meta = p.put("s0", "s1", rid, tensor)
+                if not ok:
+                    with lock:
+                        errors.append(f"{rid}: put failed")
+                    return
+                with lock:
+                    meta_map[rid] = meta
+            except Exception as e:
+                with lock:
+                    errors.append(f"{rid}: put exception {e}")
+
+        try:
+            threads = []
+            for i in range(num_workers):
+                t = threading.Thread(
+                    target=producer_worker,
+                    args=(f"cg{i}", torch.randn(512, 512, dtype=torch.float32)),
+                )
+                threads.append(t)
+                t.start()
+            for t in threads:
+                t.join()
+
+            assert len(errors) == 0, f"Producer errors: {errors}"
+            assert len(meta_map) == num_workers
+
+            time.sleep(1.0)
+
+            # Consumer: get each tensor and verify integrity
+            for rid, meta in meta_map.items():
+                result = c.get("s0", "s1", rid, meta)
+                assert result is not None, f"get returned None for {rid}"
+                buf, _ = result
+                assert isinstance(buf, ManagedBuffer), f"{rid} not ManagedBuffer"
+                recon = buf.as_tensor(dtype=torch.float32, shape=(512, 512))
+                recv_md5 = _md5(recon)
+                assert recv_md5 == md5_sent[rid], f"MD5 mismatch for {rid}"
+                buf.release()
+        finally:
+            p.close()
+            c.close()
+
+    def test_concurrent_put_get_threaded_both_sides(self):
+        """
+        True overlapping concurrency: producer threads put items one by
+        one and signal the consumer via a queue; the consumer thread
+        starts getting as soon as each item is available, so put and get
+        execute in overlapping time windows.
+        """
+        import queue
+
+        p, c = self._pair(pool_size=128 * 1024 * 1024)
+        num_items = 8
+        errors: list[str] = []
+        lock = threading.Lock()
+        md5_sent: dict[str, str] = {}
+        item_queue: queue.Queue[tuple[str, dict] | None] = queue.Queue()
+
+        def producer_worker(rid, tensor):
+            try:
+                with lock:
+                    md5_sent[rid] = _md5(tensor)
+                ok, _, meta = p.put("s0", "s1", rid, tensor)
+                if not ok:
+                    with lock:
+                        errors.append(f"{rid}: put failed")
+                    return
+                # Signal consumer immediately after put
+                item_queue.put((rid, meta))
+            except Exception as e:
+                with lock:
+                    errors.append(f"{rid}: {e}")
+
+        consumed_count = [0]  # mutable container for thread access
+
+        def consumer_thread():
+            while consumed_count[0] < num_items:
+                try:
+                    item = item_queue.get(timeout=30)
+                except queue.Empty:
+                    with lock:
+                        errors.append(f"consumer: queue.Empty after 30s, consumed {consumed_count[0]}/{num_items}")
+                    break
+                if item is None:
+                    break
+                rid, meta = item
+                try:
+                    # Small delay to let RDMA transfer initiate, but
+                    # overlap with other producers that are still putting
+                    time.sleep(0.1)
+                    result = c.get("s0", "s1", rid, meta)
+                    if result is None:
+                        with lock:
+                            errors.append(f"{rid}: get returned None")
+                        continue
+                    buf, _ = result
+                    if isinstance(buf, ManagedBuffer):
+                        recon = buf.as_tensor(dtype=torch.float32, shape=(256, 256))
+                        if _md5(recon) != md5_sent.get(rid, ""):
+                            with lock:
+                                errors.append(f"{rid}: MD5 mismatch")
+                        buf.release()
+                except Exception as e:
+                    with lock:
+                        errors.append(f"{rid}: get exception {e}")
+                finally:
+                    consumed_count[0] += 1
+
+        try:
+            # Start consumer first so it is ready to get
+            ct = threading.Thread(target=consumer_thread)
+            ct.start()
+            # Launch producers with staggered starts for overlap
+            threads = []
+            for i in range(num_items):
+                t = threading.Thread(
+                    target=producer_worker,
+                    args=(f"bi{i}", torch.randn(256, 256, dtype=torch.float32)),
+                )
+                threads.append(t)
+                t.start()
+                time.sleep(0.05)  # slight stagger so puts overlap with gets
+            for t in threads:
+                t.join()
+            ct.join(timeout=60)
+            assert len(errors) == 0, f"errors: {errors}"
+            assert consumed_count[0] == num_items, f"Consumer only processed {consumed_count[0]}/{num_items} items"
+        finally:
+            p.close()
+            c.close()
+
+    # -- Edge cases --
+
+    def test_small_tensor_1_element(self):
+        """Transfer a single-element tensor."""
+        p, c = self._pair()
+        try:
+            orig = torch.tensor([3.14], dtype=torch.float32)
+            ok, _, meta = p.put("s0", "s1", "tiny", orig)
+            assert ok
+            time.sleep(0.5)
+            result = c.get("s0", "s1", "tiny", meta)
+            assert result is not None
+            buf, _ = result
+            recon = buf.as_tensor(dtype=torch.float32, shape=(1,))
+            assert torch.equal(recon, orig)
+            buf.release()
+        finally:
+            p.close()
+            c.close()
+
+    def test_empty_bytes_rejected(self):
+        """Connector should gracefully reject empty bytes payload."""
+        port = _free_port()
+        with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c:
+            ok, sz, meta = c.put("s0", "s1", "empty_b", b"")
+            assert not ok, "Empty bytes should be rejected by connector"
+
+    # -- Large payload stress (500 MB) --
+
+    def test_large_tensor_500mb(self):
+        """Transfer ~500 MB tensor and verify MD5 integrity."""
+        p, c = self._pair(pool_size=600 * 1024 * 1024)
+        try:
+            # ~500 MB: 11180 * 11180 * 4 bytes ≈ 500 MB
+            side = 11180
+            orig = torch.randn(side, side, dtype=torch.float32)
+            md5_send = _md5(orig)
+            ok, _, meta = p.put("s0", "s1", "500mb", orig)
+            assert ok
+            time.sleep(2.0)
+            result = c.get("s0", "s1", "500mb", meta)
+            assert result is not None
+            buf, _ = result
+            recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
+            assert _md5(recon) == md5_send
+            buf.release()
+        finally:
+            p.close()
+            c.close()
+
+    # -- Allocator stress under connector --
+
+    def test_rapid_alloc_free_cycle(self):
+        """Put + cleanup in tight loop to stress allocator under real connector."""
+        port = _free_port()
+        with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c:
+            for i in range(50):
+                rid = f"cycle_{i}"
+                ok, _, _ = c.put("s0", "s1", rid, torch.randn(256, 256))
+                assert ok, f"Put failed at iteration {i}"
+                c.cleanup(rid, from_stage="s0", to_stage="s1")
+            # After all cycles, pool should be fully recovered
+            ok, _, _ = c.put("s0", "s1", "final", torch.randn(1024, 1024))
+            assert ok, "Pool not fully recovered after rapid cycles"
@@ -12,11 +12,11 @@ Strategy:
   and correctness of the model definition itself, independent of the pipeline mocks.
 """

-from unittest.mock import MagicMock, patch
-
 import pytest
 import torch
+from pytest_mock import MockerFixture

+from tests.utils import hardware_test
 from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig

 # Mock the OvisImageTransformer2DModel to avoid complex init if needed,
@@ -29,15 +29,15 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams


@pytest.fixture
-def mock_dependencies(monkeypatch):
+def mock_dependencies(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
    """
    Mock external dependencies to avoid loading real models.
    """
    device = get_local_device()

    # Mock Tokenizer
-    mock_tokenizer = MagicMock()
-    mock_tokenizer.return_value = MagicMock(
+    mock_tokenizer = mocker.MagicMock()
+    mock_tokenizer.return_value = mocker.MagicMock(
        input_ids=torch.zeros((1, 50), dtype=torch.long, device=device),
        attention_mask=torch.ones((1, 50), dtype=torch.long, device=device),
    )
@@ -45,13 +45,13 @@ def mock_dependencies(monkeypatch):
    mock_tokenizer.model_max_length = 1024

    # Mock Text Encoder
-    mock_text_encoder = MagicMock()
+    mock_text_encoder = mocker.MagicMock()
    mock_text_encoder.dtype = torch.float32
    # Output of text encoder must be on the same device as inputs (which are moved to execution_device)
    mock_text_encoder.return_value.last_hidden_state = torch.randn(1, 50, 32, device=device)

    # Mock VAE
-    mock_vae = MagicMock()
+    mock_vae = mocker.MagicMock()
    mock_vae.config.block_out_channels = [128, 256, 512, 512]  # Scale factor 8
    mock_vae.config.scale_factor_temporal = 1
    mock_vae.config.scale_factor_spatial = 8
@@ -63,8 +63,8 @@ def mock_dependencies(monkeypatch):
    mock_vae.to.return_value = mock_vae

    # Mock Scheduler
-    mock_scheduler = MagicMock()
-    mock_scheduler.config = MagicMock()
+    mock_scheduler = mocker.MagicMock()
+    mock_scheduler.config = mocker.MagicMock()
    # Timesteps on device to match latents during denoising loop interaction if needed
    mock_scheduler.timesteps = torch.tensor([1.0, 0.5, 0.0], device=device)
    mock_scheduler.set_timesteps.return_value = None
@@ -95,7 +95,7 @@ def mock_dependencies(monkeypatch):


@pytest.fixture
-def ovis_pipeline(mock_dependencies, monkeypatch):
+def ovis_pipeline(mocker: MockerFixture, mock_dependencies, monkeypatch: pytest.MonkeyPatch):
    """
    Creates an OvisImagePipeline instance with mocked components.
    """
@@ -122,8 +122,8 @@ def ovis_pipeline(mock_dependencies, monkeypatch):

    # Mock Transformer Layer separately to avoid full init
    # We patch OvisImageTransformer2DModel class in the module
-    mock_transformer_cls = MagicMock()
-    mock_transformer_instance = MagicMock()
+    mock_transformer_cls = mocker.MagicMock()
+    mock_transformer_instance = mocker.MagicMock()
    mock_transformer_instance.dtype = torch.float32
    mock_transformer_instance.in_channels = 16  # Must be 16 so num_channel_latents=4, packed=16
    # Forward return: noise prediction
@@ -144,8 +144,8 @@ def ovis_pipeline(mock_dependencies, monkeypatch):

    # Initialize pipeline
    # We use a dummy model path check override
-    with patch("os.path.exists", return_value=True):
-        pipeline = OvisImagePipeline(od_config=od_config)
+    mocker.patch("os.path.exists", return_value=True)
+    pipeline = OvisImagePipeline(od_config=od_config)

    return pipeline

@@ -223,9 +223,11 @@ def test_resolution_check(ovis_pipeline):
    assert output is not None


-def test_real_transformer_init_and_forward():
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+def test_real_transformer_init_and_forward(mocker: MockerFixture):
    """Test the real OvisImageTransformer2DModel initialization and forward pass for coverage."""
-    from unittest.mock import patch

    from vllm_omni.diffusion.models.ovis_image.ovis_image_transformer import OvisImageTransformer2DModel

@@ -249,42 +251,42 @@ def test_real_transformer_init_and_forward():

    # Mock distributed state for QKVParallelLinear initialization
    # We patch get_tp_group because get_tensor_model_parallel_rank calls it and asserts _TP is not None
-    mock_group = MagicMock()
+    mock_group = mocker.MagicMock()
    mock_group.rank_in_group = 0
    mock_group.world_size = 1

-    with patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group):
-        # Initialize real model
-        model = OvisImageTransformer2DModel(
-            od_config=od_config,
-            patch_size=1,
-            in_channels=16,
-            out_channels=16,
-            num_single_layers=1,
-            attention_head_dim=8,
-            num_attention_heads=2,
-            joint_attention_dim=32,
-            axes_dims_rope=(2, 2, 4),
-        ).to(device)
+    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
+    # Initialize real model
+    model = OvisImageTransformer2DModel(
+        od_config=od_config,
+        patch_size=1,
+        in_channels=16,
+        out_channels=16,
+        num_single_layers=1,
+        attention_head_dim=8,
+        num_attention_heads=2,
+        joint_attention_dim=32,
+        axes_dims_rope=(2, 2, 4),
+    ).to(device)

-        # Create dummy inputs
-        B, Seq, C = 1, 16, 16
-        hidden_states = torch.randn(B, Seq, C, device=device)
-        encoder_hidden_states = torch.randn(B, 10, 32, device=device)  # joint_attention_dim=32
-        timestep = torch.tensor([1], device=device)
-        img_ids = torch.zeros(Seq, 3, device=device)
-        txt_ids = torch.zeros(10, 3, device=device)
+    # Create dummy inputs
+    B, Seq, C = 1, 16, 16
+    hidden_states = torch.randn(B, Seq, C, device=device)
+    encoder_hidden_states = torch.randn(B, 10, 32, device=device)  # joint_attention_dim=32
+    timestep = torch.tensor([1], device=device)
+    img_ids = torch.zeros(Seq, 3, device=device)
+    txt_ids = torch.zeros(10, 3, device=device)

-        # Run forward
-        output = model(
-            hidden_states=hidden_states,
-            encoder_hidden_states=encoder_hidden_states,
-            timestep=timestep,
-            img_ids=img_ids,
-            txt_ids=txt_ids,
-            return_dict=False,
-        )
+    # Run forward
+    output = model(
+        hidden_states=hidden_states,
+        encoder_hidden_states=encoder_hidden_states,
+        timestep=timestep,
+        img_ids=img_ids,
+        txt_ids=txt_ids,
+        return_dict=False,
+    )

-        assert output is not None
-        assert isinstance(output, tuple)
-        assert output[0].shape == hidden_states.shape
+    assert output is not None
+    assert isinstance(output, tuple)
+    assert output[0].shape == hidden_states.shape
@@ -3,10 +3,6 @@ E2E Online tests for Qwen3-Omni model with video input and audio output.
 """

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

 import pytest
@@ -21,6 +17,10 @@ from tests.conftest import (
 from tests.utils import hardware_test
 from vllm_omni.platforms import current_omni_platform

+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+
 models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]


@@ -114,7 +114,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
        "messages": messages,
        "stream": True,
        "key_words": {
-            "audio": ["water", "chirping"],
+            "audio": ["water", "chirping", "crackling", "rain"],
            "image": ["square", "quadrate"],
        },
    }
@@ -4,66 +4,55 @@
 E2E Online tests for Qwen3-Omni model.
 """

-import concurrent.futures
 import os

 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-import threading
-import time
 from pathlib import Path

-import openai
 import pytest

 from tests.conftest import (
-    OmniServer,
-    convert_audio_to_text,
-    cosine_similarity_text,
    dummy_messages_from_mix_data,
    generate_synthetic_audio,
    generate_synthetic_image,
+    generate_synthetic_video,
+    modify_stage_config,
 )
+from tests.utils import hardware_test

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

-# CI stage config for 2*H100-80G GPUs
-stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")]
+AUDIO_KEY = ["water", "chirping", "crackling", "rain"]
+IMAGE_KEY = ["square", "quadrate"]
+VIDEO_KEY = ["sphere", "globe", "circle", "round", "ball"]

+
+def get_chunk_config(default_path):
+    path = modify_stage_config(
+        default_path,
+        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
+
+
+# CI stage config for 2*H100-80G GPUs
+default_path = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
+stage_configs = [default_path, get_chunk_config(default_path)]
 # Create parameter combinations for model and stage config
 test_params = [(model, stage_config) for model in models for stage_config in stage_configs]


-_omni_server_lock = threading.Lock()
-
-
-@pytest.fixture(scope="module")
-def omni_server(request):
-    """Start vLLM-Omni server as a subprocess with actual model weights.
-    Uses session scope so the server starts only once for the entire test session.
-    Multi-stage initialization can take 10-20+ minutes.
-    """
-    with _omni_server_lock:
-        model, stage_config_path = request.param
-
-        print(f"Starting OmniServer with model: {model}")
-
-        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")
-
-
-@pytest.fixture
-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",
-    )
-
-
 def get_system_prompt():
    return {
        "role": "system",
@@ -83,7 +72,10 @@ def get_system_prompt():
 def get_prompt(prompt_type="text_only"):
    prompts = {
        "text_only": "What is the capital of China?",
-        "mix": "What is recited in the audio? What is in this image? Describe the video briefly.",
+        "mix": "What is recited in the audio? What is in this image? What is in this video?",
+        "text_video": "What is in this video? ",
+        "text_image": "What is in this image? ",
+        "text_audio": "What is in this audio? ",
    }
    return prompts.get(prompt_type, prompts["text_only"])

@@ -93,196 +85,230 @@ def get_max_batch_size(size_type="few"):
    return batch_sizes.get(size_type, 5)


-def get_deploy_config(deploy_type="TP1"):
-    result = {
-        "TP1": {
-            "stage_args": {
-                0: {
-                    "engine_args.gpu_memory_utilization": 0.95,
-                    "engine_args.tensor_parallel_size": 1,
-                    "runtime.devices": "0",
-                },
-                2: {"runtime.devices": "1"},
-            }
-        }
-    }
-    return result.get(deploy_type, result["TP1"])
-
-
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_text_to_text_001(client: openai.OpenAI, omni_server) -> None:
-    """Test processing text, generating text output via OpenAI API."""
+def test_text_to_audio_001(omni_server, openai_client) -> None:
    messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())

-    # Test single completion
-    start_time = time.perf_counter()
-    chat_completion = client.chat.completions.create(
-        model=omni_server.model, messages=messages, max_tokens=20, modalities=["text"]
-    )
-    # Verify E2E
-    print(f"the request e2e is: {time.perf_counter() - start_time}")
-    # TODO: Verify the E2E latency after confirmation baseline.
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "modalities": ["audio"],
+        "key_words": {"text": ["beijing"]},
+    }

-    # Verify only output text
-    assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
-
-    # Verify text output success
-    text_choice = chat_completion.choices[0]
-    assert text_choice.message.content is not None, "No text output is generated"
-    assert chat_completion.usage.completion_tokens <= 20, "The output length more than the requested max_tokens."
-    assert "beijing" in text_choice.message.content.lower(), "The output do not contain keywords."
+    openai_client.send_request(request_config)


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_audio_to_text_001(client: openai.OpenAI, omni_server) -> None:
-    """Test processing text, generating text output via OpenAI API."""
-    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(1, 1)['base64']}"
-    messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
-    # Test single completion
-    start_time = time.perf_counter()
-    chat_completion = client.chat.completions.create(model=omni_server.model, messages=messages, modalities=["text"])
-    # Verify only output text
-    assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
+def test_text_to_text_audio_001(omni_server, openai_client) -> None:
+    messages = dummy_messages_from_mix_data(system_prompt=get_system_prompt(), content_text=get_prompt())

-    # Verify text output success
-    text_choice = chat_completion.choices[0]
-    assert text_choice.message.content is not None, "No text output is generated"
-    assert "water" in text_choice.message.content.lower(), "The output do not contain keywords."
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "key_words": {"text": ["beijing"]},
+    }

-    # Verify E2E
-    print(f"the request e2e is: {time.perf_counter() - start_time}")
-    # TODO: Verify the E2E latency after confirmation baseline.
+    openai_client.send_request(request_config)


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_audio_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
-    """Test processing text, generating audio output via OpenAI API."""
-
-    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(1, 5)['base64']}"
-
-    messages = dummy_messages_from_mix_data(audio_data_url=audio_data_url)
-    num_concurrent_requests = get_max_batch_size()
-    # Test single completion
-    e2e_list = list()
-    with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
-        # Submit multiple completion requests concurrently
-        futures = [
-            executor.submit(client.chat.completions.create, model=omni_server.model, messages=messages)
-            for _ in range(num_concurrent_requests)
-        ]
-        start_time = time.perf_counter()
-        # Wait for all requests to complete and collect results
-        chat_completions = list()
-        for future in concurrent.futures.as_completed(futures):
-            chat_completions.append(future.result())
-            # Verify E2E
-            current_e2e = time.perf_counter() - start_time
-            print(f"the request e2e is: {current_e2e}")
-            # TODO: Verify the E2E latency after confirmation baseline.
-            e2e_list.append(current_e2e)
-
-    print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
-    # Verify all completions succeeded
-    assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
-    for chat_completion in chat_completions:
-        # Verify audio output success
-        audio_message = chat_completion.choices[1].message
-        audio_data = audio_message.audio.data
-        assert audio_data is not None, "No audio output is generated"
-        assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
-
-        # Verify text output success
-        text_choice = chat_completion.choices[0]
-        text_content = text_choice.message.content
-        assert text_choice.message.content is not None, "No text output is generated"
-
-        # Verify text output same as audio output
-        audio_content = convert_audio_to_text(audio_data)
-        print(f"text content is: {text_content}")
-        print(f"audio content is: {audio_content}")
-        similarity = cosine_similarity_text(audio_content, text_content)
-        print(f"similarity between audio and text is: {similarity}")
-        assert similarity > 0.9, "The audio content is not same as the text"
-
-
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_image_to_text_001(client: openai.OpenAI, omni_server) -> None:
-    """Test processing text, generating text output via OpenAI API."""
+def test_image_to_text_001(omni_server, openai_client) -> None:
    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
    messages = dummy_messages_from_mix_data(image_data_url=image_data_url)
-    # Test single completion
-    start_time = time.perf_counter()
-    chat_completion = client.chat.completions.create(
-        model=omni_server.model, messages=messages, max_tokens=100, modalities=["text"]
-    )
-    # Verify E2E
-    print(f"the request e2e is: {time.perf_counter() - start_time}")
-    # TODO: Verify the E2E latency after confirmation baseline.

-    # Verify only output text
-    assert len(chat_completion.choices) == 1, "The generated content includes more than just text."
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "modalities": ["text"],
+        "key_words": {"image": IMAGE_KEY},
+    }

-    # Verify text output success
-    text_choice = chat_completion.choices[0]
-    text_content = text_choice.message.content
-    assert text_content is not None, "No text output is generated"
-    assert chat_completion.usage.completion_tokens <= 100, "The output length more than the requested max_tokens."
-    assert "square" in text_content.lower(), "The output do not contain keywords."
+    openai_client.send_request(request_config)


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
-def test_image_to_text_audio_001(client: openai.OpenAI, omni_server) -> None:
-    """Test processing text, generating audio output via OpenAI API."""
+def test_image_to_audio_001(omni_server, openai_client) -> None:
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+    messages = dummy_messages_from_mix_data(image_data_url=image_data_url)

-    num_concurrent_requests = get_max_batch_size()
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "modalities": ["audio"],
+        "key_words": {"image": IMAGE_KEY},
+    }

+    openai_client.send_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_image_to_text_audio_001(omni_server, openai_client) -> None:
    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"

    messages = dummy_messages_from_mix_data(image_data_url=image_data_url)

-    # Test single completion
-    e2e_list = list()
-    with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
-        # Submit multiple completion requests concurrently
-        futures = [
-            executor.submit(
-                client.chat.completions.create,
-                model=omni_server.model,
-                messages=messages,
-            )
-            for _ in range(num_concurrent_requests)
-        ]
-        start_time = time.perf_counter()
-        # Wait for all requests to complete and collect results
-        chat_completions = list()
-        for future in concurrent.futures.as_completed(futures):
-            chat_completions.append(future.result())
-            # Verify E2E
-            current_e2e = time.perf_counter() - start_time
-            print(f"the request e2e is: {current_e2e}")
-            # TODO: Verify the E2E latency after confirmation baseline.
-            e2e_list.append(current_e2e)
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "key_words": {"image": IMAGE_KEY},
+    }

-    print(f"the avg e2e is: {sum(e2e_list) / len(e2e_list)}")
-    # Verify all completions succeeded
-    assert len(chat_completions) == num_concurrent_requests, "Not all requests succeeded."
-    for chat_completion in chat_completions:
-        # Verify audio output success
-        audio_message = chat_completion.choices[1].message
-        audio_data = audio_message.audio.data
-        assert audio_data is not None, "No audio output is generated"
-        assert audio_message.audio.expires_at > time.time(), "The generated audio has expired."
+    openai_client.send_request(request_config, request_num=get_max_batch_size())

-        # Verify text output success
-        text_choice = chat_completion.choices[0]
-        text_content = text_choice.message.content
-        assert text_content is not None, "No text output is generated"
-        assert "square" in text_content.lower(), "The output do not contain keywords."

-        # Verify text output same as audio output
-        audio_content = convert_audio_to_text(audio_data)
-        print(f"text content is: {text_content}")
-        print(f"audio content is: {audio_content}")
-        similarity = cosine_similarity_text(audio_content, text_content)
-        print(f"similarity between audio and text is: {similarity}")
-        assert similarity > 0.9, "The audio content is not same as the text"
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_video_to_text_001(omni_server, openai_client) -> None:
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
+    messages = dummy_messages_from_mix_data(video_data_url=video_data_url)
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "modalities": ["text"],
+        "key_words": {"video": VIDEO_KEY},
+    }
+
+    openai_client.send_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_video_to_audio_001(omni_server, openai_client) -> None:
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
+    messages = dummy_messages_from_mix_data(video_data_url=video_data_url)
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "modalities": ["audio"],
+        "key_words": {"video": VIDEO_KEY},
+    }
+
+    openai_client.send_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_video_to_text_audio_001(omni_server, openai_client) -> None:
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}"
+
+    messages = dummy_messages_from_mix_data(video_data_url=video_data_url)
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "key_words": {"video": VIDEO_KEY},
+    }
+
+    openai_client.send_request(request_config, request_num=get_max_batch_size())
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        audio_data_url=audio_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_audio")
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "modalities": ["text"],
+        "key_words": {"audio": AUDIO_KEY},
+    }
+
+    openai_client.send_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}"
+
+    messages = dummy_messages_from_mix_data(
+        image_data_url=image_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_image")
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "modalities": ["audio"],
+        "key_words": {"image": IMAGE_KEY},
+    }
+
+    openai_client.send_request(request_config)
+
+
+@pytest.mark.skip(reason="There is a known issue with oom error.")
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(1280, 720, 30)['base64']}"
+
+    messages = dummy_messages_from_mix_data(
+        video_data_url=video_data_url, system_prompt=get_system_prompt(), content_text=get_prompt("text_video")
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "key_words": {"video": VIDEO_KEY},
+    }
+
+    openai_client.send_request(request_config, request_num=get_max_batch_size())
+
+
+@pytest.mark.skip(reason="There is a known issue with shape mismatch error.")
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
+    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']}"
+    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"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY},
+    }
+    openai_client.send_request(request_config, request_num=get_max_batch_size())
@@ -13,10 +13,10 @@ stage_args:
      model_arch: Qwen2_5OmniForConditionalGeneration
      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 2400
-      max_num_batched_tokens: 2400
+      max_model_len: 16384
+      max_num_batched_tokens: 16384
      max_num_seqs: 1
-      gpu_memory_utilization: 0.8
+      gpu_memory_utilization: 0.9
      skip_mm_profiling: true
      enforce_eager: true  # Now we only support eager mode
      trust_remote_code: true
@@ -43,10 +43,10 @@ stage_args:
      model_arch: Qwen2_5OmniForConditionalGeneration
      worker_type: ar
      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      max_model_len: 2400
-      max_num_batched_tokens: 2400
+      max_model_len: 16384
+      max_num_batched_tokens: 16384
      max_num_seqs: 1
-      gpu_memory_utilization: 0.8
+      gpu_memory_utilization: 0.9
      skip_mm_profiling: true
      enforce_eager: true
      trust_remote_code: true
@@ -58,7 +58,7 @@ stage_args:
      temperature: 0.9
      top_p: 0.8
      top_k: 40
-      max_tokens: 128
+      max_tokens: 4096
      seed: 42
      detokenize: True
      repetition_penalty: 1.05
@@ -66,7 +66,7 @@ stage_args:
  - stage_id: 2
    runtime:
      process: true
-      devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      devices: "2"            # Example: use a different GPU than the previous stage; use "0" if single GPU
      max_batch_size: 1
    engine_args:
      model_stage: code2wav
@@ -78,7 +78,7 @@ stage_args:
      trust_remote_code: true
      enable_prefix_caching: false
      engine_output_type: audio
-      max_num_batched_tokens: 4069
+      max_num_batched_tokens: 100000
    engine_input_source: [1]
    final_output: true
    final_output_type: audio
@@ -86,7 +86,7 @@ stage_args:
      temperature: 0.0
      top_p: 1.0
      top_k: -1
-      max_tokens: 128
+      max_tokens: 4096
      seed: 42
      detokenize: True
      repetition_penalty: 1.1
@@ -86,6 +86,7 @@ stage_args:
    max_num_batched_tokens: 100000
    hf_config_name: thinker_config
    async_scheduling: false
+    load_format: dummy
  engine_input_source: [1]
  custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
  final_output: true
@@ -10,11 +10,11 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py.
 import base64
 import io
 from argparse import Namespace
-from unittest.mock import AsyncMock, Mock

 import pytest
 from fastapi.testclient import TestClient
 from PIL import Image
+from pytest_mock import MockerFixture
 from vllm import SamplingParams

 from vllm_omni.entrypoints.openai.image_api_utils import (
@@ -128,11 +128,11 @@ class FakeAsyncOmni:


@pytest.fixture
-def mock_async_diffusion():
+def mock_async_diffusion(mocker: MockerFixture):
    """Mock AsyncOmniDiffusion instance that returns fake images"""
-    mock = Mock()
+    mock = mocker.Mock()
    mock.is_running = True  # For health endpoint
-    mock.check_health = AsyncMock()  # For LLM mode health check
+    mock.check_health = mocker.AsyncMock()  # For LLM mode health check

    async def generate(**kwargs):
        # Return n PIL images wrapped in result object
@@ -143,7 +143,7 @@ def mock_async_diffusion():
        images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
        return MockGenerationResult(images)

-    mock.generate = AsyncMock(side_effect=generate)
+    mock.generate = mocker.AsyncMock(side_effect=generate)
    return mock


@@ -6,27 +6,26 @@ Tests that standard OpenAI API parameters (max_tokens, temperature, etc.)
 are correctly applied to the comprehension stage while preserving YAML defaults.
 """

-from unittest.mock import MagicMock
-
 import pytest
+from pytest_mock import MockerFixture
 from vllm.sampling_params import SamplingParams

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


@pytest.fixture
-def mock_comprehension_stage():
+def mock_comprehension_stage(mocker: MockerFixture):
    """Create a mock comprehension stage with is_comprehension=True."""
-    stage = MagicMock()
+    stage = mocker.MagicMock()
    stage.is_comprehension = True
    stage.model_stage = "comprehension"
    return stage


@pytest.fixture
-def mock_other_stage():
+def mock_other_stage(mocker: MockerFixture):
    """Create a mock non-comprehension stage."""
-    stage = MagicMock()
+    stage = mocker.MagicMock()
    stage.is_comprehension = False
    stage.model_stage = "other"
    return stage
@@ -57,9 +56,15 @@ def default_other_params():


@pytest.fixture
-def mock_engine_client(mock_comprehension_stage, mock_other_stage, default_comprehension_params, default_other_params):
+def mock_engine_client(
+    mock_comprehension_stage,
+    mock_other_stage,
+    default_comprehension_params,
+    default_other_params,
+    mocker: MockerFixture,
+):
    """Create mock engine client with stage_list and default_sampling_params_list."""
-    engine_client = MagicMock()
+    engine_client = mocker.MagicMock()
    engine_client.stage_list = [mock_comprehension_stage, mock_other_stage]
    engine_client.default_sampling_params_list = [
        default_comprehension_params,
@@ -80,9 +85,9 @@ def serving_chat(mock_engine_client):


@pytest.fixture
-def mock_request():
+def mock_request(mocker: MockerFixture):
    """Create a mock request with all OpenAI sampling params set to None."""
-    request = MagicMock()
+    request = mocker.MagicMock()
    # OpenAI standard sampling fields
    request.temperature = None
    request.top_p = None
@@ -294,35 +299,35 @@ def test_get_comprehension_stage_index_finds_first_stage(mock_engine_client):
    assert instance._get_comprehension_stage_index() == 0


-def test_get_comprehension_stage_index_finds_second_stage():
+def test_get_comprehension_stage_index_finds_second_stage(mocker: MockerFixture):
    """Test finding comprehension stage when it's at index 1."""
    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat

    instance = object.__new__(OmniOpenAIServingChat)

-    other = MagicMock()
+    other = mocker.MagicMock()
    other.is_comprehension = False
-    comprehension = MagicMock()
+    comprehension = mocker.MagicMock()
    comprehension.is_comprehension = True

-    instance.engine_client = MagicMock()
+    instance.engine_client = mocker.MagicMock()
    instance.engine_client.stage_list = [other, comprehension]

    assert instance._get_comprehension_stage_index() == 1


-def test_get_comprehension_stage_index_raises_when_not_found():
+def test_get_comprehension_stage_index_raises_when_not_found(mocker: MockerFixture):
    """Test that ValueError is raised when no comprehension stage exists."""
    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat

    instance = object.__new__(OmniOpenAIServingChat)

-    stage1 = MagicMock()
+    stage1 = mocker.MagicMock()
    stage1.is_comprehension = False
-    stage2 = MagicMock()
+    stage2 = mocker.MagicMock()
    stage2.is_comprehension = False

-    instance.engine_client = MagicMock()
+    instance.engine_client = mocker.MagicMock()
    instance.engine_client.stage_list = [stage1, stage2]

    with pytest.raises(ValueError, match="No comprehension stage"):
@@ -1,13 +1,13 @@
 # tests/entrypoints/openai/test_serving_speech.py
 import logging
 from inspect import Signature, signature
-from unittest.mock import MagicMock, patch

 import numpy as np
 import pytest
 import torch
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
+from pytest_mock import MockerFixture

 from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
 from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio, OpenAICreateSpeechRequest
@@ -26,25 +26,24 @@ class TestAudioMixin:
    def audio_mixin(self):
        return AudioMixin()

-    def test_stereo_to_mono_conversion(self, audio_mixin):
+    def test_stereo_to_mono_conversion(self, audio_mixin, mocker: MockerFixture):
        stereo_tensor = np.random.rand(24000, 2).astype(np.float32)
        audio_obj = CreateAudio(audio_tensor=stereo_tensor)

-        with (
-            patch.object(
-                audio_mixin, "_apply_speed_adjustment", side_effect=lambda tensor, speed, sr: (tensor, sr)
-            ) as mock_speed,
-            patch("soundfile.write") as _,
-        ):
-            audio_mixin.create_audio(audio_obj)
+        mock_speed = mocker.patch.object(
+            audio_mixin, "_apply_speed_adjustment", side_effect=lambda tensor, speed, sr: (tensor, sr)
+        )
+        mocker.patch("soundfile.write")

-            # Check that the tensor passed to speed adjustment is mono
-            mock_speed.assert_called_once()
-            adjusted_tensor = mock_speed.call_args[0][0]
-            assert len(adjusted_tensor) == 24000
+        audio_mixin.create_audio(audio_obj)

-    @patch("librosa.effects.time_stretch")
-    def test_speed_adjustment(self, mock_time_stretch, audio_mixin):
+        # Check that the tensor passed to speed adjustment is mono
+        mock_speed.assert_called_once()
+        adjusted_tensor = mock_speed.call_args[0][0]
+        assert len(adjusted_tensor) == 24000
+
+    def test_speed_adjustment(self, audio_mixin, mocker: MockerFixture):
+        mock_time_stretch = mocker.patch("librosa.effects.time_stretch")
        mock_time_stretch.return_value = np.zeros(12000)
        audio_tensor = np.random.rand(24000).astype(np.float32)

@@ -53,8 +52,8 @@ class TestAudioMixin:
        mock_time_stretch.assert_called_with(y=audio_tensor, rate=2.0)
        assert adjusted_audio.shape == (12000,)

-    @patch("soundfile.write")
-    def test_unsupported_format_fallback(self, mock_write, audio_mixin, caplog):
+    def test_unsupported_format_fallback(self, audio_mixin, caplog, mocker: MockerFixture):
+        mock_write = mocker.patch("soundfile.write")
        audio_tensor = np.random.rand(24000).astype(np.float32)
        # Use a format that is not in the list of supported formats
        audio_obj = CreateAudio(audio_tensor=audio_tensor, response_format="vorbis")
@@ -66,50 +65,50 @@ class TestAudioMixin:
        write_kwargs = mock_write.call_args.kwargs
        assert write_kwargs["format"] == "WAV"

-    def test_mono_audio_preservation(self, audio_mixin):
+    def test_mono_audio_preservation(self, audio_mixin, mocker: MockerFixture):
        """Test that mono (1D) audio tensors are processed correctly and passed to writer."""
        mono_tensor = np.random.rand(24000).astype(np.float32)
        audio_obj = CreateAudio(audio_tensor=mono_tensor)

-        with patch("soundfile.write") as mock_write:
-            audio_mixin.create_audio(audio_obj)
+        mock_write = mocker.patch("soundfile.write")
+        audio_mixin.create_audio(audio_obj)

-            mock_write.assert_called_once()
-            # Verify the tensor passed to soundfile.write is the exact 1D tensor
-            output_tensor = mock_write.call_args[0][1]
-            assert output_tensor.ndim == 1
-            assert output_tensor.shape == (24000,)
-            assert np.array_equal(output_tensor, mono_tensor)
+        mock_write.assert_called_once()
+        # Verify the tensor passed to soundfile.write is the exact 1D tensor
+        output_tensor = mock_write.call_args[0][1]
+        assert output_tensor.ndim == 1
+        assert output_tensor.shape == (24000,)
+        assert np.array_equal(output_tensor, mono_tensor)

-    def test_stereo_audio_preservation(self, audio_mixin):
+    def test_stereo_audio_preservation(self, audio_mixin, mocker: MockerFixture):
        """Test that stereo (2D) audio tensors are processed correctly and preserved."""
        stereo_tensor = np.random.rand(24000, 2).astype(np.float32)
        audio_obj = CreateAudio(audio_tensor=stereo_tensor)

-        with patch("soundfile.write") as mock_write:
-            audio_mixin.create_audio(audio_obj)
+        mock_write = mocker.patch("soundfile.write")
+        audio_mixin.create_audio(audio_obj)

-            mock_write.assert_called_once()
-            # Verify the tensor passed to soundfile.write is the exact 2D tensor
-            output_tensor = mock_write.call_args[0][1]
-            assert output_tensor.ndim == 2
-            assert output_tensor.shape == (24000, 2)
-            assert np.array_equal(output_tensor, stereo_tensor)
+        mock_write.assert_called_once()
+        # Verify the tensor passed to soundfile.write is the exact 2D tensor
+        output_tensor = mock_write.call_args[0][1]
+        assert output_tensor.ndim == 2
+        assert output_tensor.shape == (24000, 2)
+        assert np.array_equal(output_tensor, stereo_tensor)

-    def test_speed_adjustment_bypass(self, audio_mixin):
+    def test_speed_adjustment_bypass(self, audio_mixin, mocker: MockerFixture):
        """Test that speed=1.0 bypasses the expensive librosa time stretching."""
        audio_tensor = np.random.rand(24000).astype(np.float32)

-        with patch("librosa.effects.time_stretch") as mock_time_stretch:
-            # speed=1.0 should return immediately without calling librosa
-            result, _ = audio_mixin._apply_speed_adjustment(audio_tensor, speed=1.0, sample_rate=24000)
+        mock_time_stretch = mocker.patch("librosa.effects.time_stretch")
+        # speed=1.0 should return immediately without calling librosa
+        result, _ = audio_mixin._apply_speed_adjustment(audio_tensor, speed=1.0, sample_rate=24000)

-            mock_time_stretch.assert_not_called()
-            assert np.array_equal(result, audio_tensor)
+        mock_time_stretch.assert_not_called()
+        assert np.array_equal(result, audio_tensor)

-    @patch("librosa.effects.time_stretch")
-    def test_speed_adjustment_stereo_handling(self, mock_time_stretch, audio_mixin):
+    def test_speed_adjustment_stereo_handling(self, audio_mixin, mocker: MockerFixture):
        """Test that speed adjustment is attempted on stereo inputs."""
+        mock_time_stretch = mocker.patch("librosa.effects.time_stretch")
        stereo_tensor = np.random.rand(24000, 2).astype(np.float32)
        # Mock return value representing a sped-up version (half length)
        mock_time_stretch.return_value = np.zeros((12000, 2), dtype=np.float32)
@@ -161,22 +160,22 @@ def create_mock_audio_output_for_test(


@pytest.fixture
-def test_app():
+def test_app(mocker: MockerFixture):
    # Mock the engine client
-    mock_engine_client = MagicMock()
+    mock_engine_client = mocker.MagicMock()
    mock_engine_client.errored = False

    async def mock_generate_fn(*args, **kwargs):
        yield create_mock_audio_output_for_test(request_id=kwargs.get("request_id"))

-    mock_engine_client.generate = MagicMock(side_effect=mock_generate_fn)
+    mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_fn)
    mock_engine_client.default_sampling_params_list = [{}]

    # Mock models to have an is_base_model method
-    mock_models = MagicMock()
+    mock_models = mocker.MagicMock()
    mock_models.is_base_model.return_value = True

-    mock_request_logger = MagicMock()
+    mock_request_logger = mocker.MagicMock()

    speech_server = OmniOpenAIServingSpeech(
        engine_client=mock_engine_client,
@@ -186,7 +185,7 @@ def test_app():

    # Patch the signature of create_speech to remove 'raw_request' for FastAPI route introspection
    original_create_speech = speech_server.create_speech
-    _ = MagicMock(side_effect=original_create_speech)
+    _ = mocker.MagicMock(side_effect=original_create_speech)

    sig = signature(original_create_speech)

@@ -253,11 +252,13 @@ class TestSpeechAPI:
        response = client.post("/v1/audio/speech", json=payload)
        assert response.status_code == 422  # Unprocessable Entity

-    @patch("vllm_omni.entrypoints.openai.serving_speech.OmniOpenAIServingSpeech.create_audio")
-    def test_speed_parameter_is_used(self, mock_create_audio, test_app):
+    def test_speed_parameter_is_used(self, test_app, mocker: MockerFixture):
+        mock_create_audio = mocker.patch(
+            "vllm_omni.entrypoints.openai.serving_speech.OmniOpenAIServingSpeech.create_audio"
+        )
        client = TestClient(test_app)

-        mock_audio_response = MagicMock()
+        mock_audio_response = mocker.MagicMock()
        mock_audio_response.audio_data = b"dummy_audio"
        mock_audio_response.media_type = "audio/wav"
        mock_create_audio.return_value = mock_audio_response
@@ -287,33 +288,55 @@ class TestTTSMethods:
    """Unit tests for TTS validation and parameter building."""

    @pytest.fixture
-    def speech_server(self):
-        mock_engine_client = MagicMock()
+    def speech_server(self, mocker: MockerFixture):
+        mock_engine_client = mocker.MagicMock()
        mock_engine_client.errored = False
        mock_engine_client.stage_list = None
-        mock_models = MagicMock()
+        mock_engine_client.tts_max_instructions_length = None
+        mock_models = mocker.MagicMock()
        mock_models.is_base_model.return_value = True
        return OmniOpenAIServingSpeech(
            engine_client=mock_engine_client,
            models=mock_models,
-            request_logger=MagicMock(),
+            request_logger=mocker.MagicMock(),
        )

-    def test_is_tts_model(self, speech_server):
-        """Test TTS model detection."""
-        # No stage_list -> False
-        assert speech_server._is_tts_model() is False
+    def test_is_tts_detection_no_stage(self, speech_server):
+        """Test TTS model detection when no TTS stage exists."""
+        # Fixture creates server with stage_list = None -> _is_tts should be False
+        assert speech_server._is_tts is False
+        assert speech_server._tts_stage is None

-        # With qwen3_tts stage -> True
-        mock_stage = MagicMock()
+    def test_is_tts_detection_with_tts_stage(self, mocker: MockerFixture):
+        """Test TTS model detection when TTS stage exists."""
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.tts_max_instructions_length = None
+
+        # Create a TTS stage
+        mock_stage = mocker.MagicMock()
        mock_stage.model_stage = "qwen3_tts"
-        speech_server.engine_client.stage_list = [mock_stage]
-        assert speech_server._is_tts_model() is True
+        mock_stage.tts_args = {}
+        mock_engine_client.stage_list = [mock_stage]

-    def test_build_tts_prompt(self, speech_server):
-        """Test TTS prompt format."""
-        prompt = speech_server._build_tts_prompt("Hello")
-        assert prompt == "<|im_start|>assistant\nHello<|im_end|>\n<|im_start|>assistant\n"
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+
+        assert server._is_tts is True
+        assert server._tts_stage is mock_stage
+
+    def test_estimate_prompt_len_fallback(self, speech_server):
+        """Test prompt length estimation falls back to 2048 when model is unavailable."""
+        tts_params = {"text": ["Hello"], "task_type": ["CustomVoice"]}
+        result = speech_server._estimate_prompt_len(tts_params)
+        # Without a real model, it should fall back to 2048.
+        assert result == 2048

    def test_validate_tts_request_basic(self, speech_server):
        """Test basic validation cases."""
@@ -357,25 +380,130 @@ class TestTTSMethods:
        assert params["language"] == ["English"]
        assert params["task_type"] == ["CustomVoice"]

-    def test_load_supported_speakers(self):
+    def test_load_supported_speakers(self, mocker: MockerFixture):
        """Test _load_supported_speakers."""
-        mock_engine_client = MagicMock()
+        mock_engine_client = mocker.MagicMock()
        mock_engine_client.errored = False
        mock_engine_client.stage_list = None

        # Mock talker_config with mixed-case speaker names
-        mock_talker_config = MagicMock()
+        mock_talker_config = mocker.MagicMock()
        mock_talker_config.spk_id = {"Ryan": 0, "Vivian": 1, "Aiden": 2}
        mock_engine_client.model_config.hf_config.talker_config = mock_talker_config

-        mock_models = MagicMock()
+        mock_models = mocker.MagicMock()
        mock_models.is_base_model.return_value = True

        server = OmniOpenAIServingSpeech(
            engine_client=mock_engine_client,
            models=mock_models,
-            request_logger=MagicMock(),
+            request_logger=mocker.MagicMock(),
        )

        # Verify speakers are normalized to lowercase
        assert server.supported_speakers == {"ryan", "vivian", "aiden"}
+
+    def test_max_instructions_length_default(self, speech_server):
+        """Test default max instructions length (500) when no config provided."""
+        # Fixture creates server with no CLI override and no TTS stage
+        assert speech_server._max_instructions_length == 500
+
+    def test_max_instructions_length_cli_override(self, mocker: MockerFixture):
+        """Test CLI override (stored in engine_client) takes highest priority."""
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.stage_list = None
+        # CLI override is stored in engine_client
+        mock_engine_client.tts_max_instructions_length = 1000
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+        # Value is cached during __init__
+        assert server._max_instructions_length == 1000
+
+    def test_max_instructions_length_stage_config(self, mocker: MockerFixture):
+        """Test stage config value is used when no CLI override."""
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.tts_max_instructions_length = None  # No CLI override
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+
+        # Mock stage with tts_args
+        mock_stage = mocker.MagicMock()
+        mock_stage.model_stage = "qwen3_tts"
+        mock_stage.tts_args = {"max_instructions_length": 750}
+        mock_engine_client.stage_list = [mock_stage]
+
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+        # Value is cached during __init__
+        assert server._max_instructions_length == 750
+
+    def test_max_instructions_length_cli_overrides_stage_config(self, mocker: MockerFixture):
+        """Test CLI override (in engine_client) takes precedence over stage config."""
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        # CLI override stored in engine_client
+        mock_engine_client.tts_max_instructions_length = 2000
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+
+        # Mock stage with tts_args
+        mock_stage = mocker.MagicMock()
+        mock_stage.model_stage = "qwen3_tts"
+        mock_stage.tts_args = {"max_instructions_length": 750}
+        mock_engine_client.stage_list = [mock_stage]
+
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+        # CLI value (2000) should override stage config (750)
+        assert server._max_instructions_length == 2000
+
+    def test_validate_instructions_length_uses_cached_value(self, mocker: MockerFixture):
+        """Test instructions length validation uses cached _max_instructions_length."""
+        mock_engine_client = mocker.MagicMock()
+        mock_engine_client.errored = False
+        mock_engine_client.stage_list = None
+        # CLI override with max length of 10 characters
+        mock_engine_client.tts_max_instructions_length = 10
+        mock_models = mocker.MagicMock()
+        mock_models.is_base_model.return_value = True
+
+        server = OmniOpenAIServingSpeech(
+            engine_client=mock_engine_client,
+            models=mock_models,
+            request_logger=mocker.MagicMock(),
+        )
+
+        # Verify cached value
+        assert server._max_instructions_length == 10
+
+        # Instructions within limit should pass
+        req = OpenAICreateSpeechRequest(
+            input="Hello",
+            task_type="VoiceDesign",
+            instructions="short",
+        )
+        assert server._validate_tts_request(req) is None
+
+        # Instructions exceeding limit should fail
+        req = OpenAICreateSpeechRequest(
+            input="Hello",
+            task_type="VoiceDesign",
+            instructions="this is too long",
+        )
+        error = server._validate_tts_request(req)
+        assert error is not None
+        assert "max 10 characters" in error
@@ -5,18 +5,20 @@ Unit tests for OpenAI-compatible video generation endpoints.
 """

 import io
-from unittest.mock import patch

 import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 from PIL import Image
 from pydantic import ValidationError
+from pytest_mock import MockerFixture

 from vllm_omni.entrypoints.openai.api_server import router
 from vllm_omni.entrypoints.openai.protocol.videos import VideoGenerationRequest, VideoResponseFormat
 from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo

+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+

 class MockVideoResult:
    def __init__(self, videos):
@@ -55,27 +57,27 @@ def _make_test_image_bytes(size=(64, 64)) -> bytes:
    return buf.getvalue()


-def test_t2v_video_generation_form(test_client):
+def test_t2v_video_generation_form(test_client, mocker: MockerFixture):
    fps_values = []

    def _fake_encode(video, fps):
        fps_values.append(fps)
        return "Zg=="

-    with patch(
+    mocker.patch(
        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
        side_effect=_fake_encode,
-    ):
-        response = test_client.post(
-            "/v1/videos",
-            data={
-                "prompt": "A cat runs across the street.",
-                "size": "640x360",
-                "seconds": "2",
-                "fps": "12",
-                "n": "2",
-            },
-        )
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "A cat runs across the street.",
+            "size": "640x360",
+            "seconds": "2",
+            "fps": "12",
+            "n": "2",
+        },
+    )

    assert response.status_code == 200
    data = response.json()
@@ -92,18 +94,18 @@ def test_t2v_video_generation_form(test_client):
    assert fps_values == [12, 12]


-def test_i2v_video_generation_form(test_client):
+def test_i2v_video_generation_form(test_client, mocker: MockerFixture):
    image_bytes = _make_test_image_bytes((48, 32))

-    with patch(
+    mocker.patch(
        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
        return_value="Zg==",
-    ):
-        response = test_client.post(
-            "/v1/videos",
-            data={"prompt": "A bear playing with yarn."},
-            files={"input_reference": ("input.png", image_bytes, "image/png")},
-        )
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={"prompt": "A bear playing with yarn."},
+        files={"input_reference": ("input.png", image_bytes, "image/png")},
+    )

    assert response.status_code == 200
    data = response.json()
@@ -119,24 +121,24 @@ def test_i2v_video_generation_form(test_client):
    assert input_image.size == (48, 32)


-def test_seconds_defaults_fps_and_frames(test_client):
+def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
    fps_values = []

    def _fake_encode(video, fps):
        fps_values.append(fps)
        return "Zg=="

-    with patch(
+    mocker.patch(
        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
        side_effect=_fake_encode,
-    ):
-        response = test_client.post(
-            "/v1/videos",
-            data={
-                "prompt": "A bird flying.",
-                "seconds": "3",
-            },
-        )
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "A bird flying.",
+            "seconds": "3",
+        },
+    )

    assert response.status_code == 200
    engine = test_client.app.state.openai_serving_video._engine_client
@@ -146,18 +148,18 @@ def test_seconds_defaults_fps_and_frames(test_client):
    assert fps_values == [24]


-def test_size_param_sets_width_height(test_client):
-    with patch(
+def test_size_param_sets_width_height(test_client, mocker: MockerFixture):
+    mocker.patch(
        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
        return_value="Zg==",
-    ):
-        response = test_client.post(
-            "/v1/videos",
-            data={
-                "prompt": "size test",
-                "size": "320x240",
-            },
-        )
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "size test",
+            "size": "320x240",
+        },
+    )

    assert response.status_code == 200
    engine = test_client.app.state.openai_serving_video._engine_client
@@ -166,23 +168,23 @@ def test_size_param_sets_width_height(test_client):
    assert captured.height == 240


-def test_sampling_params_pass_through(test_client):
-    with patch(
+def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
+    mocker.patch(
        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
        return_value="Zg==",
-    ):
-        response = test_client.post(
-            "/v1/videos",
-            data={
-                "prompt": "param pass",
-                "num_inference_steps": "30",
-                "guidance_scale": "6.5",
-                "guidance_scale_2": "8.0",
-                "true_cfg_scale": "4.0",
-                "boundary_ratio": "0.7",
-                "flow_shift": "0.25",
-            },
-        )
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "param pass",
+            "num_inference_steps": "30",
+            "guidance_scale": "6.5",
+            "guidance_scale_2": "8.0",
+            "true_cfg_scale": "4.0",
+            "boundary_ratio": "0.7",
+            "flow_shift": "0.25",
+        },
+    )

    assert response.status_code == 200
    engine = test_client.app.state.openai_serving_video._engine_client
@@ -264,19 +266,19 @@ def test_invalid_n_raises_validation_error(test_client):
        )


-def test_negative_prompt_and_seed_pass_through(test_client):
-    with patch(
+def test_negative_prompt_and_seed_pass_through(test_client, mocker: MockerFixture):
+    mocker.patch(
        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
        return_value="Zg==",
-    ):
-        response = test_client.post(
-            "/v1/videos",
-            data={
-                "prompt": "snowy mountain",
-                "negative_prompt": "blurry",
-                "seed": "123",
-            },
-        )
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "snowy mountain",
+            "negative_prompt": "blurry",
+            "seed": "123",
+        },
+    )

    assert response.status_code == 200
    engine = test_client.app.state.openai_serving_video._engine_client
@@ -3,21 +3,23 @@

 import pytest

-from vllm_omni.entrypoints import omni as omni_module
+from vllm_omni.entrypoints import utils as utils_module
 from vllm_omni.entrypoints.async_omni import AsyncOmni

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]

+MODEL = "riverclouds/qwen_image_random"
+

 def test_default_stage_config_includes_cache_backend(monkeypatch):
    """Ensure cache_backend/cache_config are preserved in default diffusion stage."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
+    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
+    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)

    omni = AsyncOmni(
-        model="dummy-model",
+        model=MODEL,
        cache_backend="cache_dit",
        cache_config='{"Fn_compute_blocks": 2}',
        vae_use_slicing=True,
@@ -41,13 +43,13 @@ def test_default_stage_config_includes_cache_backend(monkeypatch):

 def test_default_cache_config_used_when_missing(monkeypatch):
    """Ensure default cache_config is applied when cache_backend is set."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
+    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
+    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)

    omni = AsyncOmni(
-        model="dummy-model",
+        model=MODEL,
        cache_backend="cache_dit",
    )

@@ -59,13 +61,13 @@ def test_default_cache_config_used_when_missing(monkeypatch):

 def test_default_stage_devices_from_sequence_parallel(monkeypatch):
    """Ensure devices list reflects sequence parallel size when no parallel_config is provided."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
+    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
+    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
    monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
    monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)

    omni = AsyncOmni(
-        model="dummy-model",
+        model=MODEL,
        ulysses_degree=2,
        ring_degree=2,
    )
@@ -2,9 +2,9 @@ import uuid
 import warnings
 from queue import Empty, Queue
 from typing import Any
-from unittest.mock import MagicMock

 import pytest
+from pytest_mock import MockerFixture
 from vllm import SamplingParams

 from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
@@ -18,6 +18,8 @@ warnings.filterwarnings(
    category=DeprecationWarning,
 )

+MODEL = "riverclouds/qwen_image_random"
+

 class _FakeEngineArgs(dict):
    """Fake engine args that can be used both as object attributes and as **kwargs."""
@@ -60,8 +62,10 @@ class _FakeQueue:
    def put_nowait(self, item):
        self._queue.put_nowait(item)

-    def get(self):
-        return self._queue.get()
+    def get(self, timeout=None):
+        if timeout is None:
+            return self._queue.get()
+        return self._queue.get(timeout=timeout)

    def get_nowait(self):
        return self._queue.get_nowait()
@@ -70,10 +74,56 @@ class _FakeQueue:
        return self._queue.empty()


+class _FakeZmqQueue:
+    """Fake ZmqQueue that wraps _FakeQueue and matches ZmqQueue interface."""
+
+    def __init__(
+        self,
+        ctx=None,
+        socket_type=None,
+        *,
+        bind: str | None = None,
+        connect: str | None = None,
+        recv_timeout_ms: int | None = None,
+        send_timeout_ms: int | None = None,
+    ):
+        """Initialize fake ZMQ queue with same signature as real ZmqQueue."""
+        self._queue = _FakeQueue(maxsize=0)
+        # Determine endpoint from bind or connect
+        path = bind if bind is not None else connect
+        self.endpoint = path or f"fake://zmq-endpoint-{id(self)}"
+        self._recv_timeout_ms = recv_timeout_ms
+        self._send_timeout_ms = send_timeout_ms
+
+    def put(self, obj: Any) -> None:
+        """Send an object to the queue."""
+        self._queue.put(obj)
+
+    def put_nowait(self, obj: Any) -> None:
+        """Send an object to the queue without blocking."""
+        self._queue.put_nowait(obj)
+
+    def get(self, timeout: float | None = None) -> Any:
+        """Receive an object from the queue with optional timeout in seconds."""
+        return self._queue.get(timeout=timeout)
+
+    def get_nowait(self) -> Any:
+        """Receive an object from the queue without blocking."""
+        return self._queue.get_nowait()
+
+    def empty(self) -> bool:
+        """Check if the queue is empty without blocking."""
+        return self._queue.empty()
+
+    def close(self) -> None:
+        """Close the queue."""
+        pass
+
+
 class _FakeStage:
    """Lightweight Stage stub for multi-process pipeline version with queue support."""

-    def __init__(self, config, stage_init_timeout: int = 300):
+    def __init__(self, mocker: MockerFixture, config, stage_init_timeout: int = 300):
        # Handle both dict and object configs
        if isinstance(config, dict):
            config = _FakeStageConfig(config)
@@ -99,6 +149,8 @@ class _FakeStage:
        self._out_q = None
        self._proc = None  # Mock process reference
        self._stage_init_timeout = max(0, int(stage_init_timeout))
+        # NOTE: mocker fixture object **must be provided**
+        self._mocker = mocker

    def attach_queues(self, in_q, out_q):
        """Attach input and output queues."""
@@ -117,11 +169,11 @@ class _FakeStage:
    ):
        """Mock init_stage_worker: don't start real process, just send stage_ready message."""
        # Create a mock process object
-        self._proc = MagicMock()
-        self._proc.start = MagicMock()
-        self._proc.join = MagicMock()
-        self._proc.is_alive = MagicMock(return_value=False)
-        self._proc.terminate = MagicMock()
+        self._proc = self._mocker.MagicMock()
+        self._proc.start = self._mocker.MagicMock()
+        self._proc.join = self._mocker.MagicMock()
+        self._proc.is_alive = self._mocker.MagicMock(return_value=False)
+        self._proc.terminate = self._mocker.MagicMock()
        # Send stage_ready message to output queue
        if self._out_q is not None:
            try:
@@ -185,21 +237,21 @@ def fake_stage_config():
    }


-def _setup_engine_mocks(monkeypatch):
+def _setup_engine_mocks(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture):
    """Helper function to set up common engine mocks."""
-    fake_engine = MagicMock()
+    fake_engine = mocker.MagicMock()
    # Add necessary attributes to fake_engine
-    fake_engine.tokenizer = MagicMock()
+    fake_engine.tokenizer = mocker.MagicMock()
    fake_engine.log_stats = False
-    fake_engine.vllm_config = MagicMock()
-    fake_engine.vllm_config.model_config = MagicMock()
+    fake_engine.vllm_config = mocker.MagicMock()
+    fake_engine.vllm_config.model_config = mocker.MagicMock()
    fake_engine.vllm_config.model_config.io_processor_plugin = None
-    fake_engine.get_supported_tasks = MagicMock(return_value=[])
-    fake_engine.model_config = MagicMock()
+    fake_engine.get_supported_tasks = mocker.MagicMock(return_value=[])
+    fake_engine.model_config = mocker.MagicMock()
    fake_engine.model_config.io_processor_plugin = None
    # Add registry with resolve_model_cls method
-    fake_registry = MagicMock()
-    fake_registry.resolve_model_cls = MagicMock(return_value=(MagicMock(), "test_arch"))
+    fake_registry = mocker.MagicMock()
+    fake_registry.resolve_model_cls = mocker.MagicMock(return_value=(mocker.MagicMock(), "test_arch"))
    fake_engine.model_config.registry = fake_registry
    fake_engine.vllm_config.model_config.registry = fake_registry

@@ -248,21 +300,21 @@ def _setup_engine_mocks(monkeypatch):
    )


-def _setup_multiprocessing_mocks(monkeypatch):
+def _setup_multiprocessing_mocks(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture):
    """Helper function to set up multiprocessing mocks."""
    import multiprocessing as mp

    # Mock Process
-    fake_process_class = MagicMock()
-    fake_process_instance = MagicMock()
-    fake_process_instance.start = MagicMock()
-    fake_process_instance.join = MagicMock()
-    fake_process_instance.is_alive = MagicMock(return_value=False)
-    fake_process_instance.terminate = MagicMock()
+    fake_process_class = mocker.MagicMock()
+    fake_process_instance = mocker.MagicMock()
+    fake_process_instance.start = mocker.MagicMock()
+    fake_process_instance.join = mocker.MagicMock()
+    fake_process_instance.is_alive = mocker.MagicMock(return_value=False)
+    fake_process_instance.terminate = mocker.MagicMock()
    fake_process_class.return_value = fake_process_instance

    # Mock get_context to return a context with Queue that returns _FakeQueue
-    fake_ctx = MagicMock()
+    fake_ctx = mocker.MagicMock()
    fake_ctx.Queue = lambda maxsize=0: _FakeQueue(maxsize=maxsize)
    fake_ctx.Process = fake_process_class

@@ -272,8 +324,21 @@ def _setup_multiprocessing_mocks(monkeypatch):
    monkeypatch.setattr(mp, "get_context", _mock_get_context, raising=False)
    monkeypatch.setattr(mp, "Process", fake_process_class, raising=False)

+    # Mock ZmqQueue to use _FakeZmqQueue
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.zmq_utils.ZmqQueue",
+        _FakeZmqQueue,
+        raising=False,
+    )
+    # Also mock where ZmqQueue is imported/used
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.ZmqQueue",
+        _FakeZmqQueue,
+        raising=False,
+    )

-def _setup_ipc_mocks(monkeypatch):
+
+def _setup_ipc_mocks(monkeypatch: pytest.MonkeyPatch):
    """Helper function to set up IPC function mocks."""

    # Mock _encode: simple serialization
@@ -293,7 +358,7 @@ def _setup_ipc_mocks(monkeypatch):
    monkeypatch.setattr("vllm_omni.entrypoints.omni._set", _fake_set, raising=False)


-def _setup_log_mocks(monkeypatch):
+def _setup_log_mocks(monkeypatch: pytest.MonkeyPatch):
    """Helper function to set up logging and stats mocks."""
    # Mock OrchestratorMetrics to be a simple class that doesn't require file operations

@@ -324,16 +389,75 @@ def _setup_log_mocks(monkeypatch):
    )


+def _setup_connector_mocks(monkeypatch, mocker, omni_module=None):
+    """Helper function to set up connector mocks for stage-to-stage forwarding.
+
+    If omni_module is provided, mocks directly on the module. Otherwise, uses string path.
+    """
+
+    # Mock initialize_orchestrator_connectors to return fake connectors
+    def _fake_initialize_orchestrator_connectors(config_path, worker_backend=None, shm_threshold_bytes=None):
+        # Create fake connectors for all stage-to-stage edges
+        # Each connector is just a mock object that will be passed to try_send_via_connector
+        fake_connectors = {}
+        # Add connectors for common edges (0->1, 1->2, etc.)
+        for i in range(10):  # Support up to 10 stages
+            fake_connectors[(str(i), str(i + 1))] = mocker.MagicMock()
+        return None, fake_connectors
+
+    if omni_module is not None:
+        # Mock directly on the omni module where it's used (after import)
+        monkeypatch.setattr(omni_module, "initialize_orchestrator_connectors", _fake_initialize_orchestrator_connectors)
+    else:
+        # Mock via string path (before import)
+        monkeypatch.setattr(
+            "vllm_omni.entrypoints.omni.initialize_orchestrator_connectors",
+            _fake_initialize_orchestrator_connectors,
+            raising=False,
+        )
+
+
+def _setup_connector_adapter_mock(monkeypatch, omni_module):
+    """Helper function to mock try_send_via_connector on the omni module.
+
+    This must be called AFTER importing omni module, to mock the function where it's actually used.
+    """
+
+    # Mock try_send_via_connector to always succeed
+    def _fake_try_send_via_connector(
+        connector,
+        stage_id,
+        next_stage_id,
+        req_id,
+        next_inputs,
+        sampling_params,
+        original_prompt,
+        next_stage_queue_submit_fn,
+        metrics,
+    ):
+        # Simulate successful send by calling the submit function
+        task = {
+            "request_id": req_id,
+            "engine_inputs": next_inputs,
+            "sampling_params": sampling_params,
+        }
+        next_stage_queue_submit_fn(task)
+        return True
+
+    # Mock directly on the omni module where it's used
+    monkeypatch.setattr(omni_module, "try_send_via_connector", _fake_try_send_via_connector)
+
+
@pytest.fixture(autouse=True)
-def mock_get_config(monkeypatch):
+def mock_get_config(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture):
    """Auto-mock get_config and related model loading functions to avoid model path validation."""
    # CRITICAL: Mock tokenizer-related imports FIRST, before any module imports
    # This prevents ImportError when async_omni is imported (which happens via omni_stage)
    import sys

-    fake_tokenizer = MagicMock()
-    fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3])
-    fake_tokenizer.decode = MagicMock(return_value="test")
+    fake_tokenizer = mocker.MagicMock()
+    fake_tokenizer.encode = mocker.MagicMock(return_value=[1, 2, 3])
+    fake_tokenizer.decode = mocker.MagicMock(return_value="test")

    # Mock init_tokenizer_from_configs (used in async_omni)
    def _mock_init_tokenizer_from_configs(model_config=None, **kwargs):
@@ -404,7 +528,7 @@ def mock_get_config(monkeypatch):
        setattr(async_omni_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)

    # Now mock get_config and other functions
-    fake_hf_config = MagicMock()
+    fake_hf_config = mocker.MagicMock()
    fake_hf_config.model_type = "qwen2_5_omni"

    def _mock_get_config(model, **kwargs):
@@ -446,13 +570,24 @@ def mock_get_config(monkeypatch):
    )


-def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config):
+def test_initialize_stage_configs_called_when_none(
+    monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config
+):
    """Test that stage configs are auto-loaded when stage_configs_path is None."""

-    def _fake_loader(model: str, base_engine_args=None):
-        return [
-            _FakeStageConfig(fake_stage_config),
-            _FakeStageConfig(fake_stage_config),
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        cfg0 = dict(fake_stage_config)
+        cfg0["stage_id"] = 0
+        cfg1 = dict(fake_stage_config)
+        cfg1["stage_id"] = 1
+        return None, [
+            _FakeStageConfig(cfg0),
+            _FakeStageConfig(cfg1),
        ]

    # Remove modules from cache BEFORE setting mocks
@@ -467,14 +602,15 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
            del sys.modules[module_name]

    # Set up mocks
-    _setup_engine_mocks(monkeypatch)
-    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
    _setup_ipc_mocks(monkeypatch)
    _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch, mocker)

-    # Mock load_stage_configs_from_model
+    # Mock load_and_resolve_stage_configs
    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
        _fake_loader,
        raising=False,
    )
@@ -482,7 +618,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
    # Replace OmniStage
    monkeypatch.setattr(
        "vllm_omni.entrypoints.omni_stage.OmniStage",
-        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs),
        raising=False,
    )

@@ -490,12 +626,14 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
    import vllm_omni.entrypoints.omni as omni_module

    # Patch the imported function and class in the module
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
-    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, mocker, omni_module)

    from vllm_omni.entrypoints.omni import Omni

-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
    # Verify: auto-loaded stage_configs and stage_list have consistent count
    assert isinstance(omni.stage_configs, list)
    assert len(omni.stage_configs) == 2
@@ -511,11 +649,18 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
    assert len(omni._stages_ready) == 2


-def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
+def test_generate_raises_on_length_mismatch(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config):
    """Test that generate raises ValueError when sampling_params_list length doesn't match."""

-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        cfg0 = dict(fake_stage_config)
+        cfg0["stage_id"] = 0
+        return None, [_FakeStageConfig(cfg0)]

    import sys

@@ -527,47 +672,50 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
        if module_name in sys.modules:
            del sys.modules[module_name]

-    _setup_engine_mocks(monkeypatch)
-    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
    _setup_ipc_mocks(monkeypatch)
    _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch, mocker)

    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-        _fake_loader,
-        raising=False,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
        _fake_loader,
        raising=False,
    )
    monkeypatch.setattr(
        "vllm_omni.entrypoints.omni_stage.OmniStage",
-        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs),
        raising=False,
    )

    import vllm_omni.entrypoints.omni as omni_module

-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
-    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))

    from vllm_omni.entrypoints.omni import Omni

-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
    with pytest.raises(ValueError):
        omni.generate(prompts=["hi"], sampling_params_list=[])


-def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
+def test_generate_pipeline_and_final_outputs(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config):
    """Test multi-stage generation pipeline with queue polling."""
    stage_cfg0 = dict(fake_stage_config)
+    stage_cfg0["stage_id"] = 0
    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["stage_id"] = 1
    stage_cfg1["processed_input"] = ["processed-for-stage-1"]

-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]

    import sys

@@ -579,31 +727,30 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
        if module_name in sys.modules:
            del sys.modules[module_name]

-    _setup_engine_mocks(monkeypatch)
-    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
    _setup_ipc_mocks(monkeypatch)
    _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch, mocker)

    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-        _fake_loader,
-        raising=False,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
        _fake_loader,
        raising=False,
    )
    monkeypatch.setattr(
        "vllm_omni.entrypoints.omni_stage.OmniStage",
-        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs),
        raising=False,
    )

    import vllm_omni.entrypoints.omni as omni_module

-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
-    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, mocker, omni_module)
+    _setup_connector_adapter_mock(monkeypatch, omni_module)

    # Mock uuid.uuid4() to return a predictable value for request ID generation
    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -612,7 +759,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):

    from vllm_omni.entrypoints.omni import Omni

-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)

    # Generate the expected request ID format: "0_<uuid>"
    expected_request_id = f"0_{test_uuid}"
@@ -659,15 +806,24 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
    assert omni.stage_list[1].process_engine_inputs([], []) is not None


-def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
+def test_generate_no_final_output_returns_empty(
+    monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config
+):
    """Test that generate returns empty list when all stages have final_output=False."""
    stage_cfg0 = dict(fake_stage_config)
-    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg0["stage_id"] = 0
    stage_cfg0["final_output"] = False
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["stage_id"] = 1
    stage_cfg1["final_output"] = False

-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]

    import sys

@@ -679,26 +835,30 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
        if module_name in sys.modules:
            del sys.modules[module_name]

-    _setup_engine_mocks(monkeypatch)
-    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
    _setup_ipc_mocks(monkeypatch)
    _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch, mocker)

    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
        _fake_loader,
        raising=False,
    )
    monkeypatch.setattr(
        "vllm_omni.entrypoints.omni_stage.OmniStage",
-        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs),
        raising=False,
    )

    import vllm_omni.entrypoints.omni as omni_module

-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
-    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, mocker, omni_module)
+    _setup_connector_adapter_mock(monkeypatch, omni_module)

    # Mock uuid.uuid4() to return a predictable value for request ID generation
    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -707,7 +867,7 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):

    from vllm_omni.entrypoints.omni import Omni

-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)

    # Generate the expected request ID format: "0_<uuid>"
    expected_request_id = f"0_{test_uuid}"
@@ -738,15 +898,24 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
    assert outputs == []


-def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_config):
+def test_generate_sampling_params_none_use_default(
+    monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config
+):
    """Test that generate uses default sampling params when sampling_params_list is None."""
    stage_cfg0 = dict(fake_stage_config)
-    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg0["stage_id"] = 0
    stage_cfg0["final_output"] = False
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["stage_id"] = 1
    stage_cfg1["final_output"] = False

-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]

    import sys

@@ -758,26 +927,30 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
        if module_name in sys.modules:
            del sys.modules[module_name]

-    _setup_engine_mocks(monkeypatch)
-    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
    _setup_ipc_mocks(monkeypatch)
    _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch, mocker)

    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
        _fake_loader,
        raising=False,
    )
    monkeypatch.setattr(
        "vllm_omni.entrypoints.omni_stage.OmniStage",
-        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs),
        raising=False,
    )

    import vllm_omni.entrypoints.omni as omni_module

-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
-    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, mocker, omni_module)
+    _setup_connector_adapter_mock(monkeypatch, omni_module)

    # Mock uuid.uuid4() to return a predictable value for request ID generation
    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -786,7 +959,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi

    from vllm_omni.entrypoints.omni import Omni

-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)

    # Generate the expected request ID format: "0_<uuid>"
    expected_request_id = f"0_{test_uuid}"
@@ -810,11 +983,18 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
    omni.generate(prompts=["p"], sampling_params_list=None)


-def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
+def test_wait_for_stages_ready_timeout(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config):
    """Test that _wait_for_stages_ready handles timeout correctly."""

-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        cfg0 = dict(fake_stage_config)
+        cfg0["stage_id"] = 0
+        return None, [_FakeStageConfig(cfg0)]

    import sys

@@ -826,13 +1006,14 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
        if module_name in sys.modules:
            del sys.modules[module_name]

-    _setup_engine_mocks(monkeypatch)
-    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
    _setup_ipc_mocks(monkeypatch)
    _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch, mocker)

    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
        _fake_loader,
        raising=False,
    )
@@ -841,36 +1022,43 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
    class _FakeStageNoReady(_FakeStage):
        def init_stage_worker(self, *args, **kwargs):
            # Don't send stage_ready message
-            self._proc = MagicMock()
-            self._proc.start = MagicMock()
-            self._proc.join = MagicMock()
-            self._proc.is_alive = MagicMock(return_value=False)
-            self._proc.terminate = MagicMock()
+            self._proc = self._mocker.MagicMock()
+            self._proc.start = self._mocker.MagicMock()
+            self._proc.join = self._mocker.MagicMock()
+            self._proc.is_alive = self._mocker.MagicMock(return_value=False)
+            self._proc.terminate = self._mocker.MagicMock()

    monkeypatch.setattr(
        "vllm_omni.entrypoints.omni_stage.OmniStage",
-        lambda cfg, **kwargs: _FakeStageNoReady(cfg, **kwargs),
+        lambda cfg, **kwargs: _FakeStageNoReady(mocker, cfg, **kwargs),
        raising=False,
    )

    import vllm_omni.entrypoints.omni as omni_module

-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
-    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(mocker, cfg, **kwargs))

    from vllm_omni.entrypoints.omni import Omni

    # Use very short timeout
-    omni = Omni(model="any", init_timeout=0.01)
+    omni = Omni(model=MODEL, init_timeout=0.01)
    # Verify that no stages are ready
    assert len(omni._stages_ready) == 0


-def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
+def test_generate_handles_error_messages(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config):
    """Test that generate handles error messages from stages correctly."""

-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        cfg0 = dict(fake_stage_config)
+        cfg0["stage_id"] = 0
+        return None, [_FakeStageConfig(cfg0)]

    import sys

@@ -882,26 +1070,27 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
        if module_name in sys.modules:
            del sys.modules[module_name]

-    _setup_engine_mocks(monkeypatch)
-    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
    _setup_ipc_mocks(monkeypatch)
    _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch, mocker)

    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
        _fake_loader,
        raising=False,
    )
    monkeypatch.setattr(
        "vllm_omni.entrypoints.omni_stage.OmniStage",
-        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs),
        raising=False,
    )

    import vllm_omni.entrypoints.omni as omni_module

-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
-    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))

    # Mock uuid.uuid4() to return a predictable value for request ID generation
    test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -910,7 +1099,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):

    from vllm_omni.entrypoints.omni import Omni

-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)

    # Generate the expected request ID format: "0_<uuid>"
    expected_request_id = f"0_{test_uuid}"
@@ -941,11 +1130,18 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
    assert len(outputs) == 1


-def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
+def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config):
    """Test that close() sends shutdown signal to all input queues."""

-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        cfg0 = dict(fake_stage_config)
+        cfg0["stage_id"] = 0
+        return None, [_FakeStageConfig(cfg0)]

    import sys

@@ -957,30 +1153,31 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
        if module_name in sys.modules:
            del sys.modules[module_name]

-    _setup_engine_mocks(monkeypatch)
-    _setup_multiprocessing_mocks(monkeypatch)
+    _setup_engine_mocks(monkeypatch, mocker)
+    _setup_multiprocessing_mocks(monkeypatch, mocker)
    _setup_ipc_mocks(monkeypatch)
    _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch, mocker)

    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
        _fake_loader,
        raising=False,
    )
    monkeypatch.setattr(
        "vllm_omni.entrypoints.omni_stage.OmniStage",
-        lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
+        lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs),
        raising=False,
    )

    import vllm_omni.entrypoints.omni as omni_module

-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
-    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))

    from vllm_omni.entrypoints.omni import Omni

-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)

    # Call close
    omni.close()
@@ -1,8 +1,8 @@
 import os
 import sys
-from unittest.mock import MagicMock

 import pytest
+from pytest_mock import MockerFixture

 from vllm_omni.entrypoints.stage_utils import set_stage_devices

@@ -45,16 +45,18 @@ def _make_dummy_torch(call_log):
    return _Torch


-def _make_mock_platform(device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"):
-    """Create a mock platform for testing."""
-    mock_platform = MagicMock()
+def _make_mock_platform(mocker, device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"):
+    """Create a mock platform for testing.
+    mocker object has to be passed in to utilize this helper function.
+    """
+    mock_platform = mocker.MagicMock()
    mock_platform.device_type = device_type
    mock_platform.device_control_env_var = env_var
    return mock_platform


@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
-def test_set_stage_devices_respects_logical_ids(monkeypatch):
+def test_set_stage_devices_respects_logical_ids(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
    # Preserve an existing logical mapping and ensure devices "0,1" map through it.
    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,7")
    call_log: list[int] = []
@@ -62,7 +64,7 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch):
    monkeypatch.setitem(sys.modules, "torch", dummy_torch)

    # Mock the platform at the source module where it's defined
-    mock_platform = _make_mock_platform(device_type="cuda", env_var="CUDA_VISIBLE_DEVICES")
+    mock_platform = _make_mock_platform(mocker, device_type="cuda", env_var="CUDA_VISIBLE_DEVICES")
    monkeypatch.setattr(
        "vllm_omni.platforms.current_omni_platform",
        mock_platform,
@@ -74,7 +76,7 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch):


@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
-def test_set_stage_devices_npu_platform(monkeypatch):
+def test_set_stage_devices_npu_platform(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
    """Test that set_stage_devices works correctly for NPU platform."""
    monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "4,5")
    call_log: list[int] = []
@@ -99,7 +101,7 @@ def test_set_stage_devices_npu_platform(monkeypatch):
    monkeypatch.setitem(sys.modules, "torch", _NpuTorch)

    # Mock NPU platform at the source module where it's defined
-    mock_platform = _make_mock_platform(device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES")
+    mock_platform = _make_mock_platform(mocker, device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES")
    monkeypatch.setattr(
        "vllm_omni.platforms.current_omni_platform",
        mock_platform,
@@ -0,0 +1,271 @@
+"""Unit tests for vllm_omni.entrypoints.utils module."""
+
+from collections import Counter
+from dataclasses import dataclass
+
+import pytest
+from pytest_mock import MockerFixture
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.entrypoints.utils import (
+    _convert_dataclasses_to_dict,
+    _filter_dict_like_object,
+    filter_dataclass_kwargs,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestFilterDictLikeObject:
+    """Test suite for _filter_dict_like_object function."""
+
+    def test_simple_dict(self):
+        """Test filtering a simple dictionary with no callables."""
+        input_dict = {"key1": "value1", "key2": 42, "key3": [1, 2, 3]}
+        result = _filter_dict_like_object(input_dict)
+
+        assert result == input_dict
+        assert isinstance(result, dict)
+
+    def test_dict_with_nested_values(self):
+        """Test filtering dict with nested dict and list values."""
+        input_dict = {
+            "level1": {
+                "level2": {"key": "value"},
+                "list": [1, 2, 3],
+            },
+            "simple": "string",
+        }
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Nested dicts and lists should be recursively processed
+        assert result["simple"] == "string"
+        assert isinstance(result["level1"], dict)
+
+    def test_dict_with_dataclass_values(self):
+        """Test filtering dict containing dataclass values."""
+
+        @dataclass
+        class TestDataclass:
+            field1: str
+            field2: int
+
+        obj = TestDataclass(field1="test", field2=42)
+        input_dict = {"data": obj, "normal": "value"}
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Dataclass should be converted to dict by recursive _convert_dataclasses_to_dict
+        assert "data" in result
+        assert "normal" in result
+        assert result["normal"] == "value"
+
+    def test_dict_with_counter_values(self):
+        """Test filtering dict containing Counter objects."""
+        counter_obj = Counter({"a": 1, "b": 2})
+        input_dict = {"counter": counter_obj, "normal": "value"}
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Counter should be converted to regular dict
+        assert "counter" in result
+        assert "normal" in result
+        assert result["normal"] == "value"
+
+    def test_empty_dict(self):
+        """Test filtering an empty dictionary."""
+        result = _filter_dict_like_object({})
+        assert result == {}
+        assert isinstance(result, dict)
+
+    def test_dict_with_set_values(self):
+        """Test filtering dict with set values."""
+        input_dict = {"set_key": {1, 2, 3}, "normal": "value"}
+
+        result = _filter_dict_like_object(input_dict)
+
+        assert "set_key" in result
+        assert "normal" in result
+        # Set should be converted to list by _convert_dataclasses_to_dict
+        assert result["normal"] == "value"
+
+    def test_dict_with_none_values(self):
+        """Test filtering dict with None values."""
+        input_dict = {"key1": None, "key2": "value", "key3": 0}
+
+        result = _filter_dict_like_object(input_dict)
+
+        assert result == input_dict
+
+    def test_dict_with_mixed_types(self):
+        """Test filtering dict with mixed value types."""
+        input_dict = {
+            "string": "hello",
+            "int": 42,
+            "float": 3.14,
+            "bool": True,
+            "none": None,
+            "list": [1, 2, 3],
+            "tuple": (1, 2, 3),
+            "set": {1, 2, 3},
+            "dict": {"nested": "value"},
+        }
+
+        result = _filter_dict_like_object(input_dict)
+
+        assert "string" in result
+        assert "int" in result
+        assert "float" in result
+        assert "bool" in result
+        assert "none" in result
+        assert "list" in result
+        assert "tuple" in result
+        assert "set" in result
+        assert "dict" in result
+
+    def test_dict_preserves_key_types(self):
+        """Test that dict key types are preserved."""
+        input_dict = {
+            "string_key": "value1",
+            42: "value2",
+            (1, 2): "value3",  # tuple as key
+        }
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Keys should remain the same
+        assert "string_key" in result
+        assert 42 in result
+        assert (1, 2) in result
+
+    def test_dict_with_recursive_structure(self, mocker: MockerFixture):
+        """Test filtering dict with recursive/complex nested structure."""
+        input_dict = {
+            "level1": {
+                "level2": {
+                    "level3": {"key": "value"},
+                    "callable": lambda x: x,
+                }
+            },
+            "normal": "value",
+        }
+
+        mocker.patch("vllm_omni.entrypoints.utils.logger")
+        result = _filter_dict_like_object(input_dict)
+
+        # Normal key should exist
+        assert "normal" in result
+        # Level1 should exist
+        assert "level1" in result
+
+    def test_integration_with_convert_dataclasses(self, mocker: MockerFixture):
+        """Test that _filter_dict_like_object integrates properly with _convert_dataclasses_to_dict."""
+
+        @dataclass
+        class Config:
+            name: str
+            count: int
+
+        input_dict = {
+            "config": Config(name="test", count=5),
+            "func": lambda x: x,
+            "normal": "value",
+        }
+
+        mocker.patch("vllm_omni.entrypoints.utils.logger")
+        result = _filter_dict_like_object(input_dict)
+
+        # Callable should be filtered
+        assert "func" not in result
+        # Config should be converted to dict
+        assert "config" in result
+        assert "normal" in result
+
+
+class TestConvertDataclassesToDict:
+    """Test suite for _convert_dataclasses_to_dict function."""
+
+    def test_uses_filter_dict_like_object(self, mocker: MockerFixture):
+        """Test that _convert_dataclasses_to_dict uses _filter_dict_like_object for dicts."""
+        input_dict = {
+            "normal": "value",
+            "callable": lambda x: x,
+        }
+
+        mocker.patch("vllm_omni.entrypoints.utils.logger")
+        result = _convert_dataclasses_to_dict(input_dict)
+
+        # Callable should be filtered out by _filter_dict_like_object
+        assert "normal" in result
+        assert "callable" not in result
+
+
+class TestFilterDataclassKwargs:
+    """Test basic functionality of filter_dataclass_kwargs."""
+
+    def test_simple_filtering(self):
+        """Test basic dataclass kwargs filtering."""
+
+        @dataclass
+        class SimpleConfig:
+            name: str
+            count: int
+
+        kwargs = {"name": "test", "count": 42, "invalid": "should_be_removed"}
+        result = filter_dataclass_kwargs(SimpleConfig, kwargs)
+
+        assert "name" in result
+        assert "count" in result
+        assert "invalid" not in result
+
+    def test_invalid_dataclass_raises_error(self):
+        """Test that non-dataclass raises ValueError."""
+        with pytest.raises(ValueError, match="is not a dataclass"):
+            filter_dataclass_kwargs(dict, {})
+
+    def test_invalid_kwargs_type_raises_error(self):
+        """Test that non-dict kwargs raises ValueError."""
+
+        @dataclass
+        class SimpleConfig:
+            name: str
+
+        with pytest.raises(ValueError, match="kwargs must be a dictionary"):
+            filter_dataclass_kwargs(SimpleConfig, "invalid")
+
+    def test_filters_omni_engine_args_unknown_fields(self):
+        """Test that OmniEngineArgs kwargs are filtered to valid fields only."""
+        kwargs = {
+            "model": "dummy",
+            "stage_id": 1,
+            "engine_output_type": "image",
+            "unknown_field": "drop_me",
+        }
+
+        result = filter_dataclass_kwargs(OmniEngineArgs, kwargs)
+
+        assert "model" in result
+        assert "stage_id" in result
+        assert "engine_output_type" in result
+        assert "unknown_field" not in result
+
+    def test_filters_omni_diffusion_config_union_dataclass(self):
+        """Test that OmniDiffusionConfig filters nested dataclass in Union fields."""
+        kwargs = {
+            "model": "dummy",
+            "cache_config": {
+                "rel_l1_thresh": 0.3,
+                "extra_param": "should_drop",
+            },
+            "unknown_top": "drop_me",
+        }
+
+        result = filter_dataclass_kwargs(OmniDiffusionConfig, kwargs)
+
+        assert "model" in result
+        assert "cache_config" in result
+        assert "unknown_top" not in result
+        assert result["cache_config"]["rel_l1_thresh"] == 0.3
+        assert "extra_param" not in result["cache_config"]
@@ -12,7 +12,8 @@ from pathlib import Path

 import pytest

-from tests.conftest import OmniServer, convert_audio_file_to_text, cosine_similarity_text
+from tests.conftest import convert_audio_file_to_text, cosine_similarity_text
+from tests.utils import hardware_test

 models = ["Qwen/Qwen2.5-Omni-7B"]

@@ -21,26 +22,7 @@ stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_config

 example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen2_5_omni")
 # Create parameter combinations for model and stage config
-test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
-
-
-@pytest.fixture(scope="module")
-def omni_server(request):
-    """Start vLLM-Omni server as a subprocess with actual model weights.
-    Uses module scope so the server starts only once per test module.
-    Multi-stage initialization can take 10-20+ minutes.
-    """
-    model, stage_config_path = request.param
-
-    print(f"Starting OmniServer with model: {model}")
-    print("This may take 10-20+ minutes for initialization...")
-
-    with OmniServer(
-        model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"], port=8091
-    ) as server:
-        print("OmniServer started successfully")
-        yield server
-        print("OmniServer stopped")
+test_params = [(8091, model, stage_config) for model in models for stage_config in stage_configs]


 def run_cmd(command):
@@ -67,6 +49,9 @@ def extract_content_after_keyword(keywords, text):
    return matches[0]


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_send_multimodal_request_001(omni_server) -> None:
    command = [
@@ -99,6 +84,9 @@ def test_send_multimodal_request_001(omni_server) -> None:
    # TODO: Verify the E2E latency after confirmation baseline.


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_send_multimodal_request_002(omni_server) -> None:
    command = [
@@ -131,6 +119,9 @@ def test_send_multimodal_request_002(omni_server) -> None:
    # TODO: Verify the E2E latency after confirmation baseline.


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_send_multimodal_request_003(omni_server) -> None:
    command = ["bash", os.path.join(example_dir, "run_curl_multimodal_generation.sh"), "mixed_modalities"]
@@ -151,6 +142,9 @@ def test_send_multimodal_request_003(omni_server) -> None:
    # TODO: Verify the E2E latency after confirmation baseline.


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_modality_control_001(omni_server) -> None:
    command = [
@@ -178,6 +172,9 @@ def test_modality_control_001(omni_server) -> None:
    # TODO: Verify the E2E latency after confirmation baseline.


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_modality_control_002(omni_server) -> None:
    command = [
@@ -203,6 +200,9 @@ def test_modality_control_002(omni_server) -> None:
    # TODO: Verify the E2E latency after confirmation baseline.


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_modality_control_003(omni_server) -> None:
    command = [
@@ -237,6 +237,9 @@ def test_modality_control_003(omni_server) -> None:


@pytest.mark.skip(reason="There is a known issue with stream error.")
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_stream_001(omni_server) -> None:
    command = [
@@ -12,7 +12,8 @@ from pathlib import Path

 import pytest

-from tests.conftest import OmniServer, convert_audio_file_to_text, cosine_similarity_text
+from tests.conftest import convert_audio_file_to_text, cosine_similarity_text
+from tests.utils import hardware_test

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

@@ -21,26 +22,7 @@ stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_config

 example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen3_omni")
 # Create parameter combinations for model and stage config
-test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
-
-
-@pytest.fixture(scope="module")
-def omni_server(request):
-    """Start vLLM-Omni server as a subprocess with actual model weights.
-    Uses module scope so the server starts only once per test module..
-    Multi-stage initialization can take 10-20+ minutes.
-    """
-    model, stage_config_path = request.param
-
-    print(f"Starting OmniServer with model: {model}")
-    print("This may take 10-20+ minutes for initialization...")
-
-    with OmniServer(
-        model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"], port=8091
-    ) as server:
-        print("OmniServer started successfully")
-        yield server
-        print("OmniServer stopped")
+test_params = [(8091, model, stage_config) for model in models for stage_config in stage_configs]


 def run_cmd(command):
@@ -67,6 +49,9 @@ def extract_content_after_keyword(keywords, text):
    return matches[0]


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_send_multimodal_request_001(omni_server) -> None:
    command = [
@@ -92,6 +77,9 @@ def test_send_multimodal_request_001(omni_server) -> None:
    # TODO: Verify the E2E latency after confirmation baseline.


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_send_multimodal_request_002(omni_server) -> None:
    command = [
@@ -120,6 +108,9 @@ def test_send_multimodal_request_002(omni_server) -> None:
    # TODO: Verify the E2E latency after confirmation baseline.


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_send_multimodal_request_003(omni_server) -> None:
    command = ["bash", os.path.join(example_dir, "run_curl_multimodal_generation.sh"), "use_image"]
@@ -134,6 +125,9 @@ def test_send_multimodal_request_003(omni_server) -> None:
    # TODO: Verify the E2E latency after confirmation baseline.


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_modality_control_001(omni_server) -> None:
    command = [
@@ -155,6 +149,9 @@ def test_modality_control_001(omni_server) -> None:
    # TODO: Verify the E2E latency after confirmation baseline.


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_modality_control_002(omni_server) -> None:
    command = [
@@ -175,6 +172,9 @@ def test_modality_control_002(omni_server) -> None:
    # TODO: Verify the E2E latency after confirmation baseline.


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_modality_control_003(omni_server) -> None:
    command = [
@@ -202,6 +202,9 @@ def test_modality_control_003(omni_server) -> None:
    # TODO: Verify the E2E latency after confirmation baseline.


+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_stream_001(omni_server) -> None:
    command = [
@@ -46,7 +46,8 @@ def modify_stage(default_path, updates, deletes):


 def create_unique_server_params(configs: list[dict[str, Any]]) -> list[tuple[str, str, str]]:
-    unique_params = set()
+    unique_params = []
+    seen = set()
    for config in configs:
        test_name = config["test_name"]
        model = config["server_params"]["model"]
@@ -55,9 +56,13 @@ def create_unique_server_params(configs: list[dict[str, Any]]) -> list[tuple[str
        delete = config["server_params"].get("delete", None)
        update = config["server_params"].get("update", None)
        stage_config_path = modify_stage(stage_config_path, update, delete)
-        unique_params.add((test_name, model, stage_config_path))

-    return list(unique_params)
+        server_param = (test_name, model, stage_config_path)
+        if server_param not in seen:
+            seen.add(server_param)
+            unique_params.append(server_param)
+
+    return unique_params


 def create_test_parameter_mapping(configs: list[dict[str, Any]]) -> dict[str, dict]:
@@ -95,6 +100,7 @@ def omni_server(request):
        print(f"Starting OmniServer with test: {test_name}, model: {model}")

        with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
+            server.test_name = test_name
            print("OmniServer started successfully")
            yield server
            print("OmniServer stopping...")
@@ -150,10 +156,15 @@ def get_benchmark_params_for_server(test_name: str) -> list:

 def create_benchmark_indices():
    indices = []
-    for test_name, config_data in server_to_benchmark_mapping.items():
-        params_list = config_data["benchmark_params"]
-        for idx in range(len(params_list)):
-            indices.append((test_name, idx))
+    seen = set()
+    for config in BENCHMARK_CONFIGS:
+        test_name = config["test_name"]
+        if test_name not in seen:
+            seen.add(test_name)
+            params_list = get_benchmark_params_for_server(test_name)
+            for idx in range(len(params_list)):
+                indices.append((test_name, idx))
+
    return indices


@@ -164,6 +175,10 @@ benchmark_indices = create_benchmark_indices()
 def benchmark_params(request, omni_server):
    """Benchmark parameters fixture with proper parametrization"""
    test_name, param_index = request.param
+
+    if test_name != omni_server.test_name:
+        pytest.skip(f"Skipping parameter for {test_name} - current server is {omni_server.test_name}")
+
    all_params = get_benchmark_params_for_server(test_name)

    if not all_params:
@@ -172,6 +187,10 @@ def benchmark_params(request, omni_server):
    if param_index >= len(all_params):
        raise ValueError(f"No benchmark parameters found for index {param_index} in test: {test_name}")

+    current = param_index + 1
+    total = len(all_params)
+    print(f"\n  Running benchmark {current}/{total} for {test_name}")
+
    return {"test_name": test_name, "params": all_params[param_index]}
Author	SHA1	Message	Date
Kevin H. Luu	5bdc2fa345	Update test-merge.yml Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2026-02-26 03:54:05 -08:00
TJian	881c675087	[ROCm] [CI] [Docker] Point to use the latest vLLM v0.16.0 stable version (#1500 ) Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>	2026-02-26 19:16:14 +08:00
Kevin H. Luu	d05212cd5f	Use pull through cache image for H100 pool (#1518 ) Signed-off-by: Kevin H. Luu <khluu000@gmail.com>	2026-02-26 19:09:42 +08:00
Juan Pablo Zuluaga Gao Han Hongsheng Liu	66457c39fa	[Bugfix][Qwen3TTS] (#1289 ) Signed-off-by: pablo <juanz9312@gmail.com> Co-authored-by: Gao Han <gaohan19@huawei.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-02-26 18:58:35 +08:00
Yupu	c9786c5c0c	[Model] Support OmniGen2 (#513 ) Signed-off-by: Yupu <feng.yu.pu0330@gmail.com>	2026-02-26 10:34:28 +00:00
Junhong Liu	3eb3aa54ad	[Bugfix] fix record audio generated frame in offline infer (#1312 ) Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com> Signed-off-by: Junhong Liu <ljh_lbj@163.com>	2026-02-26 17:23:55 +08:00
Alex Brooks	bc26de834e	[Bugfix] Fix LoRA Scaling on Active Adapters (#1421 ) Signed-off-by: Alex Brooks <albrooks@redhat.com>	2026-02-26 17:22:35 +08:00
Yuanheng Zhao	ed08a216c1	[Misc] Migrate L1 tests to use pytest-mock (#1315 ) Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com> Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>	2026-02-26 17:08:39 +08:00
Canlin Guo	9913c00144	[Feature] Support Wan2.2 output with irregular shapes (#1279 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com>	2026-02-26 17:03:43 +08:00
Andy Zhou	a2d5b61703	[Doc] format lora related docs for the user's end (#1009 ) Signed-off-by: AndyZhou952 <jzhoubc@connect.ust.hk> Signed-off-by: Andy Zhou <46011930+AndyZhou952@users.noreply.github.com>	2026-02-26 16:48:29 +08:00
junuxyzandHongsheng Liu	caed95fb45	[Bugfix] Fix tuple/list KV cache extraction crash (#1405 ) Signed-off-by: junuxyz <216036880+junuxyz@users.noreply.github.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-02-26 15:20:40 +08:00
Yuanheng Zhao	904b97d5bd	[Fix][Chore] Qwen3-TTS Modeling Minor Code Sanity Improvements (#1482 ) Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>	2026-02-26 14:39:08 +08:00
Zhou Taichang	986b356077	[CI] Update Dockerfile for vllm-omni CI image and remove obsolete dep… (#1491 ) Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-02-26 14:36:03 +08:00
Canlin Guo	22c8cde328	[NPU] Upgrade to v0.16.0 (#1375 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com>	2026-02-26 14:26:50 +08:00
wangyu	1ae2e1e4c4	[Test] Reduce Perf test case and fix modify stage config (#1449 ) Signed-off-by: yenuo26 <410167048@qq.com>	2026-02-26 14:20:57 +08:00
亦瑾	5e7974c15b	[Bugfix] Fix SDPA attention mask dtype and shape (Fix #857 ) (#1349 ) Signed-off-by: jader <yjader@foxmail.com>	2026-02-26 11:54:28 +08:00
Junhong LiuandGao Han	919ad28da5	[BugFix] process request.num_cached_tokens if it equals to the initial value (#1468 ) Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com> Co-authored-by: Gao Han <gaohan19@huawei.com>	2026-02-26 10:59:18 +08:00
Sangchun Ha	315daeba8c	[Bugfix] Fix case-sensitive task_type matching in Qwen3TTSModelForGeneration (#1455 ) Signed-off-by: Sangchun Ha <seomk9896@gmail.com>	2026-02-26 10:27:13 +08:00
Zhou Taichang	8e950daf6a	[Debug] Multi-Request for Qwen 3 Omni use_audio_in_video (#1433 ) Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-02-26 09:26:25 +08:00
Yanick Schraner	97f4c5b1e7	[Bugfix] Add TTS configuration options (#1177 ) Signed-off-by: Yanick Schraner <yanick.schraner@bs.ch>	2026-02-26 00:14:44 +08:00
Dong WandHongsheng Liu	1eebb374f1	[Model]Add new nextstep_1(Diffusion) model(only T2I) (#612 ) Signed-off-by: Dong Wang <dongw2019@gmail.com> Signed-off-by: sniper35 <dongw2019@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-02-26 00:09:11 +08:00
SHIJIN ZHANG	466a7f7a62	[Bugfix]: Fix merging updated additional information to ensure dict type (#1296 ) Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>	2026-02-26 00:06:49 +08:00
Junhong Liu	43939df2a5	[Bugfix] Robust Audio Data Handling in _create_audio_choice (#1222 ) Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com>	2026-02-25 22:59:49 +08:00
Didan Deng	1e86404456	[Doc] Refine Diffusion Tutorial Documents (#1305 ) Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>	2026-02-25 15:43:08 +08:00
WeiQing Chen	bf9754a8b1	[Misc] update wechat image (#1464 ) Signed-off-by: David Chen <530634352@qq.com>	2026-02-25 12:06:27 +08:00
Zhou Taichang	296d995ec2	[CI] Remove overwhelming debug log (#1463 ) Signed-off-by: tzhouam <tzhouam@connect.ust.hk>	2026-02-25 12:06:01 +08:00
Juan Pablo Zuluagaandpablo	4de077ea62	[Qwen3TTS][Feat] Code2Wav batched decoding (#1426 ) Signed-off-by: pablo <pablo@agigo.ai> Co-authored-by: pablo <pablo@agigo.ai>	2026-02-25 11:16:46 +13:00
汪志鹏	82e1bf2804	[Doc]: Add bagel single/multi node usage with mooncake document (#1450 )	2026-02-24 15:54:57 +00:00
Junhong LiuandHongsheng Liu	c0908dd3fa	[Bugfix] Forward all vllm-omni serve command parameters to model (#985 ) Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com> Signed-off-by: Junhong Liu <ljh_lbj@163.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-02-24 22:20:01 +08:00
Junhong LiuandHongsheng Liu	a8cec43ce6	[Doc] Optimize vLLM-Omni metrics documentation (#1311 ) Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com> Signed-off-by: Junhong Liu <ljh_lbj@163.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-02-24 20:27:29 +08:00
wuhang 汪志鹏 Copilot	36b8f80266	[Feature] Support Stage Based Deployment CLI (#939 ) Signed-off-by: wuhang <wuhang6@huawei.com> Signed-off-by: princepride <wangzhipeng628@gmail.com> Signed-off-by: wuhang <whlbx@hotmail.com> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>	2026-02-24 16:55:23 +08:00
Nicolò LucchesiandHongsheng Liu	06d347b47f	[Misc] Extend Diffusion Benchmark script to other backends (#875 ) Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-02-24 16:44:33 +08:00
Canlin GuoandHongsheng Liu	af11b02081	[Minor][Refactor] Pass seq_token_counts explicitly (#1425 ) Signed-off-by: gcanlin <canlinguosdu@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-02-24 14:10:00 +08:00
NATURE	1589931fa6	[Feature] vLLM-Omni RDMA connector (#1019 ) Signed-off-by: natureofnature <wzliu@connect.hku.hk>	2026-02-24 11:13:52 +08:00
Alicia	7b02b856a6	[CI] Set L2 & L3 tests running conditions. (#1344 ) Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>	2026-02-24 10:50:10 +08:00
Ding ZuhaoandWang Zhipeng: princepride	d6a3551a62	[Bug] [Bagel] Fix kv transfer bug (#1437 ) Signed-off-by: Ding Zuhao <e1583181@u.nus.edu> Co-authored-by: Wang Zhipeng: princepride <wangzhipeng628@gmail.com>	2026-02-23 15:58:22 +00:00
zhou zhuoxin	40f72bb870	[Feat]: add ASCII art logo for vLLM-Omni (#1430 )	2026-02-23 22:10:09 +08:00
Ding Zuhaoand汪志鹏	0ff443f069	[Feature][BAGEL] Combine multi-branch cfg into a single batch to accelerate inference. (#1429 ) Signed-off-by: Ding Zuhao <e1583181@u.nus.edu> Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>	2026-02-23 17:41:47 +08:00
TJian	5722333fba	[ROCM] [CI] fix dockerfile.rocm to support nightly build and also fix amd ci v0.16.0rc1 (#1380 ) Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>	2026-02-23 15:00:12 +08:00
Zhou Taichang	b42287d14b	[Doc] Update GPU installation commands (#1434 )	2026-02-23 13:54:57 +08:00
Hongsheng Liu	a3f2d4cbfe	Revert "[misc] Feature/pr reviewer auto trigger&update model" (#1432 )	2026-02-23 07:52:09 +08:00
Hongsheng LiuandClaude Opus 4.6	267170115b	[misc] Feature/pr reviewer auto trigger&update model (#1431 ) Signed-off-by: hsliu <liuhongsheng4@huawei.com> Signed-off-by: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Hunter Liu <hunter@liu.sh> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>	2026-02-22 23:33:45 +08:00
UsamaKenway	d8a32a6a87	[Misc] Add Qwen2.5-Omni-3B model support to Gradio demo (#1382 ) Signed-off-by: UsamaKenway <usamakenway@gmail.com>	2026-02-22 21:17:09 +08:00
Hongsheng LiuandClaude Opus 4.6	6c85ddcb96	[Feature] Add automated PR reviewer bot with GLM integration (#1424 ) Signed-off-by: hsliu <liuhongsheng4@huawei.com> Signed-off-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>	2026-02-22 13:53:59 +08:00
Sy03andHongsheng Liu	efbe4116ff	[Feature] : Support disaggregated inference pipeline for Qwen3_TTS (#1161 ) Signed-off-by: Sy03 <1370724210@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>	2026-02-21 07:36:44 +08:00
Alex Brooks	5295f0f81b	[Bugfix] Fix Hardcoded Datatypes in Z-image (#1393 ) Signed-off-by: Alex Brooks <albrooks@redhat.com>	2026-02-20 06:55:24 +00:00
Yan MaandDaniel Huang	4f9731c920	[XPU] Add XPU Dockerfile and related docs (#1162 ) Signed-off-by: Yan Ma <yan.ma@intel.com> Signed-off-by: Daniel Huang <daniel1.huang@intel.com> Co-authored-by: Daniel Huang <daniel1.huang@intel.com>	2026-02-20 13:33:48 +08:00
Chendi.Xue	d2f015edce	[0.16.0] remove cuda hard-code for Hunyuan Image3 (#1402 ) Signed-off-by: Chendi Xue <chendi.xue@intel.com>	2026-02-20 13:23:31 +08:00
Chendi.Xue	0b577a7263	[v0.16.0][BUG FIX]Fix hunyuan MOE after update to 0.16.0 (#1401 ) Signed-off-by: Chendi Xue <chendi.xue@intel.com>	2026-02-19 23:57:11 +13:00
wangyu	a42b7487b0	[Test] Fix expansion and example test case for qwen3-omni (#1358 ) Signed-off-by: yenuo26 <410167048@qq.com>	2026-02-15 00:43:38 +13:00