Merge branch 'main' into fix/tp-local-kv-recv-consensus

Signed-off-by: lishunyang12 <lishunyang03@gmail.com>
Signed-off-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: lishunyang12 <lishunyang03@gmail.com>

.buildkite/amd/test-amd-merge.yml

@@ -89,7 +89,7 @@ steps:
       mirror_hardwares: [amdproduction]
       grade: Blocking
       commands:
-        - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
+        - pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py
 
     - label: "Diffusion Model Test"
       agent_pool: mi325_1

.buildkite/amd/test-amd-ready.yml

@@ -35,7 +35,7 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
+    - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
 
 
 - label: "Voxtral TTS CUDA Unit Test"

.buildkite/cuda/test-merge.yml

@@ -8,14 +8,24 @@ steps:
   - group: ":card_index_dividers: Simple Test"
     depends_on: upload-merge-pipeline
     steps:
-      - label: "Simple · Diffusion & Model Executor Test"
+      - label: "Simple · Model Executor Test"
         commands:
-          - "timeout 20m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        mirror_hardwares: l4_1
+
+      - label: "Simple · Diffusion Test"
+        commands:
+          - "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        mirror_hardwares: l4_1
+
+      - label: "Simple · Engine&Entrypoints Test"
+        commands:
+          - "timeout 20m pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
         mirror_hardwares: l4_1
 
       - label: "Simple · Other Test"
         commands:
-          - "timeout 20m pytest -sv -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "timeout 20m pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
         mirror_hardwares: l4_1
 
   - group: ":card_index_dividers: Diffusion Test"
@@ -27,12 +37,6 @@ steps:
           - pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model"
         mirror_hardwares: l4_1
 
-      - label: "Diffusion · Offloader Test"
-        timeout_in_minutes: 40
-        commands:
-          - timeout 40m pytest -s -v tests/diffusion/offloader -m "advanced_model and cuda"
-        mirror_hardwares: l4_1
-
   - label: "Distributed Test"
     depends_on: upload-merge-pipeline
     timeout_in_minutes: 60
@@ -135,39 +139,32 @@ steps:
         timeout_in_minutes: 50
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'advanced_model and cuda' --run-level 'advanced_model'
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 2
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                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
+        mirror_hardwares: h100_2
+
+      - label: "Omni · MiniCPM-o 4.5 Duplex Test"
+        source_file_dependencies:
+          - tests/assets/minicpmo_4_5/response_required_16k.wav
+          - tests/assets/minicpmo_4_5/soft_interrupt_16k.wav
+          - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
+          - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
+          - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
+          - tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py
+          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
+          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
+          - tests/helpers/minicpmo_4_5_duplex.py
+          - vllm_omni/deploy/minicpmo_4_5_duplex.yaml
+          - vllm_omni/experimental/fullduplex/
+          - vllm_omni/model_executor/models/minicpmo_4_5/
+          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
+          - vllm_omni/core/sched/
+          - vllm_omni/outputs/multimodal_accumulation.py
+          - vllm_omni/engine/stage_init_utils.py
+          - vllm_omni/engine/async_omni_engine.py
+          - vllm_omni/entrypoints/cli/serve.py
+        timeout_in_minutes: 50
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
+        mirror_hardwares: h100_2
 
       - label: "TTS · VoxCPM2 Test"
         source_file_dependencies:
@@ -185,22 +182,6 @@ steps:
             "
         mirror_hardwares: l4_1
 
-      - label: "TTS · SoulX-Singer Test"
-        source_file_dependencies:
-          - tests/e2e/offline_inference/test_soulxsinger.py
-          - vllm_omni/diffusion/models/soulx_singer/
-          - vllm_omni/deploy/soulxsinger_svc.yaml
-          - vllm_omni/deploy/soulxsinger_svs.yaml
-        timeout_in_minutes: 30
-        commands:
-          - |
-            timeout 30m bash -c "
-              export VLLM_LOGGING_LEVEL=DEBUG
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/offline_inference/test_soulxsinger.py -k 'svc or precomputed' -m 'advanced_model and cuda' --run-level 'advanced_model'
-            "
-        mirror_hardwares: l4_1
-
       - label: "Diffusion · Qwen Image Edit Test"
         source_file_dependencies:
           - tests/e2e/online_serving/test_qwen_image_edit.py

.buildkite/cuda/test-nightly.yml

@@ -15,7 +15,13 @@ steps:
       - label: ":full_moon: Omni · Function Test with H100"
         timeout_in_minutes: 120
         commands:
-          - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
+          - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
+        mirror_hardwares: h100_2
+
+      - label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test"
+        timeout_in_minutes: 50
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni" --run-level "full_model"
         mirror_hardwares: h100_2
 
       - label: ":full_moon: Omni · Doc Test with L4"
@@ -449,36 +455,35 @@ steps:
         mirror_hardwares: h100_2
 
   - group: ":card_index_dividers: Diffusion Test"
+    key: nightly-diffusion-test-group
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test"
     steps:
+      - label: ":full_moon: Diffusion · Distributed Test"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda' --run-level "full_model"
+        mirror_hardwares: l4_4
 
       - label: ":full_moon: Diffusion · Offloader Test"
-        timeout_in_minutes: 40
+        timeout_in_minutes: 60
         commands:
-          - timeout 40m pytest -s -v tests/diffusion/offloader -m "full_model and cuda"
+          - pytest -sv tests/diffusion/offloader -m "full_model and cuda"
         mirror_hardwares: l4_1
 
       - label: ":full_moon: Diffusion · Quantization Test with H100"
+        timeout_in_minutes: 60
         commands:
-          - timeout 60m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
+          - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
         mirror_hardwares: h100_1
 
       - label: ":full_moon: Diffusion · Quantization Test with L4"
         timeout_in_minutes: 60
         commands:
+          - pip install "vllm-gguf-plugin>=0.0.3"
           - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
         mirror_hardwares: l4_1
 
-      - label: ":full_moon: Diffusion · Wan VAE Spatial Shard Correctness Test"
-        timeout_in_minutes: 60
-        commands:
-          - >-
-            pytest -sv tests/diffusion/distributed/test_wan_spatial_shard.py
-            -m "full_model and diffusion and H100 and distributed_cuda"
-            --run-level "full_model"
-        mirror_hardwares: h100_2
-
   - label: ":bar_chart: Testcase Statistics"
     key: nightly-testcase-statistics
     timeout_in_minutes: 120

.buildkite/cuda/test-ready.yml

@@ -8,22 +8,26 @@ steps:
   - group: ":card_index_dividers: Simple Test"
     depends_on: upload-ready-pipeline
     steps:
-      - label: "Simple · Diffusion & Model Executor Test"
+      - label: "Simple · Model Executor Test"
         commands:
-          - "timeout 20m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        mirror_hardwares: l4_1
+
+      - label: "Simple · Diffusion Test"
+        commands:
+          - "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        mirror_hardwares: l4_1
+
+      - label: "Simple · Engine&Entrypoints Test"
+        commands:
+          - "timeout 20m pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
         mirror_hardwares: l4_1
 
       - label: "Simple · Other Test"
         commands:
-          - "timeout 20m pytest -sv -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "timeout 20m pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
         mirror_hardwares: l4_1
 
-  - label: "Custom Pipeline Test"
-    depends_on: upload-ready-pipeline
-    commands:
-      - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
-    mirror_hardwares: l4_1
-
   - group: ":card_index_dividers: Diffusion Test"
     depends_on: upload-ready-pipeline
     steps:
@@ -104,11 +108,34 @@ steps:
   - group: ":card_index_dividers: E2E Test"
     depends_on: upload-ready-pipeline
     steps:
+      - label: "Custom Pipeline Test"
+        source_file_dependencies:
+          - tests/e2e/features/custom_pipeline/
+          - tests/e2e/features/helpers/custom_pipeline.py
+          - vllm_omni/diffusion/
+          - vllm_omni/engine/
+          - vllm_omni/entrypoints/
+          - vllm_omni/inputs/
+          - vllm_omni/outputs/
+        commands:
+          - timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
+        agents:
+          queue: "gpu_1_queue"
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
 
       - label: "RLHF · VeRL-Omni E2E Test"
         source_file_dependencies:
-          - tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
-          - tests/e2e/offline_inference/custom_pipeline/
+          - tests/e2e/features/rlhf_test/test_verl_omni_e2e.py
+          - tests/e2e/features/helpers/verl_omni_server.py
+          - tests/e2e/features/helpers/custom_pipeline.py
           - vllm_omni/diffusion/
           - vllm_omni/engine/
           - vllm_omni/entrypoints/
@@ -117,7 +144,7 @@ steps:
         commands:
           - |
             timeout 20m bash -c "
-              pytest -s -v tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
+              pytest -s -v tests/e2e/features/rlhf_test/test_verl_omni_e2e.py
             "
         mirror_hardwares: l4_1
 
@@ -143,34 +170,7 @@ steps:
             timeout 20m bash -c "
               pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py -m 'core_model' --run-level 'core_model'
             "
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 2
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
+        mirror_hardwares: h100_2
 
       - label: "Omni · MiniCPM-o 4.5 Online Test"
         source_file_dependencies:
@@ -181,34 +181,34 @@ steps:
             timeout 20m bash -c "
               pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5.py -m 'core_model' --run-level 'core_model'
             "
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 2
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
+        mirror_hardwares: h100_2
+
+      - label: "Omni · MiniCPM-o 4.5 Duplex Test"
+        source_file_dependencies:
+          - tests/assets/minicpmo_4_5/response_required_16k.wav
+          - tests/assets/minicpmo_4_5/soft_interrupt_16k.wav
+          - tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
+          - tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
+          - tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
+          - tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py
+          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
+          - tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
+          - tests/helpers/minicpmo_4_5_duplex.py
+          - vllm_omni/deploy/minicpmo_4_5_duplex.yaml
+          - vllm_omni/experimental/fullduplex/
+          - vllm_omni/model_executor/models/minicpmo_4_5/
+          - vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
+          - vllm_omni/core/sched/
+          - vllm_omni/outputs/multimodal_accumulation.py
+          - vllm_omni/engine/stage_init_utils.py
+          - vllm_omni/engine/async_omni_engine.py
+          - vllm_omni/entrypoints/cli/serve.py
+        commands:
+          - |
+            timeout 20m bash -c "
+              pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and cuda' --run-level 'core_model'
+            "
+        mirror_hardwares: h100_2
 
       - label: "TTS · VoxCPM2 Test"
         source_file_dependencies:
@@ -344,13 +344,3 @@ steps:
               pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model'
             "
         mirror_hardwares: h100_1
-
-      - label: "Diffusion · GGUF Plugin Test"
-        soft_fail: true
-        source_file_dependencies:
-          - tests/plugin/test_gguf_diffusion.py
-          - vllm_omni/diffusion/model_loader
-        commands:
-          - pip install "vllm-gguf-plugin>=0.0.3"
-          - timeout 30m pytest -sv tests/plugin/test_gguf_diffusion.py
-        mirror_hardwares: l4_1

.buildkite/intel/pipeline-intel.yml

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.25.0"
+          VLLM_VERSION: "v0.26.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: false

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py

@@ -1,12 +1,8 @@
 from collections.abc import Callable
 from enum import Enum, auto
-from typing import (
-    Any,
-    Literal,
-    NotRequired,
-    TypeAlias,
-    TypedDict,
-)
+from typing import Any, Literal, TypeAlias
+
+from typing_extensions import NotRequired, TypedDict
 
 AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]
 

docker/Dockerfile.npu

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.25.0
+ARG VLLM_ASCEND_TAG=v0.26.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm
@@ -26,6 +26,10 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
     python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
+RUN python3 -m pip install -v --no-deps \
+      "stepaudio2-minicpmo>=0.1.1" \
+      "step-audio2>=1.0.0"
+
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 
 ENTRYPOINT []

docker/Dockerfile.npu.a3

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.25.0-a3
+ARG VLLM_ASCEND_TAG=v0.26.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm
@@ -26,6 +26,10 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi
     export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
     python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
+RUN python3 -m pip install -v --no-deps \
+      "stepaudio2-minicpmo>=0.1.1" \
+      "step-audio2>=1.0.0"
+
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 
 ENTRYPOINT []

docker/Dockerfile.npu.ci

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.25.0
+ARG VLLM_ASCEND_TAG=v0.26.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
@@ -35,6 +35,7 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
       "psutil>=7.2.0" \
       "soundfile>=0.13.1" \
       "pyttsx3>=2.99" \
+      "opencc==1.4.1" \
       "opencv-python>=4.12.0.88" \
       "imageio[ffmpeg]>=0.6.0" \
       "jiwer>=3.0.0" \

docker/Dockerfile.npu.ci.a3

@@ -1,5 +1,5 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.25.0-a3
+ARG VLLM_ASCEND_TAG=v0.26.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
@@ -35,6 +35,7 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
       "psutil>=7.2.0" \
       "soundfile>=0.13.1" \
       "pyttsx3>=2.99" \
+      "opencc==1.4.1" \
       "opencv-python>=4.12.0.88" \
       "imageio[ffmpeg]>=0.6.0" \
       "jiwer>=3.0.0" \

docs/contributing/ci/test_system_overview.md

@@ -56,7 +56,11 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td>/</td>
       <td><code>core_model and cpu</code></td>
       <td rowspan="2">&lt;15min</td>
-      <td>/tests/{component_name}/test_xxx</td>
+      <td>
+        <code>tests/{component}/…</code> mirroring <code>vllm_omni/{component}/</code><br>
+        (e.g. <code>tests/diffusion/</code>, <code>tests/engine/</code>, <code>tests/entrypoints/</code>)<br>
+        Do <strong>not</strong> add new top-level <code>tests/</code> dirs unrelated to a component.
+      </td>
       <td>
         <a href="test_writing_guide.md#l1--l2-level-testing-unit-testing-and-basic-end-to-end-verification">L1 &amp; L2</a><br>
         Section 1 L1&amp;L2: Purpose, Test Content, Directory Location, Example
@@ -68,16 +72,18 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td><strong>L2</strong><br>(E2E across models & GPU-required UT)</td>
       <td>Online (basic deployment scenarios):<br>dummy, normal inference function (output format, stream), some instance startup UT</td>
       <td>High-priority models + online basic scenarios; request success, non-empty output, format match (no Whisper/accuracy)</td>
-      <td>High-priority features (using random lightweight models)</td>
+      <td>High-priority features and component tests that require launching instances</td>
       <td>High-priority interfaces (using random lightweight models)</td>
       <td><code>core_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
       <td>
-        <strong>Model tests:</strong><br>
-        /tests/e2e/online_serving/test_{model_name}.py<br>
-        <strong>Feature tests:</strong><br>
-        /tests/{component_name}/test_xxx<br>
-        <strong>Interface tests:</strong><br>
-        /tests/entrypoints/test_xxx
+        <strong>Model E2E:</strong><br>
+        <code>tests/e2e/online_serving/test_{model}.py</code><br>
+        <code>tests/e2e/offline_inference/test_{model}.py</code><br>
+        <strong>Feature integration:</strong><br>
+        <code>tests/e2e/features/&lt;feature&gt;/</code><br>
+        (e.g. <code>fullduplex/</code>, <code>custom_pipeline/</code>, <code>rlhf_test/</code>)<br>
+        <strong>Component / interface Test (GPU):</strong><br>
+        <code>tests/{component}/…</code>, <code>tests/entrypoints/…</code>
       </td>
       <td>
         <a href="test_writing_guide.md#l1--l2-level-testing-unit-testing-and-basic-end-to-end-verification">L1 &amp; L2</a><br>
@@ -90,18 +96,19 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td><strong>L3</strong><br>(Important Perf & Integration & Accuracy)</td>
       <td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
       <td>High/medium-priority models + key online/offline scenarios; real weights, Whisper/similarity, preset voice gender, basic accuracy</td>
-      <td>Medium-priority features (using random lightweight models)</td>
+      <td>Medium-priority features and component tests that require launching instances</td>
       <td>Medium-priority interfaces (using random lightweight models)</td>
       <td><code>advanced_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
       <td>&lt;30min</td>
       <td>
-        <strong>Model tests:</strong><br>
-        /tests/e2e/online_serving/test_{model_name}.py<br>
-        /tests/e2e/offline_inference/test_{model_name}.py<br>
-        <strong>Feature tests:</strong><br>
-        /tests/{component_name}/test_xxx<br>
-        <strong>Interface tests:</strong><br>
-        /tests/entrypoints/test_xxx
+        <strong>Model E2E:</strong><br>
+        <code>tests/e2e/online_serving/test_{model}.py</code><br>
+        <code>tests/e2e/offline_inference/test_{model}.py</code><br>
+        <code>tests/e2e/accuracy/</code><br>
+        <strong>Feature integration:</strong><br>
+        <code>tests/e2e/features/&lt;feature&gt;/</code><br>
+        <strong>Component / interface Test:</strong><br>
+        <code>tests/{component}/…</code>, <code>tests/entrypoints/…</code>
       </td>
       <td>
         <a href="test_writing_guide.md#l3-level-testing-core-integration-performance-and-accuracy-verification">L3</a><br>
@@ -114,26 +121,23 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td><strong>L4</strong><br>(Perf & Integration & Accuracy)</td>
       <td>Online: full functional scenarios + performance test + doc test + accuracy test</td>
       <td>High-priority models: function, performance, accuracy, and doc testing<br>Medium-priority models: function and doc testing</td>
-      <td>Low-priority features (using real weights)</td>
+      <td>Low-priority features and component tests that require launching instances</td>
       <td>Low-priority interfaces (using real weights)</td>
       <td><code>full_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
       <td>&lt;3 hour</td>
       <td>
-        <strong>Model tests:</strong><br>
-        /tests/e2e/online_serving/test_{model_name}_expansion.py<br>
-        <strong>Feature tests:</strong><br>
-        /tests/{component_name}/test_xxx<br>
-        <strong>Interface tests:</strong><br>
-        /tests/entrypoints/test_xxx<br>
+        <strong>Model E2E:</strong><br>
+        <code>tests/e2e/online_serving/test_{model}_expansion.py</code><br>
+        <code>tests/e2e/offline_inference/test_{model}_expansion.py</code><br>
+        <code>tests/e2e/accuracy/test_{model}.py</code><br>
+        <strong>Feature integration:</strong><br>
+        <code>tests/e2e/features/&lt;feature&gt;/</code><br>
+        <strong>Component / interface Test:</strong><br>
+        <code>tests/{component}/…</code>, <code>tests/entrypoints/…</code><br>
         <strong>Performance:</strong><br>
-        /tests/dfx/perf/tests/test_qwen3_omni_*.json (Omni), test_tts.json (TTS),<br>
-        test_voxcpm2.json, test_higgs_audio_v3.json, and<br>
-        /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)<br>
-        <strong>Doc Test:</strong><br>
-        tests/examples/online_serving/test_{model_name}.py<br>
-        tests/examples/offline_inference/test_{model_name}.py<br>
-        <strong>Accuracy Test:</strong><br>
-        /tests/e2e/accuracy/test_{model_name}.py
+        <code>tests/dfx/perf/tests/</code><br>
+        <strong>Doc examples:</strong><br>
+        <code>tests/examples/online_serving/</code>, <code>tests/examples/offline_inference/</code>
       </td>
       <td>
         <a href="test_writing_guide.md#l4-level-testing-full-functionality-performance-and-documentation-testing">L4</a><br>
@@ -152,11 +156,9 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td> Depends on reality </td>
       <td>
         <strong>Stability:</strong><br>
-        /tests/dfx/stability/tests/test_qwen3_omni.json<br>
-        /tests/dfx/stability/tests/test_wan22.json<br>
+        <code>tests/dfx/stability/tests/</code><br>
         <strong>Reliability:</strong><br>
-        tests/dfx/reliability/test_reliability_{model_key}.py<br>
-        (e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>, <code>test_reliability_hunyuan_image.py</code>, <code>test_reliability_voxcpm2.py</code>)
+        <code>tests/dfx/reliability/test_reliability_{model}.py</code>
       </td>
       <td>
         <a href="test_writing_guide.md#l5-level-testing-stability-and-reliability-testing">L5</a><br>
@@ -168,7 +170,9 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
   </tbody>
 </table>
 
-For per-level test authoring (directories, markers, examples), see [Test Writing Guide](./test_writing_guide.md).
+**Test Dir placement (summary):** component / unit under `tests/{component}/` mirroring `vllm_omni/`; model E2E under `tests/e2e/online_serving/`, `tests/e2e/offline_inference/`, and `tests/e2e/accuracy/`; feature integration under `tests/e2e/features/<feature>/`; doc example tests under `tests/examples/online_serving/` and `tests/examples/offline_inference/`; performance under `tests/dfx/perf/`; stability under `tests/dfx/stability/`; reliability under `tests/dfx/reliability/`. Do **not** add new top-level directories under `tests/` that are unrelated to a `vllm_omni` component (or to the established `e2e` / `dfx` / `helpers` / `examples` / `buildkite` layout).
+
+For per-level test authoring (markers, examples), see [Test Writing Guide](./test_writing_guide.md).
 
 ## Common Specifications
 

docs/contributing/ci/test_writing_guide.md

@@ -4,137 +4,20 @@ This document describes how to author tests for each CI level (L1–L5). For the
 
 ## Test Directory Structure
 
-Legend: `✅` = test exists, `⬜` = suggested to add.
-```
-vllm_omni/                                    tests/
-├── config/                             →     ├── config/
-│   ├── model.py                              │   └── test_model.py                    ⬜
-│   └── lora.py                               │   └── test_lora.py                      ⬜
-│
-├── core/                               →     ├── core/
-│   └── sched/                                 │   └── sched/
-│       ├── omni_ar_scheduler.py               │       ├── test_omni_ar_scheduler.py    ⬜
-│       ├── omni_generation_scheduler.py       │       ├── test_omni_generation_scheduler.py  ⬜
-│       └── output.py                          │       └── test_output.py               ✅ currently in entrypoints/test_omni_new_request_data.py (tests output.OmniNewRequestData)
-│
-├── diffusion/                          →     ├── diffusion/
-│   ├── diffusion_engine.py                    │   ├── test_diffusion_engine.py          ⬜
-│   ├── attention/                             │   ├── attention/
-│   │   ├── layer.py                            │   │   ├── test_attention_sp.py         ✅
-│   │   └── backends/                           │   │   └── test_flash_attn.py           ✅
-│   ├── distributed/                           │   ├── distributed/
-│   │   └── ...                                 │   │   ├── test_comm.py                 ✅
-│   │                                            │   │   ├── test_cfg_parallel.py        ✅
-│   │                                            │   │   └── test_sp_plan_hooks.py       ✅
-│   ├── lora/                                   │   ├── lora/
-│   │   └── ...                                 │   │   ├── test_base_linear.py          ✅
-│   │                                            │   │   └── test_lora_manager.py        ✅
-│   ├── models/                                 │   ├── models/
-│   │   ├── qwen_image/                         │   │   ├── qwen_image/                 (e2e coverage)
-│   │   ├── ovis_image/                         │   │   ├── ovis_image/
-│   │   │   └── ...                             │   │   │   └── test_ovis_image.py     ✅
-│   │   ├── z_image/                            │   │   └── z_image/
-│   │   └── ...                                 │   │       └── test_zimage_tp_constraints.py  ✅
-│   └── worker/                                 │   └── worker/
-│       ├── diffusion_worker.py                 │       └── test_diffusion_worker.py   ✅ file at diffusion/test_diffusion_worker.py
-│       └── diffusion_model_runner.py            │
-│
-├── distributed/                         →     ├── distributed/
-│   └── omni_connectors/                         │   └── omni_connectors/
-│       ├── adapter.py                           │       ├── test_adapter_and_flow.py   ✅
-│       ├── kv_transfer_manager.py               │       ├── test_basic_connectors.py   ✅
-│       ├── connectors/                           │       ├── test_kv_flow.py             ✅
-│       └── utils/                               │       └── test_omni_connector_configs.py  ✅
-│
-├── engine/                             →     ├── engine/
-│   ├── input_processor.py                      │   ├── test_input_processor.py         ⬜  (no processor.py in source)
-│   ├── output_processor.py                     │   └── test_output_processor.py         ⬜
-│   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
-│
-├── entrypoints/                        →     ├── entrypoints/
-│   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
-│   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
-│   │   └── ...                                  │   │   └── test_*.py                  ⬜
-│   └── openai/                                  │   └── openai_api/                    # maps to entrypoints/openai/
-│       ├── api_server.py                        │       ├── test_api_server.py         ⬜  (e2e indirect coverage)
-│       ├── serving_chat.py                       │       ├── test_serving_chat_sampling_params.py  ✅
-│       ├── serving_speech.py                     │       ├── test_serving_speech.py     ✅
-│       └── image_api_utils.py                   │       └── test_image_server.py      ✅
-│
-├── inputs/                             →     ├── inputs/
-│   ├── data.py                                 │   ├── test_data.py                   ⬜
-│   ├── parse.py                                │   ├── test_parse.py                 ⬜
-│   └── preprocess.py                            │   └── test_preprocess.py            ✅ currently in entrypoints/test_omni_input_preprocessor.py
-│
-├── model_executor/                     →     ├── model_executor/
-│   ├── layers/                                  │   ├── layers/
-│   │   └── mrope.py                             │   │   └── test_mrope.py              ⬜
-│   ├── model_loader/                            │   ├── model_loader/
-│   │   └── weight_utils.py                      │   │   └── test_weight_utils.py      ⬜
-│   ├── models/                                  │   ├── models/
-│   │   ├── qwen2_5_omni/                         │   │   ├── qwen2_5_omni/
-│   │   │   ├── qwen2_5_omni_thinker.py           │   │   │   ├── test_audio_length.py  ✅
-│   │   │   ├── qwen2_5_omni_talker.py            │   │   │   ├── test_qwen2_5_omni_thinker.py  ⬜
-│   │   │   └── qwen2_5_omni_token2wav.py         │   │   │   ├── test_qwen2_5_omni_talker.py  ⬜
-│   │   └── qwen3_omni/                          │   │   │   └── test_qwen2_5_omni_token2wav.py  ⬜
-│   │       └── ...                               │   │   └── qwen3_omni/
-│   ├── stage_configs/                           │   │       └── test_*.py              ⬜
-│   │   └── *.yaml                               │   └── stage_configs/                 (used by e2e, test_*.py can be added)  ⬜
-│   └── stage_input_processors/                  │   └── stage_input_processors/
-│       └── ...                                  │       └── test_*.py                 ⬜
-│
-├── sample/                             →     ├── sample/
-│   └── __init__.py                             │   └── test_*.py                      ⬜
-│
-├── utils/                              →     ├── utils/
-│   └── __init__.py                             │   └── test_*.py                       ⬜  (no platform_utils.py currently)
-│
-├── worker/                             →     ├── worker/
-│   ├── gpu_ar_model_runner.py                  │   ├── test_gpu_ar_model_runner.py    ⬜
-│   ├── gpu_ar_worker.py                        │   ├── test_gpu_ar_worker.py           ⬜
-│   ├── gpu_generation_model_runner.py          │   ├── test_gpu_generation_model_runner.py  ✅
-│   ├── gpu_generation_worker.py                │   ├── test_gpu_generation_worker.py  ⬜
-│   ├── gpu_model_runner.py                     │   ├── test_omni_gpu_model_runner.py   ✅
-│   └── mixins.py                               │   └── (npu under platforms/npu/worker/)  # not worker/npu/
-│
-├── platforms/                          →     (no tests/platforms/, e2e and stage_configs provide indirect coverage)
-│   ├── cuda/
-│   ├── npu/worker/                             # NPU worker here, not vllm_omni/worker/npu/
-│   ├── rocm/
-│   └── xpu/worker/
-│
-├── outputs.py                          →     test_outputs.py                         ✅ (at tests root)
-├── (logger, patch, request, version)    →     (no corresponding unit test)
-│
-└── e2e (tests side only)               →     ├── e2e/
-                                               ├── online_serving/                     ✅ non-empty
-                                               │   ├── test_qwen2_5_omni.py
-                                               │   ├── test_async_omni.py
-                                               │   ├── test_qwen3_omni.py
-                                               │   ├── test_qwen3_omni_expansion.py
-                                               │   ├── test_mimo_audio.py
-                                               │   └── test_images_generations_lora.py
-                                               └── offline_inference/                  ✅
-                                                   ├── test_qwen2_5_omni.py
-                                                   ├── test_qwen3_omni.py
-                                                   ├── test_bagel_text2img.py
-                                                   ├── test_z_image.py
-                                                   ├── test_wan22.py
-                                                   ├── test_zimage_tensor_parallel.py
-                                                   ├── test_cache_dit.py
-                                                   ├── test_teacache.py
-                                                   ├── test_stable_audio_expansion.py
-                                                   ├── test_diffusion_cpu_offload.py
-                                                   ├── test_diffusion_layerwise_offload.py
-                                                   ├── test_diffusion_lora.py
-                                                   ├── test_sequence_parallel.py
-                                                   └── stage_configs/                  (legacy schema, still
-                                                       ├── bagel_*.yaml                 present for unmigrated
-                                                       └── npu/, rocm/, etc.            models)
+Canonical placement by CI level is in the **Test Dir** column of the level matrix in [Test System Overview](./test_system_overview.md).
 
-# Migrated models (qwen3_omni_moe, qwen2_5_omni, qwen3_tts) live under
-# vllm_omni/deploy/ instead — see docs/configuration/stage_configs.md.
-```
+Summary:
+
+- **Component / unit:** `tests/{component}/…` mirroring `vllm_omni/{component}/`
+- **Model E2E:** `tests/e2e/online_serving/`, `tests/e2e/offline_inference/`, `tests/e2e/accuracy/`
+- **Feature integration:** `tests/e2e/features/<feature>/`
+- **Doc example tests:** `tests/examples/online_serving/`, `tests/examples/offline_inference/`
+- **Performance:** `tests/dfx/perf/tests/` (JSON configs; runners under `tests/dfx/perf/scripts/`)
+- **Stability:** `tests/dfx/stability/tests/`
+- **Reliability:** `tests/dfx/reliability/`
+- Do **not** add new top-level `tests/` directories unrelated to a component (or to `e2e` / `dfx` / `helpers` / `examples` / `buildkite`)
+
+Deploy YAMLs: `vllm_omni/deploy/` via `tests.helpers.stage_config.get_deploy_config_path` (see [stage configs](../../configuration/stage_configs.md)).
 
 ## Markers for Tests
 
@@ -373,13 +256,13 @@ L2 level testing builds upon L1 by introducing GPU resources and verifying that
 
 #### 1.3 Test Directory and Execution Files
 
-A clear directory structure is key to managing test cases efficiently.
+A clear directory structure is key to managing test cases efficiently. See the **Test Dir** column in [Test System Overview](./test_system_overview.md) for the canonical placement rules.
 
--   ***L1 Test Directory***: `/tests/{component_name}/test_xxx.py`
--   -   Here, `{component_name}` corresponds to modules in the source code, such as `distributed`, `entrypoints`, etc., and `test_xxx.py` is the specific test file.
--   ***L2 Test Directory***:
--   -   Online Serving: `/tests/e2e/online_serving/test_{model_name}.py`
-    -   Offline Inference: `/tests/e2e/offline_inference/test_{model_name}.py`
+-   ***L1 (component / unit)***: `tests/{component_name}/…` mirroring `vllm_omni/{component_name}/` (e.g. `tests/diffusion/`, `tests/entrypoints/`). Do **not** invent new top-level `tests/` folders unrelated to a component.
+-   ***L2 (basic model E2E)***:
+    -   Online serving: `tests/e2e/online_serving/test_{model_name}.py`
+    -   Offline inference: `tests/e2e/offline_inference/test_{model_name}.py`
+-   ***Feature-level integration*** (cross-cutting, not a single-model smoke): `tests/e2e/features/<feature>/` (e.g. `fullduplex/`, `custom_pipeline/`, `rlhf_test/`).
 
 #### 1.4 Execution Method and Example
 
@@ -485,10 +368,12 @@ L3 level testing executes after code is merged into the main branch. Its core pu
 
 #### 2.3 Test Directory and Execution Files
 
--   ***Functional Testing***:
--   -   Online Serving: `/tests/e2e/online_serving/test_{model_name}_expansion.py`
-    -   Offline Inference: `/tests/e2e/offline_inference/test_{model_name}_expansion.py`
-    -   (Note: `_expansion.py` likely means it contains more comprehensive scenario cases compared to L2 tests).
+-   ***Functional / model E2E***:
+    -   Online serving: `tests/e2e/online_serving/test_{model_name}_expansion.py`
+    -   Offline inference: `tests/e2e/offline_inference/test_{model_name}_expansion.py`
+    -   Accuracy / similarity: `tests/e2e/accuracy/`
+    -   (`_expansion.py` holds broader scenario coverage than the L2 smoke file.)
+-   ***Feature-level integration***: `tests/e2e/features/<feature>/` (not under a new top-level `tests/` directory).
 
 #### 2.4 Execution Method and Example
 

docs/design/feature/async_diffusion_output.md

@@ -0,0 +1,174 @@
+# Async Diffusion Output
+
+## Table of Contents
+
+1. [Overview](#overview)
+2. [Performance](#performance)
+3. [Architecture](#architecture)
+4. [Model Coverage](#model-coverage)
+5. [Configuration](#configuration)
+6. [Related Files](#related-files)
+
+## Overview
+
+The async diffusion output feature moves the D2H (Device→Host) copy and SHM packing from the Worker main thread to a background daemon thread, so the GPU default stream can immediately start the next request's forward pass. This eliminates the GPU bubble caused by synchronous D2H/packing.
+
+**Core benefit**: the previous request's output D2H/packing overlaps with the next request's forward pass.
+
+**Automatically enabled** when `step_execution=False` (default, request-mode). No extra configuration needed. When `step_execution=True` (step-mode), the original synchronous path is used — pump thread and Worker background thread are not started.
+
+## Performance
+
+### HunyuanImage-3.0 (TP4)
+
+| Resolution | Async Off (QPS) | Async On (QPS) | Change |
+|---|---|---|---|
+| 1024×1024 | 0.4773 | 0.4802 | +0.60% |
+| 768×768 | 0.8370 | 0.8533 | +1.95% |
+
+**Analysis**: smaller resolution → larger QPS gain. D2H/packing is relatively fixed overhead; forward is variable (scales with resolution). Smaller images → shorter forward → D2H/packing occupies a larger fraction of total time → more benefit from overlapping it.
+
+## Architecture
+
+### Execution Timeline
+
+```
+Before (synchronous D2H):
+[forward req1] [D2H+SHM pack req1] [forward req2] [D2H+SHM pack req2]
+                                 ^^^^^^^^^^^^^^^^ GPU bubble
+
+After (async D2H):
+[forward req1] [forward req2]
+[D2H+SHM pack req1 (side stream)] [D2H+SHM pack req2 (side stream)]
+                                 ^^^^^^^^^^^^^^^^ bubble eliminated
+```
+
+### Data Flow
+
+```
+                    ┌─────────────────────────────────────────────────────┐
+                    │                 Worker Process                      │
+                    │                                                    │
+  execute_model ──> │ WorkerBusyLoop ──> forward() ──> DiffusionOutput   │
+                    │       │                              │ (GPU tensor)│
+                    │       │                    ┌─────────┴──────────┐  │
+                    │       │                    │  async output path │  │
+                    │       │                    │                    │  │
+                    │       │              compute_done ──> result_mq │  │
+                    │       │                    │                    │  │
+                    │       │                    ▼                    │  │
+                    │       │            AsyncOutputThread            │  │
+                    │       │         (side CUDA stream)             │  │
+                    │       │         wait_event(gpu_event)          │  │
+                    │       │         pack_diffusion_output_shm()    │  │
+                    │       │         D2H + SHM write                │  │
+                    │       │                    │                    │  │
+                    │       │              output_ready ──> result_mq │  │
+                    │       │                    │                    │  │
+                    │       ▼                    ▼                    │  │
+                    │  (dequeue next request)                        │  │
+                    └─────────────────────────────────────────────────────┘
+                                         │
+                                         ▼
+                    ┌─────────────────────────────────────────────────────┐
+                    │              Executor (main process)               │
+                    │                                                    │
+                    │  ResultPumpThread (sole reader of result_mq)       │
+                    │       │                                            │
+                    │       ├── AsyncDiffusionOutput(COMPUTE_DONE)       │
+                    │       │   → resolve _rpc_futures[rpc_id]           │
+                    │       │   → RunnerOutput(async_output_id=...)      │
+                    │       │                                            │
+                    │       ├── AsyncDiffusionOutput(OUTPUT_READY)       │
+                    │       │   → resolve _output_futures[async_output_id]│
+                    │       │   → or batch split via _batch_split_map   │
+                    │       │                                            │
+                    │       └── Non-async message                        │
+                    │           → _sync_result_buffer (for other RPCs)   │
+                    │                                                    │
+                    │  wait_output_ready(async_output_id) → Future       │
+                    └─────────────────────────────────────────────────────┘
+```
+
+### Key Components
+
+1. **AsyncDiffusionOutput** (`data.py`): Protocol envelope for `result_mq`. `kind` field routes messages:
+   - `COMPUTE_DONE` — forward finished, GPU can start next request
+   - `OUTPUT_READY` — D2H/SHM packing finished, final output available
+   - `RPC_RESULT` — ordinary RPC return (including error propagation)
+
+2. **Worker AsyncOutputThread** (`diffusion_worker.py`): Background daemon thread that:
+   - Waits for `gpu_event` (cross-stream sync ensuring forward wrote the tensor)
+   - Runs `pack_diffusion_output_shm()` on a side CUDA stream
+   - Enqueues `OUTPUT_READY` to `result_mq` when packing completes
+
+3. **ResultPumpThread** (`multiproc_executor.py`): Sole reader of `result_mq` in request-mode. Dispatches messages:
+   - `COMPUTE_DONE` / `RPC_RESULT` → resolves `_rpc_futures[rpc_id]`
+   - `OUTPUT_READY` → resolves `_output_futures[async_output_id]` (with batch split via `_batch_split_map`)
+   - Non-async messages → `_sync_result_buffer` for other RPCs
+
+4. **collective_rpc() two-path dispatch** (`multiproc_executor.py`):
+   - **Path 1**: `execute_model` / `execute_model_batch` → generates `rpc_id`, registers Future, waits for pump to deliver `compute_done`
+   - **Path 2**: All other RPCs → `_sync_result_buffer` (pump-fed) or `_result_mq` (step-mode, no pump)
+
+5. **record_device_event** (`platforms/`): Cross-platform GPU event recording for side-stream synchronization. Implemented for CUDA, NPU, ROCm, XPU, MUSA; base class returns `None` (safe no-op).
+
+6. **IPC side-stream D2H** (`ipc.py`): `pack_diffusion_output_shm()` and helpers accept `d2h_stream` parameter; use `pin_memory` + `copy_(non_blocking=True)` on the side stream instead of synchronous `.cpu()`.
+
+### Batch Split
+
+When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently.
+
+## Model Coverage
+
+### Currently Supported (request-mode, `step_execution=False`)
+
+All models running in request-mode (`step_execution=False`, the default) automatically use async output.
+
+**Verified models**:
+
+| Model | Type | `supports_request_batch` |
+|---|---|---|
+| **HunyuanImage-3.0** | Image | `False` |
+| **Qwen-Image** | Image | `True` |
+
+Other models with `step_execution=False` are also supported but not yet verified.
+
+### Not Yet Supported (step-mode, `step_execution=True`)
+
+When `step_execution=True` (or `streaming_output=True`, which auto-enables step-mode), models use `execute_stepwise` instead of `execute_model`, which is not in the async Path 1 whitelist. Async output is not applicable.
+
+**Verified**: Helios (`step_execution=True`, does not benefit from this feature).
+
+Other models with `step_execution=True` are also not applicable.
+
+Adapting step-mode requires additional design because:
+
+- Each step's `RunnerOutput` must be available synchronously for the scheduler to advance
+- Intermediate latent tensors may not need D2H (only the final step needs D2H + postprocess)
+- Steps are serial (step N+1 depends on step N output), so the overlap benefit is limited to "final step's D2H overlaps with the next request's forward"
+
+**Planned approach**: first cover only the final step's async D2H, keep intermediate steps synchronous. This matches the request-mode benefit pattern.
+
+## Configuration
+
+No configuration needed. Async output is automatically enabled when `step_execution=False` (default).
+
+| `step_execution` | Mode | Async Output | Pump Thread | Worker BG Thread |
+|---|---|---|---|---|
+| `False` (default) | request-mode | ✅ enabled | ✅ started | ✅ started |
+| `True` | step-mode | ❌ disabled | ❌ not started | ❌ not started |
+
+## Related Files
+
+- `vllm_omni/diffusion/data.py`: `AsyncDiffusionOutput`, `AsyncOutputKind`, `DiffusionOutput.async_output_id`
+- `vllm_omni/diffusion/worker/diffusion_worker.py`: `WorkerProc._return_result()`, `_async_output_loop()`, `_generate_async_output_id()`
+- `vllm_omni/diffusion/executor/multiproc_executor.py`: `ResultPumpThread`, `collective_rpc()` two-path dispatch, `wait_output_ready()`, `_batch_split_map`
+- `vllm_omni/diffusion/diffusion_engine.py`: `step()` / `step_streaming()` / `add_req_and_wait_for_response()` async output waiting
+- `vllm_omni/diffusion/sched/request_scheduler.py`: `update_from_output()` — `async_output_id` → `FINISHED_COMPLETED`
+- `vllm_omni/diffusion/ipc.py`: `pack_diffusion_output_shm()` side-stream D2H path
+- `vllm_omni/platforms/interface.py`: `OmniPlatform.record_device_event()` base class
+- `vllm_omni/platforms/cuda/platform.py`: `CudaOmniPlatform.record_device_event()`
+- `vllm_omni/platforms/npu/platform.py`: `NPUOmniPlatform.record_device_event()`
+- `vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py`: `get_hunyuan_image3_post_process_func()`, postprocess removal
+- `vllm_omni/diffusion/registry.py`: `_DIFFUSION_POST_PROCESS_FUNCS` postprocess registry

docs/design/feature/tensor_parallel.md

@@ -264,7 +264,7 @@ Complete examples in the codebase:
 | **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Standard TP | Full implementation with validation |
 | **FLUX** | `vllm_omni/diffusion/models/flux/flux_transformer.py` | Dual-stream | Image + text streams |
 | **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Standard TP | With RoPE |
-| **TP Tests** | `tests/e2e/offline_inference/test_zimage_parallelism.py` | E2E testing | TP correctness and performance |
+| **TP Tests** | `tests/diffusion/distributed/test_tensor_parallel.py` | E2E testing | TP correctness and performance |
 | **Constraint Tests** | `tests/diffusion/models/z_image/test_zimage_tp_constraints.py` | Unit testing | Validation logic |
 
 ---

docs/design/metrics.md

@@ -148,7 +148,13 @@ Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, retur
 
 `OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead.
 
-The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second.
+The orchestrator side does not add its own throttle on top. Scheduler gauges are
+recorded when `raw_outputs.scheduler_stats is not None` (i.e. this replica's
+scheduler passed its own 1 Hz gate), while output-batch request/token metrics
+are recorded from `IterationStats` whenever a non-empty output batch is
+processed. A previous global `_last_stats_ts` on the orchestrator was removed
+because it starved every replica other than the first to emit within each
+second.
 
 ## Metric Definitions
 

docs/getting_started/installation/npu/npu.inc.md

@@ -9,8 +9,8 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built
 
 ```bash
 # vLLM-Ascend image
-# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.25.0
-# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.25.0-a3
+# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.26.0
+# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.26.0-a3
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -28,11 +28,11 @@ docker run --rm \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
     -v /root/.cache:/root/.cache \
     -p 8000:8000 \
-    -it quay.io/atlas-ci/vllm-ascend:0.25.0-a3 bash
+    -it quay.io/atlas-ci/vllm-ascend:v0.26.0-a3 bash
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.25.0 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.26.0 https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
@@ -84,8 +84,8 @@ Supported images as following.
 Here's an example deployment command that has been verified on 4 x NPUs:
 
 ```bash
-# Atlas A2: quay.io/ascend/vllm-omni:v0.25.0
-# Atlas A3: quay.io/ascend/vllm-omni:v0.25.0-a3
+# Atlas A2: quay.io/ascend/vllm-omni:v0.26.0
+# Atlas A3: quay.io/ascend/vllm-omni:v0.26.0-a3
 docker run --rm \
     --name vllm-omni-a3 \
     --shm-size=4g \
@@ -110,7 +110,7 @@ docker run --rm \
     -v /etc/ascend_install.info:/etc/ascend_install.info \
     -v ~/.cache:/root/.cache \
     -p 8091:8091 \
-    -it quay.io/ascend/vllm-omni:v0.25.0-a3 bash
+    -it quay.io/ascend/vllm-omni:v0.26.0-a3 bash
 ```
 
 !!! tip
@@ -122,9 +122,9 @@ Or build IMAGE from **source code**:
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 # A2
-docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu .
+# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu .
 # A3
-# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 .
+docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 .
 ```
 
 # --8<-- [end:pre-built-images]

docs/serving/speech_api.md

@@ -276,20 +276,42 @@ Client -> Server:
 
 | Message | Description |
 |---------|-------------|
-| `{"type": "session.config", ...}` | Session configuration (sent once, first message) |
+| `{"type": "session.config", ...}` | Session configuration (first message; may be resent between utterances to change it) |
 | `{"type": "input.text", "text": "..."}` | Text chunk |
-| `{"type": "input.done"}` | End of input, flushes remaining buffer |
+| `{"type": "input.done"}` | End of utterance: flushes the buffer and keeps the connection open |
+| `{"type": "session.close"}` | End of connection |
 
 Server -> Client:
 
 | Message | Description |
 |---------|-------------|
-| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
+| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
 | Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
-| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
-| `{"type": "session.done", "total_sentences": N}` | Session complete |
+| `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
+| `{"type": "session.done", "utterance_index": 0, "total_sentences": N}` | Flushed utterance complete |
 | `{"type": "error", "message": "..."}` | Non-fatal error |
 
+### Flushing vs. Closing
+
+`input.done` is a flush, not a disconnect. The server synthesizes the buffered
+text, emits `session.done`, and then waits on the same connection for the next
+utterance, so a client that speaks repeatedly (for example one driven by an
+upstream LLM) pays the WebSocket handshake once instead of once per utterance.
+
+* The session config is sticky. Send `input.text` again straight after
+  `session.done` to reuse it, or send another `session.config` first to change
+  voice, format, or reference audio. A `session.config` sent while text is
+  still buffered is rejected so no pending input is silently dropped.
+* An utterance is the flush unit, not a linguistic one: it is whatever text was
+  buffered when `input.done` arrived, of any length, synthesized as one request.
+  `utterance_index` counts those flushes across the connection, so it tells you
+  which `input.done` a frame belongs to. `sentence_index` counts within one
+  flush and so pairs with `total_sentences`, which means every utterance reports
+  `sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer).
+* End the connection with `session.close`, or by closing the socket. An idle
+  connection is still closed after the server's idle timeout, which now also
+  applies to the gap between utterances.
+
 ### Session Config Parameters
 
 All REST API parameters are supported, plus:

docs/user_guide/diffusion/attention_backends.md

@@ -14,7 +14,7 @@ The full set of backends and their platform defaults is in the **Backend Options
 
 | Value | Notes |
 |---|---|
-| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
+| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention and **SAGE** quantized attention — see below. Requires `flashinfer`. |
 | `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
 | `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
 | `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
@@ -167,6 +167,48 @@ sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `T
 Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
 `TRTLLM_ATTN` raises rather than silently degrading.
 
+## TRTLLM_ATTN SAGE Quantization
+
+On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen
+FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run
+in low precision. Enable it through the typed `quant` block on the attention spec:
+
+| Key | Valid values | Meaning |
+|---|---|---|
+| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). |
+| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). |
+| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). |
+
+V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so
+neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels.
+
+The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the
+fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`),
+`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`),
+`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises.
+
+```bash
+vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
+      "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}'
+```
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
+
+AttentionConfig(
+    default=AttentionSpec(
+        backend="TRTLLM_ATTN",
+        quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16),
+    ),
+)
+```
+
+
+**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
+availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
+(B300); `int8` QK kernels are compiled for **SM100 only**.
+
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
 
 Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
@@ -180,22 +222,6 @@ Same prompt and seed across runs. `Total generation time` from `text_to_video.py
 
 Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
 
-## Known Limitations
-
-### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
-
-LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
-
-**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
-    --model Lightricks/LTX-2 ...
-```
-
-### FA4 not yet integrated
-
-FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
 
 ## Choosing a Backend Manually
 
@@ -353,64 +379,3 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_t
     --tensor-parallel-size 2 \
     --output outputs/hv15_sage3.mp4
 ```
-
-### Mixed backends across roles
-
-Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
-    --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-    --prompt "A dog running across a field of golden wheat." \
-    --diffusion-attention-config.per_role.self.backend FLASH_ATTN \
-    --diffusion-attention-config.per_role.cross.backend TORCH_SDPA \
-    --tensor-parallel-size 2 \
-    --output outputs/wan22_mixed.mp4
-```
-
-### Compare against FlashAttention
-
-Unset the backend override, or explicitly use `FLASH_ATTN`:
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
-    --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-    --prompt "A dog running across a field of golden wheat." \
-    --height 704 --width 1280 --num-frames 49 \
-    --num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
-    --tensor-parallel-size 2 \
-    --output outputs/wan22_fa3.mp4
-```
-
-## Validation Guidance
-
-Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
-
-Always compare:
-
-- End-to-end runtime
-- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
-- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
-
-At minimum, keep the same:
-
-- model
-- prompt
-- seed
-- resolution
-- frame count / step count
-- parallel config (TP / CFG-parallel / Ulysses degrees)
-
-## Reproducing the Benchmark Table
-
-The end-to-end numbers above were collected by running `text_to_video.py` /
-`text_to_image.py` with the same prompt and seed while varying
-`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
-backends without loading a model:
-
-```bash
-python benchmarks/diffusion/bench_attention_backends.py --preset hv15
-```
-
-It runs all three BF16 backends on representative DiT attention shapes and
-prints a ranking table at the end.

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -179,7 +179,134 @@ class Flux2Transformer2DModel(nn.Module):
 - Support single GPU only for now
 
 
-### Implementation Notes
+## Distributed Layerwise Offloading
+
+### How It Works
+
+Distributed layerwise offloading extends single-GPU layerwise offloading to
+multi-device deployments.  Each DP rank stores only **1/dp_size** of the model
+weights in host memory; full layer weights are reconstructed at runtime via
+**AllGather** on a dedicated communication stream, overlapped with computation
+via a fixed double-buffer scheme.
+
+**Key features:**
+- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
+  reconstructed per-layer via `all_gather_into_tensor`
+- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any
+  time, regardless of model size
+- **DP multi-concurrency**: N concurrent requests processed in parallel
+  (AllGather only gathers weight shards, which are request-independent)
+- **mmap weight loading**: weights loaded as mmap views pointing to shared OS
+  page cache, eliminating O(dp_size × model_size) RSS during model creation
+- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU
+  (CANN/HCCL) via vLLM-Omni's platform abstraction
+
+**Execution Flow (per device):**
+
+```
+Compute Stream:  [Layer N]          [Layer N+1]          [Layer N+2]
+H2D Stream:      [H2D Shard N+1]   [H2D Shard N+2]
+AllGather:       [AG N+1]          [AG N+2]
+Slot Usage:      Slot0: Layer N    Slot1: Layer N+1
+```
+
+### Usage
+
+**CLI:**
+```bash
+# 4× GPU/NPU with AllGather (recommended)
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4
+
+# Without AllGather (each rank loads full weights, no sharding)
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4 \
+  --dlo-no-use-allgather
+
+# With SP instead of DP (long sequences)
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --usp 4
+```
+
+**Python API:**
+```python
+from vllm_omni import Omni
+
+m = Omni(
+    model="/path/to/model",
+    enable_distributed_layerwise_offload=True,
+    dlo_use_allgather=True,  # default
+)
+```
+
+### CLI Flags
+
+| Flag | Description | Default |
+|------|-------------|---------|
+| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
+| `--data-parallel-size N` | Number of DP ranks for weight sharding | `1` |
+| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
+| `--dlo-no-use-allgather` | Full weights per rank (no sharding, no AllGather) | `false` |
+
+### How model weights are loaded (mmap path)
+
+When DLO + AllGather is active, the offloader:
+
+1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`)
+   from the normally-created transformer
+2. **Converts to meta device** via `to_empty(device="meta")`, releasing random
+   initialization weights
+3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`,
+   which return views into the OS page cache (shared across all ranks, 0 RSS)
+4. **Calls `post_load_weights()`** to apply model-specific dtype conversions
+   (e.g. Cosmos3's `time_embedder` → FP32)
+5. **Restores non-persistent buffers** from saved copies
+
+This approach requires **zero model-specific code changes** — no pipeline or
+transformer modifications are needed.
+
+### OffloadPlan (declarative topology metadata)
+
+Models can optionally declare an `OffloadPlan` class variable to provide
+topology metadata (block attribute names, submodules to offload) without
+any offload-specific logic:
+
+```python
+from vllm_omni.diffusion.offloader import OffloadPlan
+
+class MyPipeline(nn.Module):
+    _dit_modules = ["transformer"]
+    _offload_plan = OffloadPlan(
+        block_attrs={"transformer": ("blocks",)},
+        offload_submodules={"context_encoder": "layers"},
+    )
+```
+
+When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs`
+and heuristic attribute search.  This is backward-compatible — existing models
+work without any changes.
+
+### DP Multi-concurrency
+
+When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches
+up to `dp_size` requests per denoise step.  Each DP rank processes a different
+request while AllGather synchronizes only weight shards (request-independent).
+
+**Requirements for concurrent requests:**
+- `num_inference_steps` must be specified explicitly (None is not allowed)
+- All concurrent requests must have the same `num_inference_steps` value
+- These constraints exist because AllGather is a collective that requires
+  every rank to participate at each denoise step
+
+### Limitations
+- Online quantization (FP8) is incompatible with mmap loading — falls back to
+  regular `load_weights()` automatically
+- Tensor Parallel is not supported (DLO uses DP-based sharding)
+- HSDP + AllGather is rejected (would double-shard weights)
+- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
 
 **Module Discovery**
 
@@ -208,7 +335,10 @@ OffloadBackend (base class)
 ├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
 │                              (delegates to a pipeline's enable_omni_model_cpu_offload
 │                               for split models like Cosmos3)
-└── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
+├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
+│                          (single-GPU, full weights on host)
+└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
+                                         (multi-GPU, 1/dp_size sharded weights + AllGather)
 ```
 
 Factory function `get_offload_backend()` selects the appropriate backend based on
@@ -221,20 +351,22 @@ reasoner/generator components inside the model forward pass.
 
 ## Supported Models
 
-| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attrs (Layerwise specific) |
-|--------------|----------------|-----------|---------------------|-------------------|-----------------------------------|
-| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` |
-| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` |
-| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` |
-| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` |
-| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` |
-| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | `"down_blocks"`, `"up_blocks"` |
-| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
-| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
-| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
-| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | `"layers"` |
-| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` |
+| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) |
+|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------|
+| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
+| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
+| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` |
+| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` |
+| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` |
+| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` |
+| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` |
+| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
+| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
+| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` |
+| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` |
+| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` |
 
 **Notes:**
 - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
 - Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
+- Distributed Layerwise Offloading works with any model that supports Layerwise Offloading — no additional model changes required.  See [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples.

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -306,7 +306,7 @@ The gradio script supports the following arguments:
     ``````
 ??? abstract "qwen3_omni_moe_thinking.yaml"
     ``````yaml
-    --8<-- "examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml"
+    --8<-- "vllm_omni/deploy/qwen3_omni_moe_thinking.yaml"
     ``````
 ??? abstract "run_curl_multimodal_generation.sh"
     ``````sh

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -422,12 +422,19 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
 
 ### Streaming WebSocket
-The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence:
+The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
 ```bash
 python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
 python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
 ```
 
+`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
+```bash
+python qwen3_tts/streaming_speech_client.py \
+    --text "First utterance." \
+    --text "Second utterance, same connection."
+```
+
 ### Gradio demos
 ```bash
 ./qwen3_tts/run_gradio_demo.sh                              # CustomVoice (default)

examples/model_display_names.yml

@@ -8,5 +8,4 @@ mimo_audio: MiMo-Audio
 ming_flash_omni: Ming-flash-omni 2.0
 qwen2_5_omni: Qwen2.5-Omni
 qwen3_omni: Qwen3-Omni
-sensenova_u1: SenseNova-U1
 step_audio2: Step-Audio2

examples/offline_inference/image_to_video/image_to_video.py

@@ -334,6 +334,7 @@ def main():
     is_ltx23 = "ltx23" in model_class_name_lower or "ltx-2.3" in model_name
     is_ltx2 = is_ltx2_distilled or is_ltx23 or "ltx2" in model_class_name_lower or "ltx-2" in model_name
     is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
+    is_cosmos_edge = is_cosmos and ("edge" in model_name or "edge" in model_class_name_lower)
 
     image = PIL.Image.open(args.image).convert("RGB") if args.image else None
     last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None
@@ -347,7 +348,19 @@ def main():
 
     # Per-model generation defaults, applied only when the matching flag is omitted.
     # Cosmos3 would otherwise silently inherit the Wan2.2 defaults (wrong size/steps/shift).
-    if is_cosmos:
+    if is_cosmos_edge:
+        # Cosmos3-Edge native defaults: 480x832, gs 5.0, flow_shift 3.0 — NOT the Nano/Super
+        # 720p / gs 6.0 / flow_shift 10.0 below, which produce degenerate output on Edge.
+        d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
+            24,
+            5.0,
+            189,
+            35,
+            3.0,
+            480 * 832,
+            16,
+        )
+    elif is_cosmos:
         d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
             24,
             6.0,

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -182,10 +182,13 @@ def run_e2e():
                 print(f"Multimodal output keys: {mm.keys()}")
                 if "audio" in mm:
                     audio_out = mm["audio"]
+                    # The model emits its own rate (24 kHz for CosyVoice3); take
+                    # it from the payload so the file is not mislabeled.
+                    sample_rate = int(mm["sr"]) if "sr" in mm.keys() else 24000
                     print(f"Generated Audio Shape: {audio_out.shape}")
                     out_path = f"output_{i}.wav"
-                    sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
-                    print(f"Saved audio to {out_path}")
+                    sf.write(out_path, audio_out.cpu().numpy().squeeze(), sample_rate)
+                    print(f"Saved audio to {out_path} (sr={sample_rate})")
             else:
                 print("No multimodal output found.")
         except Exception as e:

examples/offline_inference/text_to_video/text_to_video.py

@@ -56,6 +56,19 @@ _MODEL_PRESETS = {
         "flow_shift": 10.0,
         "output": "cosmos3_t2v_output.mp4",
     },
+    # Cosmos3-Edge is a compact checkpoint with its own native resolution / guidance /
+    # flow_shift. It must NOT inherit the Nano/Super "cosmos" numbers above (720p / gs 6.0 /
+    # flow_shift 10.0), which produce degenerate output on Edge.
+    "cosmos3_edge": {
+        "height": 480,
+        "width": 832,
+        "num_frames": 189,
+        "num_inference_steps": 35,
+        "guidance_scale": 5.0,
+        "fps": 24,
+        "flow_shift": 3.0,
+        "output": "cosmos3_edge_t2v_output.mp4",
+    },
     "helios": {
         "height": 384,
         "width": 640,
@@ -103,6 +116,10 @@ def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
         return _MODEL_PRESETS["ltx2"]
     if "vace" in model_lower:
         return _MODEL_PRESETS["vace"]
+    # Edge must be matched before the generic cosmos branch (its "cosmos" substring would
+    # otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset).
+    if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower):
+        return _MODEL_PRESETS["cosmos3_edge"]
     if "cosmos" in model_lower:
         return _MODEL_PRESETS["cosmos"]
     if "hunyuan" in model_lower:

examples/online_serving/minicpmo/README.md

@@ -23,17 +23,18 @@ including `librosa`.
 
 The deploy config auto-loads via `--omni`.
 The default `vllm_omni/deploy/minicpmo_4_5.yaml` keeps all three stages on
-logical device 0 with memory budgets of 55%, 22%, and 22%. The profile admits
-at most four sequences per stage and bounds startup video profiling to 32
-frames per video. For throughput,`minicpmo_4_5_batching.yaml` gives the Thinker
+logical device 0 with memory budgets of 55%, 15%, and 15%, leaving headroom
+for runtime kernels such as the HiFi-GAN vocoder's cuDNN workspace. The
+profile admits at most four sequences per stage and bounds startup video
+profiling to 32 frames per video. For throughput,
+`minicpmo_4_5_2gpu.yaml` gives the Thinker
 GPU 0 (90%) and colocates the Talker (55%) and Code2Wav (35%) on GPU 1. That
 profile admits at most four concurrent sequences per stage.
 
 | deploy config | GPUs | Notes |
 |---|---|---|
 | `minicpmo_4_5.yaml` (default) | 1 | Memory-constrained compatibility layout. |
-| `minicpmo_4_5_batching.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. |
-| `minicpmo_4_5_2gpu.yaml` | 2 | Backward-compatible alias for `minicpmo_4_5_batching.yaml`. |
+| `minicpmo_4_5_2gpu.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. |
 | `minicpmo_4_5_3gpu.yaml` | 3 | One GPU per stage. |
 | `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
 | `minicpmo_4_5_duplex.yaml` | 1 | Experimental native full-duplex overlay. |
@@ -89,17 +90,18 @@ python openai_chat_completion_client_for_multimodal_generation.py \
     --prompt "Briefly introduce yourself."
 ```
 
-Streaming text + audio:
+Streaming text + audio (use `--stream`):
 
 ```bash
-python streaming_chat_completion.py \
-    --base-url http://localhost:8099/v1 \
-    --output minicpmo_stream.wav
+python openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type text \
+    --prompt "Briefly introduce yourself." \
+    --port 8099 \
+    --stream
 ```
 
-The example prints text deltas immediately and joins the independently encoded
-24 kHz WAV audio deltas into one valid WAV file. Add `--text-only` to skip
-Talker and Code2Wav.
+The client prints text deltas as they arrive and saves streamed audio chunks
+to WAV files.
 
 Shared helpers also work if you pass MiniCPM defaults yourself:
 

examples/online_serving/minicpmo/realtime_web/app/static/playback_worklet.js

@@ -6,6 +6,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
     this.playedFrames = 0;
     this.underrunFrames = 0;
     this.drain = null;
+    this.terminalFadeFrames = 0;
     this.started = false;
     this.activeResponseId = null;
     this.initialBufferFrames = Math.round(sampleRate * 0.2);
@@ -31,6 +32,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
       }
     } else if (message.type === 'drain') {
       this.drain = { responseId: message.responseId || null };
+      this.terminalFadeFrames = Math.min(this.fadeFrames, this.bufferedFrames());
       if (!this.started && this.bufferedFrames() > 0) {
         this.bufferWaitFrames = 0;
         this.startPlayback();
@@ -42,6 +44,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
       this.playedFrames = 0;
       this.underrunFrames = 0;
       this.drain = null;
+      this.terminalFadeFrames = 0;
       this.started = false;
       this.activeResponseId = null;
       this.bufferWaitFrames = this.initialBufferFrames;
@@ -76,6 +79,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
     this.playedFrames = 0;
     this.underrunFrames = 0;
     this.drain = null;
+    this.terminalFadeFrames = 0;
     this.started = false;
     this.activeResponseId = null;
     this.bufferWaitFrames = this.initialBufferFrames;
@@ -116,6 +120,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
     while (target < output.length && this.queue.length > 0) {
       const pcm = this.queue[0];
       const count = Math.min(output.length - target, pcm.length - this.offset);
+      const remainingBeforeChunk = this.bufferedFrames();
       for (let index = 0; index < count; index += 1) {
         let sample = pcm[this.offset + index] / 32768;
         if (this.fadeInFrames > 0) {
@@ -123,6 +128,16 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
           sample *= elapsed / this.fadeFrames;
           this.fadeInFrames -= 1;
         }
+        const remainingFrames = remainingBeforeChunk - index;
+        if (
+          this.drain
+          && this.terminalFadeFrames > 0
+          && remainingFrames <= this.terminalFadeFrames
+        ) {
+          sample *= this.terminalFadeFrames === 1
+            ? 0
+            : (remainingFrames - 1) / (this.terminalFadeFrames - 1);
+        }
         output[target + index] = sample;
       }
       target += count;

examples/online_serving/minicpmo/streaming_chat_completion.py

@@ -1,71 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Stream MiniCPM-o 4.5 text and audio from /v1/chat/completions."""
-
-from __future__ import annotations
-
-import argparse
-import base64
-import io
-import wave
-from pathlib import Path
-
-from openai import OpenAI
-
-
-def append_wav_chunk(encoded: str, pcm_parts: list[bytes]) -> tuple[int, int, int]:
-    """Decode one base64 WAV delta and append its PCM frames."""
-    with wave.open(io.BytesIO(base64.b64decode(encoded)), "rb") as wav:
-        channels = wav.getnchannels()
-        sample_width = wav.getsampwidth()
-        sample_rate = wav.getframerate()
-        pcm_parts.append(wav.readframes(wav.getnframes()))
-    return channels, sample_width, sample_rate
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--base-url", default="http://localhost:8099/v1")
-    parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
-    parser.add_argument("--prompt", default="Say hello, then introduce vLLM in one sentence.")
-    parser.add_argument("--output", type=Path, default=Path("minicpmo_stream.wav"))
-    parser.add_argument("--text-only", action="store_true")
-    args = parser.parse_args()
-
-    modalities = ["text"] if args.text_only else ["text", "audio"]
-    stream = OpenAI(base_url=args.base_url, api_key="EMPTY").chat.completions.create(
-        model=args.model,
-        messages=[{"role": "user", "content": args.prompt}],
-        modalities=modalities,
-        stream=True,
-        extra_body={"chat_template_kwargs": {"use_tts_template": not args.text_only}},
-    )
-
-    pcm_parts: list[bytes] = []
-    audio_format: tuple[int, int, int] | None = None
-    for chunk in stream:
-        for choice in chunk.choices:
-            content = getattr(choice.delta, "content", None)
-            if not content:
-                continue
-            if getattr(chunk, "modality", None) == "audio":
-                current_format = append_wav_chunk(content, pcm_parts)
-                if audio_format is not None and current_format != audio_format:
-                    raise ValueError(f"Audio format changed: {audio_format} -> {current_format}")
-                audio_format = current_format
-            elif getattr(chunk, "modality", None) == "text":
-                print(content, end="", flush=True)
-
-    print()
-    if pcm_parts and audio_format is not None:
-        channels, sample_width, sample_rate = audio_format
-        with wave.open(str(args.output), "wb") as wav:
-            wav.setnchannels(channels)
-            wav.setsampwidth(sample_width)
-            wav.setframerate(sample_rate)
-            wav.writeframes(b"".join(pcm_parts))
-        print(f"Audio saved to {args.output} ({sample_rate} Hz)")
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml

@@ -1,36 +0,0 @@
-# Stage config for running Qwen3-Omni-MoE-Thinking (text-only output)
-# This config is for models like Qwen3-Omni-30B-A3B-Thinking that only have the
-# thinker component and do not support audio output.
-#
-# Single stage: Thinker (multimodal understanding + text generation)
-
-# The following config has been verified on 2x H100-80G GPUs.
-stage_args:
-  - stage_id: 0
-    runtime:
-      devices: "0,1"
-    engine_args:
-      model_stage: thinker
-      max_num_seqs: 1
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: text
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: true
-      hf_config_name: thinker_config
-      tensor_parallel_size: 2
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05

examples/online_serving/qwen3_omni/run_gradio_demo.sh

@@ -13,7 +13,7 @@ set -e
 MODEL="Qwen/Qwen3-Omni-30B-A3B-Instruct"
 SERVER_PORT=8091
 GRADIO_PORT=7861
-STAGE_CONFIGS_PATH=""
+DEPLOY_CONFIG=""
 SERVER_HOST="0.0.0.0"
 GRADIO_IP="127.0.0.1"
 GRADIO_SHARE=false
@@ -33,8 +33,8 @@ while [[ $# -gt 0 ]]; do
             GRADIO_PORT="$2"
             shift 2
             ;;
-        --stage-configs-path)
-            STAGE_CONFIGS_PATH="$2"
+        --deploy-config)
+            DEPLOY_CONFIG="$2"
             shift 2
             ;;
         --server-host)
@@ -56,7 +56,7 @@ while [[ $# -gt 0 ]]; do
             echo "  --model MODEL                 Model name/path (default: Qwen/Qwen3-Omni-30B-A3B-Instruct)"
             echo "  --server-port PORT            Port for vLLM server (default: 8091)"
             echo "  --gradio-port PORT            Port for Gradio demo (default: 7861)"
-            echo "  --stage-configs-path PATH     Path to custom stage configs YAML file (optional)"
+            echo "  --deploy-config PATH          Path to custom deploy config YAML file (optional)"
             echo "  --server-host HOST            Host for vLLM server (default: 0.0.0.0)"
             echo "  --gradio-ip IP                IP for Gradio demo (default: 127.0.0.1)"
             echo "  --share                       Share Gradio demo publicly"
@@ -87,8 +87,8 @@ echo "=========================================="
 
 # Build vLLM server command
 SERVER_CMD=("vllm" "serve" "$MODEL" "--omni" "--port" "$SERVER_PORT" "--host" "$SERVER_HOST")
-if [ -n "$STAGE_CONFIGS_PATH" ]; then
-    SERVER_CMD+=("--stage-configs-path" "$STAGE_CONFIGS_PATH")
+if [ -n "$DEPLOY_CONFIG" ]; then
+    SERVER_CMD+=("--deploy-config" "$DEPLOY_CONFIG")
 fi
 
 # Function to cleanup on exit

examples/online_serving/text_to_speech/README.md

@@ -537,12 +537,19 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
 
 ### Streaming WebSocket
-The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence:
+The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
 ```bash
 python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
 python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
 ```
 
+`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
+```bash
+python qwen3_tts/streaming_speech_client.py \
+    --text "First utterance." \
+    --text "Second utterance, same connection."
+```
+
 To receive word-level timestamps, launch the server with a forced aligner:
 ```bash
 vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \

examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py

@@ -7,6 +7,12 @@ Usage:
     # Send full text at once
     python streaming_speech_client.py --text "Hello world. How are you? I am fine."
 
+    # Several utterances over one connection: input.done flushes each of them
+    # and the connection stays open, so only the first pays the handshake
+    python streaming_speech_client.py \
+        --text "First utterance." \
+        --text "Second utterance, same connection."
+
     # Pick a built-in speaker for CustomVoice models
     python streaming_speech_client.py \
         --text "打开电子枪，关闭扫描，切换到低倍模式。" \
@@ -64,18 +70,27 @@ except ImportError:
     raise SystemExit(1)
 
 
-def save_audio_file(output_dir: str, sentence_index: int, response_format: str, chunks: list[bytes]) -> str:
+def frame_basename(msg: dict) -> str:
+    """Name a file after the utterance and sentence a frame belongs to.
+
+    Every utterance is sentence 0, so the utterance index is what keeps the
+    files of one connection from overwriting each other.
+    """
+    return f"utterance_{msg['utterance_index']:03d}_sentence_{msg['sentence_index']:03d}"
+
+
+def save_audio_file(output_dir: str, basename: str, response_format: str, chunks: list[bytes]) -> str:
     audio_bytes = b"".join(chunks)
     if response_format == "pcm":
-        raw_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.pcm")
+        raw_path = os.path.join(output_dir, f"{basename}.pcm")
         with open(raw_path, "wb") as f:
             f.write(audio_bytes)
 
-        wav_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.wav")
+        wav_path = os.path.join(output_dir, f"{basename}.wav")
         save_pcm_wav(wav_path, audio_bytes)
         return wav_path
 
-    filename = os.path.join(output_dir, f"sentence_{sentence_index:03d}.{response_format}")
+    filename = os.path.join(output_dir, f"{basename}.{response_format}")
     with open(filename, "wb") as f:
         f.write(audio_bytes)
     return filename
@@ -95,183 +110,207 @@ def timestamp_words_text(timestamps: list[dict] | None) -> str:
     return " ".join(str(item.get("word", "")) for item in timestamps if item.get("word"))
 
 
+async def send_utterance(ws, text: str, simulate_stt: bool, stt_delay: float) -> None:
+    """Send one utterance's text, then input.done to flush it."""
+    if simulate_stt:
+        words = text.split(" ")
+        for i, word in enumerate(words):
+            chunk = word + (" " if i < len(words) - 1 else "")
+            await ws.send(
+                json.dumps(
+                    {
+                        "type": "input.text",
+                        "text": chunk,
+                    }
+                )
+            )
+            print(f"  Sent: {chunk!r}")
+            await asyncio.sleep(stt_delay)
+    else:
+        await ws.send(
+            json.dumps(
+                {
+                    "type": "input.text",
+                    "text": text,
+                }
+            )
+        )
+        print(f"Sent full text: {text!r}")
+
+    await ws.send(json.dumps({"type": "input.done"}))
+    print("Sent input.done")
+
+
+async def receive_utterance(
+    ws,
+    output_dir: str,
+    response_format: str,
+    word_timestamps: bool,
+    save_chunks: bool,
+) -> None:
+    """Consume frames until session.done marks the flushed utterance complete."""
+    current_label = "utterance 0"
+    current_sentence_text = ""
+    current_chunks: list[bytes] = []
+    current_timestamps: list[dict] = []
+    # Distinguish `null` (aligner failed) from `[]` (silence / no tokens):
+    # the sentence is a failure only if no frame ever carried timestamps.
+    alignment_frame_seen = False
+
+    while True:
+        message = await ws.recv()
+
+        if isinstance(message, bytes):
+            current_chunks.append(message)
+            print(f"  Received audio chunk for {current_label}: {len(message)} bytes")
+        else:
+            # JSON frame
+            msg = json.loads(message)
+            msg_type = msg.get("type")
+
+            if msg_type == "audio.start":
+                current_label = f"utterance {msg['utterance_index']}"
+                current_sentence_text = msg.get("sentence_text", "")
+                current_chunks = []
+                current_timestamps = []
+                alignment_frame_seen = False
+                print(f"  [{current_label}] Generating: {current_sentence_text!r}")
+            elif msg_type == "audio.chunk":
+                audio = base64.b64decode(msg["audio_b64"])
+                current_chunks.append(audio)
+                chunk_timestamps = msg.get("timestamps")
+                # Sentence-level alignment sends a trailing timestamp-only
+                # frame with empty audio; don't write an empty chunk file.
+                if save_chunks and audio:
+                    chunk_base = os.path.join(
+                        output_dir,
+                        f"{frame_basename(msg)}_chunk_{msg['chunk_id']:03d}",
+                    )
+                    chunk_pcm = f"{chunk_base}.pcm"
+                    chunk_wav = f"{chunk_base}.wav"
+                    chunk_json = f"{chunk_base}_timestamps.json"
+                    with open(chunk_pcm, "wb") as f:
+                        f.write(audio)
+                    save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000)))
+                    with open(chunk_json, "w", encoding="utf-8") as f:
+                        json.dump(
+                            {
+                                "utterance_index": msg["utterance_index"],
+                                "sentence_index": msg["sentence_index"],
+                                "chunk_id": msg["chunk_id"],
+                                "chunk_start_ms": msg.get("chunk_start_ms"),
+                                "chunk_end_ms": msg.get("chunk_end_ms"),
+                                "sentence_text": current_sentence_text,
+                                "covered_text": timestamp_words_text(chunk_timestamps),
+                                "timestamps": chunk_timestamps,
+                            },
+                            f,
+                            ensure_ascii=False,
+                            indent=2,
+                        )
+                if chunk_timestamps is None:
+                    print(f"  [{current_label}] chunk {msg['chunk_id']}: {len(audio)} bytes, timestamps unavailable")
+                else:
+                    alignment_frame_seen = True
+                    current_timestamps.extend(chunk_timestamps)
+                    print(
+                        f"  [{current_label}] chunk {msg['chunk_id']}: "
+                        f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)"
+                    )
+                if save_chunks and audio:
+                    print(f"    saved chunk -> {chunk_wav}")
+            elif msg_type == "audio.done":
+                filename = save_audio_file(
+                    output_dir,
+                    frame_basename(msg),
+                    response_format,
+                    current_chunks,
+                )
+                if word_timestamps:
+                    # `null` on failure, else the accumulated list (`[]` for silence).
+                    sentence_timestamps = current_timestamps if alignment_frame_seen else None
+                    ts_filename = os.path.join(
+                        output_dir,
+                        f"{frame_basename(msg)}_timestamps.json",
+                    )
+                    with open(ts_filename, "w", encoding="utf-8") as f:
+                        json.dump(
+                            {
+                                "utterance_index": msg["utterance_index"],
+                                "sentence_index": msg["sentence_index"],
+                                "sentence_text": current_sentence_text,
+                                "covered_text": timestamp_words_text(sentence_timestamps),
+                                "timestamps": sentence_timestamps,
+                            },
+                            f,
+                            ensure_ascii=False,
+                            indent=2,
+                        )
+                print(
+                    f"  [{current_label}] Done"
+                    f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
+                    f" error={msg.get('error', False)}"
+                    f" -> {filename}"
+                )
+                if word_timestamps:
+                    print(f"  [{current_label}] Timestamps -> {ts_filename}")
+                current_chunks = []
+                current_timestamps = []
+                alignment_frame_seen = False
+            elif msg_type == "session.done":
+                print(f"\nUtterance {msg['utterance_index']} complete: {msg['total_sentences']} sentence(s) generated")
+                return
+            elif msg_type == "error":
+                print(f"  ERROR: {msg['message']}")
+            else:
+                print(f"  Unknown message: {msg}")
+
+
 async def stream_tts(
     url: str,
-    text: str,
+    texts: list[str],
     config: dict,
     output_dir: str,
     simulate_stt: bool = False,
     stt_delay: float = 0.1,
     save_chunks: bool = False,
 ) -> None:
-    """Connect to the streaming TTS endpoint and process audio responses."""
+    """Connect to the streaming TTS endpoint and process audio responses.
+
+    Every text is synthesized over the same connection: input.done flushes an
+    utterance without closing, so only the first one pays the handshake.
+    """
     os.makedirs(output_dir, exist_ok=True)
+    response_format = config.get("response_format", "wav")
+    word_timestamps = bool(config.get("word_timestamps", False))
 
     async with websockets.connect(url) as ws:
-        # 1. Send session config
+        # 1. Send session config; it stays in effect for every utterance below.
         config_msg = {"type": "session.config", **config}
         await ws.send(json.dumps(config_msg))
         print(f"Sent session config: {config}")
 
-        # 2. Send text (either all at once or word-by-word)
-        async def send_text():
-            if simulate_stt:
-                words = text.split(" ")
-                for i, word in enumerate(words):
-                    chunk = word + (" " if i < len(words) - 1 else "")
-                    await ws.send(
-                        json.dumps(
-                            {
-                                "type": "input.text",
-                                "text": chunk,
-                            }
-                        )
-                    )
-                    print(f"  Sent: {chunk!r}")
-                    await asyncio.sleep(stt_delay)
-            else:
-                await ws.send(
-                    json.dumps(
-                        {
-                            "type": "input.text",
-                            "text": text,
-                        }
-                    )
-                )
-                print(f"Sent full text: {text!r}")
-
-            # 3. Signal end of input
-            await ws.send(json.dumps({"type": "input.done"}))
-            print("Sent input.done")
-
-        # Run sender and receiver concurrently
-        sender_task = asyncio.create_task(send_text())
-
-        response_format = config.get("response_format", "wav")
-        word_timestamps = bool(config.get("word_timestamps", False))
-        current_sentence_index = 0
-        current_sentence_text = ""
-        current_chunks: list[bytes] = []
-        current_timestamps: list[dict] = []
-        # Distinguish `null` (aligner failed) from `[]` (silence / no tokens):
-        # the sentence is a failure only if no frame ever carried timestamps.
-        alignment_frame_seen = False
-
-        try:
-            while True:
-                message = await ws.recv()
-
-                if isinstance(message, bytes):
-                    current_chunks.append(message)
-                    print(f"  Received audio chunk for sentence {current_sentence_index}: {len(message)} bytes")
-                else:
-                    # JSON frame
-                    msg = json.loads(message)
-                    msg_type = msg.get("type")
-
-                    if msg_type == "audio.start":
-                        current_sentence_index = msg["sentence_index"]
-                        current_sentence_text = msg.get("sentence_text", "")
-                        current_chunks = []
-                        current_timestamps = []
-                        alignment_frame_seen = False
-                        print(f"  [sentence {msg['sentence_index']}] Generating: {current_sentence_text!r}")
-                    elif msg_type == "audio.chunk":
-                        audio = base64.b64decode(msg["audio_b64"])
-                        current_chunks.append(audio)
-                        chunk_timestamps = msg.get("timestamps")
-                        # Sentence-level alignment sends a trailing timestamp-only
-                        # frame with empty audio; don't write an empty chunk file.
-                        if save_chunks and audio:
-                            chunk_base = os.path.join(
-                                output_dir,
-                                f"sentence_{msg['sentence_index']:03d}_chunk_{msg['chunk_id']:03d}",
-                            )
-                            chunk_pcm = f"{chunk_base}.pcm"
-                            chunk_wav = f"{chunk_base}.wav"
-                            chunk_json = f"{chunk_base}_timestamps.json"
-                            with open(chunk_pcm, "wb") as f:
-                                f.write(audio)
-                            save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000)))
-                            with open(chunk_json, "w", encoding="utf-8") as f:
-                                json.dump(
-                                    {
-                                        "sentence_index": msg["sentence_index"],
-                                        "chunk_id": msg["chunk_id"],
-                                        "chunk_start_ms": msg.get("chunk_start_ms"),
-                                        "chunk_end_ms": msg.get("chunk_end_ms"),
-                                        "sentence_text": current_sentence_text,
-                                        "covered_text": timestamp_words_text(chunk_timestamps),
-                                        "timestamps": chunk_timestamps,
-                                    },
-                                    f,
-                                    ensure_ascii=False,
-                                    indent=2,
-                                )
-                        if chunk_timestamps is None:
-                            print(
-                                f"  [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
-                                f"{len(audio)} bytes, timestamps unavailable"
-                            )
-                        else:
-                            alignment_frame_seen = True
-                            current_timestamps.extend(chunk_timestamps)
-                            print(
-                                f"  [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
-                                f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)"
-                            )
-                        if save_chunks and audio:
-                            print(f"    saved chunk -> {chunk_wav}")
-                    elif msg_type == "audio.done":
-                        filename = save_audio_file(
-                            output_dir,
-                            msg["sentence_index"],
-                            response_format,
-                            current_chunks,
-                        )
-                        if word_timestamps:
-                            # `null` on failure, else the accumulated list (`[]` for silence).
-                            sentence_timestamps = current_timestamps if alignment_frame_seen else None
-                            ts_filename = os.path.join(
-                                output_dir,
-                                f"sentence_{msg['sentence_index']:03d}_timestamps.json",
-                            )
-                            with open(ts_filename, "w", encoding="utf-8") as f:
-                                json.dump(
-                                    {
-                                        "sentence_index": msg["sentence_index"],
-                                        "sentence_text": current_sentence_text,
-                                        "covered_text": timestamp_words_text(sentence_timestamps),
-                                        "timestamps": sentence_timestamps,
-                                    },
-                                    f,
-                                    ensure_ascii=False,
-                                    indent=2,
-                                )
-                        print(
-                            f"  [sentence {msg['sentence_index']}] Done"
-                            f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
-                            f" error={msg.get('error', False)}"
-                            f" -> {filename}"
-                        )
-                        if word_timestamps:
-                            print(f"  [sentence {msg['sentence_index']}] Timestamps -> {ts_filename}")
-                        current_chunks = []
-                        current_timestamps = []
-                        alignment_frame_seen = False
-                    elif msg_type == "session.done":
-                        print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated")
-                        break
-                    elif msg_type == "error":
-                        print(f"  ERROR: {msg['message']}")
-                    else:
-                        print(f"  Unknown message: {msg}")
-        finally:
-            sender_task.cancel()
+        for text in texts:
+            # 2. Send text (either all at once or word-by-word), then flush.
+            sender_task = asyncio.create_task(send_utterance(ws, text, simulate_stt, stt_delay))
             try:
-                await sender_task
-            except asyncio.CancelledError:
-                pass  # Task cancellation is expected during shutdown
+                await receive_utterance(
+                    ws,
+                    output_dir=output_dir,
+                    response_format=response_format,
+                    word_timestamps=word_timestamps,
+                    save_chunks=save_chunks,
+                )
+            finally:
+                sender_task.cancel()
+                try:
+                    await sender_task
+                except asyncio.CancelledError:
+                    pass  # Task cancellation is expected during shutdown
+
+        # 3. Release the connection once there is nothing left to say.
+        await ws.send(json.dumps({"type": "session.close"}))
+        print("Sent session.close")
 
     print(f"\nAudio files saved to: {output_dir}/")
 
@@ -286,7 +325,8 @@ def main():
     parser.add_argument(
         "--text",
         required=True,
-        help="Text to synthesize",
+        action="append",
+        help="Text to synthesize; repeat to send several utterances over one connection",
     )
     parser.add_argument(
         "--output-dir",
@@ -385,7 +425,7 @@ def main():
     asyncio.run(
         stream_tts(
             url=args.url,
-            text=args.text,
+            texts=args.text,
             config=config,
             output_dir=args.output_dir,
             simulate_stt=args.simulate_stt,

pyproject.toml

@@ -75,6 +75,7 @@ dev = [
 ]
 
 demo = [
+    "gradio>=6.7.0",
     "opencv-python>=4.12.0.88",  # cli/stream_client frame reading
     "requests>=2.28.0",  # cli/stream_client HTTP
     "cn2an>=0.5.22",  # Chinese number normalization for TTS (e.g. "1次" -> "一次")

realtime_video_prompt_interaction_protocol.md

@@ -1,138 +0,0 @@
-# Realtime Video Prompt Interaction Protocol
-
-This document specifies the prompt-update subset of the realtime video
-interaction protocol exposed by:
-
-```text
-WebSocket /v1/realtime/video
-```
-
-A WebSocket connection contains exactly one video-generation request. Clients
-therefore do not send a request or session ID with an interaction. The server
-assigns a `request_id` in `video.start` and includes it in subsequent
-server-to-client messages.
-
-## Interaction Request
-
-After `video.start`, a client may submit a prompt update:
-
-```json
-{
-  "type": "session.interaction",
-  "interaction": {
-    "event_id": "ui-update-42",
-    "event": {
-      "prompt": "A snowy street illuminated by warm lanterns"
-    },
-    "transition_chunks": 3
-  }
-}
-```
-
-`interaction` has the following fields:
-
-- `event_id` (optional string): session-unique idempotency key. The
-  server preserves a client-provided value or generates a UUID when omitted.
-- `event.prompt` (required string): non-empty target prompt.
-  over which the prompt reaches its target.
-- `transition_chunks` (optional non-negative integer): model-generation chunks
-  over which the prompt reaches its target. If omitted, the model's default transition is used.
-
-Note: Although `["interaction"]["event"]` only contains `prompt`, still consider using `OmniTextPrompt` for this part (with total=false) so that other fields will be supported in the future. Currently, `vllm_omni/diffusion/worker/diffusion_model_runner.py` already takes care of enforcing subset format, routing to meaningful input handing routines, and raising `NotImplementedError`. Only need to adapt this design to the input data structure above.
-
-### Transition Semantics
-
-Prompt transitions are target updates and begin at the next eligible model
-generation boundary after the interaction is processed.
-
-- `transition_chunks = 0` is a hard update: the first affected generation
-  chunk uses the target prompt.
-- `transition_chunks = N`, where `N > 0`, reaches the target in the `N`th
-  affected generation chunk. For example, a three-chunk linear transition uses
-  progress values `1/3`, `2/3`, and `1`.
-
-### Idempotency and Replacement
-
-Within a session:
-
-- repeating an `event_id` with the same payload is a retry and must not
-  apply the update twice;
-- repeating an `event_id` with a different payload is rejected with
-  `event_id_conflict`;
-- a newer prompt target supersedes any older pending prompt target;
-- when a new target interrupts an active transition, its transition starts
-  from the prompt state resolved at its effective boundary.
-
-## Queue Acknowledgement
-
-The server quickly acknowledges admission to the interaction queue:
-
-```json
-{
-  "type": "session.interaction.queued",
-  "request_id": "video_stream-...",
-  "event_id": "ui-update-42"
-}
-```
-
-`queued` confirms syntax validation, idempotency handling, and queue admission.
-It does not claim that the model has scheduled or applied the update. A later
-failure is reported as a correlated error:
-
-```json
-{
-  "type": "error",
-  "request_id": "video_stream-...",
-  "event_id": "ui-update-42",
-  "code": "interaction_failed",
-  "message": "The active model does not support prompt updates"
-}
-```
-
-## Video Chunk Metadata
-
-Each emitted binary payload is preceded by a JSON metadata message. The
-metadata message and its binary payload form one ordered pair and must not be
-interleaved with another outbound message.
-
-```json
-{
-  "type": "video.chunk_metadata",
-  "request_id": "video_stream-...",
-  "kind": "media",
-  "transport_chunk_index": 7,
-  "generation_chunk_index": 6,
-  "num_frames": 9,
-  "byte_length": 184320,
-  "started_event_ids": ["ui-update-42"],
-  "active_event_ids": ["ui-update-42"],
-  "completed_event_ids": []
-}
-```
-
-- `transport_chunk_index` identifies the binary WebSocket payload.
-- `generation_chunk_index` identifies the model-generation chunk represented
-  by the payload.
-- `started_event_ids` lists interactions that first affect this
-  generation chunk.
-- `active_event_ids` lists interactions contributing to this generation
-  chunk, including an ongoing transition.
-- `completed_event_ids` lists interactions that reach their target in
-  this generation chunk.
-
-Encoder trailer bytes are emitted as:
-
-```json
-{
-  "type": "video.chunk_metadata",
-  "request_id": "video_stream-...",
-  "kind": "trailer",
-  "transport_chunk_index": 8,
-  "generation_chunk_index": null,
-  "num_frames": 0,
-  "byte_length": 824,
-  "started_event_ids": [],
-  "active_event_ids": [],
-  "completed_event_ids": []
-}
-```

recipes/JD/JoyAI-VL-Interaction.md

@@ -168,7 +168,7 @@ the model + orchestrator + WebUI together — point `WEBUI_DIR` at your clone.
 python examples/online_serving/joyvl_interaction/cli/run_cli_demo.py \
   path/to/video.mp4 --query "Alert me if a fire breaks out"
 
-pytest tests/fullduplex   # framework + JoyVL unit tests
+pytest tests/e2e/features/fullduplex   # framework + JoyVL unit tests
 ```
 
 ## Testing with an RTSP stream (optional)

recipes/LTX/LTX-2.md

@@ -165,6 +165,40 @@ Pass these fields in online `extra_params` or offline `--extra-body`:
 --extra-body '{"video_cfg_scale":3.0,"audio_cfg_scale":7.0}'
 ```
 
+### Guidance Parallelism
+
+For LTX, `cfg_parallel_size` is the number of ranks used to execute the
+complete guidance plan in parallel. Despite the legacy `cfg` name, it covers
+text CFG, STG, and cross-modality guidance passes; guidance rescaling is
+applied after all pass predictions are gathered.
+
+The default one-stage recipe has four Transformer passes per denoise step:
+`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced
+configurations are therefore:
+
+| `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes |
+|---:|---:|---:|---|
+| `1` | 4 | 100% | Single-rank fused guidance batch |
+| `2` | 2 | 100% | Recommended two-rank configuration |
+| `4` | 1 | 100% | One guidance pass per rank |
+
+Other positive sizes are accepted. When the pass count is not divisible by
+`cfg_parallel_size`, ranks are padded to an equal number of execution slots
+and LTX emits a warning with the expected utilization. For example, four
+passes on three ranks use six slots and have 66.7% guidance-slot utilization.
+
+Start an LTX-2.3 server with two-way guidance parallelism:
+
+```bash
+vllm serve diffusers/LTX-2.3-Diffusers --omni \
+  --cfg-parallel-size 2 --stage-init-timeout 600
+```
+
+The total device count is the product of `cfg_parallel_size` and the other
+configured parallel dimensions. Distilled pipelines use positive-only
+guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for
+them.
+
 ### Python API
 
 Normal per-request values belong in `OmniDiffusionSamplingParams`. Pass the
@@ -273,8 +307,8 @@ bundled offline CLI do not currently expose `sigmas`.
 - Fused one-stage requests must resolve to identical LTX guidance and sigma
   schedules. Keep concurrent requests guidance-homogeneous or use
   `--max-num-seqs 1`.
-- `--cfg-parallel-size 2` supports only the CFG-only plan: disable STG and
-  modality guidance and set both rescale fields to `0.0`.
+- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
+  modality guidance, and rescale-compatible prediction gathering.
 
 ## Operational Notes
 

recipes/NVIDIA/GR00T-N1.7.md

@@ -69,8 +69,13 @@ Notes:
 ## Verification
 
 ```python
-from tests.gr00t.openpi_client_helper import run_policy_session, validate_session_result
-validate_session_result(run_policy_session(host="127.0.0.1", port=8000))
+from tests.helpers.runtime import OpenAIClientHandler
+
+handler = OpenAIClientHandler(host="127.0.0.1", port=8000, log_stats=False)
+response = handler.send_robot_openpi_ws_request(
+    {"run_default_policy_session": True, "session_id": "gr00t-smoke"}
+)[0]
+assert response.operation_responses[-1]["status"] == "reset successful"
 ```
 
 Or run the e2e test suite:

recipes/OpenBMB/MiniCPM-o-4_5.md

@@ -30,8 +30,6 @@ also provided.
   `hf_config.version="4.5"`):
   - Default single-GPU compatibility layout (auto-loaded):
     [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
-  - Recommended 2-GPU continuous-batching layout:
-    [`vllm_omni/deploy/minicpmo_4_5_batching.yaml`](../../vllm_omni/deploy/minicpmo_4_5_batching.yaml)
   - 2-GPU and 3-GPU layouts:
     [`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml),
     [`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
@@ -68,7 +66,7 @@ Code2Wav consumes them through a shared-memory async connector.
 MiniCPM-o 4.5 now requires the three-stage topology: the Talker owns
 request-local codec generation and Code2Wav owns waveform state and
 reference-voice prompt features. `minicpmo_4_5.yaml` remains the stable
-single-GPU entry point; `minicpmo_4_5_batching.yaml` is the recommended
+single-GPU entry point; `minicpmo_4_5_2gpu.yaml` is the recommended
 two-GPU profile. The removed fused two-stage implementation is not retained as
 a fallback because it would duplicate state machines and correctness paths.
 
@@ -79,7 +77,9 @@ a fallback because it would duplicate state machines and correctness paths.
 The default
 [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
 co-locates Thinker, codec-only Talker, and Code2Wav on GPU 0. Their
-`gpu_memory_utilization` budgets are 0.55, 0.22, and 0.22. All stages admit
+`gpu_memory_utilization` budgets are 0.55, 0.15, and 0.15. The remaining
+device memory is left available for runtime workspaces such as the HiFi-GAN
+vocoder's cuDNN kernels. All stages admit
 up to four sequences. Startup video profiling is bounded to 32 frames per
 video. Use the two-GPU profile for production throughput.
 
@@ -105,7 +105,7 @@ The deploy config is auto-loaded by the model registry — no
 For the recommended two-GPU layout, add:
 
 ```bash
---deploy-config vllm_omni/deploy/minicpmo_4_5_batching.yaml
+--deploy-config vllm_omni/deploy/minicpmo_4_5_2gpu.yaml
 ```
 
 #### Performance comparison
@@ -175,16 +175,18 @@ in another choice's `message.audio.data` (24 kHz mono, see Notes). With
 `modalities: ["text", "audio"]` you typically get two `choices` entries
 (one text, one audio).
 
-**Streaming text + speech**:
+**Streaming text + speech** (use `--stream`):
 
 ```bash
-python examples/online_serving/minicpmo/streaming_chat_completion.py \
-    --base-url http://localhost:8099/v1 \
-    --output minicpmo_stream.wav
+python examples/online_serving/minicpmo/openai_chat_completion_client_for_multimodal_generation.py \
+    --query-type text \
+    --prompt "Say hello, then introduce vLLM in one sentence." \
+    --port 8099 \
+    --stream
 ```
 
-The client prints text deltas as they arrive and reconstructs one valid WAV
-from the independently encoded audio deltas.
+The client prints text deltas as they arrive and saves streamed audio chunks
+to WAV files.
 
 **Gradio demo (text + image + audio + video UI)**:
 
@@ -202,9 +204,10 @@ speech output (TTS)"** checkbox on / off.
 
 #### Notes
 
-- Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.22, and
-  0.22 of GPU 0. The larger Thinker share protects its multimodal KV cache;
-  all three model processes still share one CUDA device.
+- Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.15, and
+  0.15 of GPU 0. The larger Thinker share protects its multimodal KV cache,
+  while the unreserved memory remains available to runtime kernels; all three
+  model processes still share one CUDA device.
 - `--trust-remote-code` is required — the HF repo ships a custom
   `MiniCPMO` config / model class.
 - Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains

recipes/cosmos3/Cosmos3-DistOffload.md

@@ -0,0 +1,191 @@
+# Cosmos3 — Distributed Layerwise Offload
+
+> Sharding weights across DP ranks with H2D + AllGather overlap for multi-NPU/GPU deployment
+
+## When to use this recipe
+
+Use this recipe when deploying Cosmos3-Super (64B / 124 GB) or Cosmos3-Nano (17B / 33 GB) across multiple devices where the model does not fit on a single device's HBM, or when you need higher throughput via DP multi-concurrency.
+
+**Key benefits:**
+- Each rank stores only 1/dp_size of the model weights on host memory
+- Only 2 transformer blocks reside on each device at any time (fixed double-buffer)
+- DP multi-concurrency: N concurrent requests with near-linear throughput scaling
+- mmap weight loading: shared page cache eliminates O(dp_size × model_size) RSS
+
+## Prerequisites
+
+- 2+ NPU or GPU devices (Ascend 910B3, NVIDIA B300, etc.)
+- vLLM-Omni with distributed layerwise offload support
+- Cosmos3-Nano or Cosmos3-Super checkpoint (local path or HuggingFace repo ID)
+
+## Serve commands
+
+### Cosmos3-Nano with DP=4 + AllGather (recommended for max throughput)
+
+```bash
+vllm serve /path/to/Cosmos3-Nano \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --no-guardrails --init-timeout 3600 \
+  --vae-use-tiling --vae-patch-parallel-size 1 \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4
+```
+
+### Cosmos3-Super with DP=2 + AllGather (for models > single-card HBM)
+
+```bash
+vllm serve /path/to/Cosmos3-Super \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --no-guardrails --init-timeout 3600 \
+  --vae-use-tiling --vae-patch-parallel-size 1 \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 2
+```
+
+### Cosmos3-Super with DP=4 + AllGather
+
+```bash
+vllm serve /path/to/Cosmos3-Super \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --no-guardrails --init-timeout 3600 \
+  --vae-use-tiling --vae-patch-parallel-size 1 \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4
+```
+
+### Without AllGather (full weights per rank, no sharding)
+
+Use `--dlo-no-use-allgather` when you want each rank to load full weights independently —
+no AllGather synchronization, no concurrent request requirement, but N× host memory.
+
+```bash
+vllm serve /path/to/Cosmos3-Nano \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --no-guardrails --init-timeout 3600 \
+  --vae-use-tiling --vae-patch-parallel-size 1 \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4 \
+  --dlo-no-use-allgather
+```
+
+### With SP (sequence parallel) instead of DP
+
+For long-sequence workloads where SP parallelism is preferred:
+
+```bash
+vllm serve /path/to/Cosmos3-Super \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --no-guardrails --init-timeout 3600 \
+  --vae-use-tiling --vae-patch-parallel-size 1 \
+  --enable-distributed-layerwise-offload \
+  --usp 4
+```
+
+## Request examples
+
+### Text-to-Image (T2I)
+
+```bash
+curl -s -o output.png -X POST "http://localhost:8000/v1/images/generations" \
+  -H "Content-Type: multipart/form-data" \
+  -F "model=/path/to/Cosmos3-Nano" \
+  -F "prompt=A robot arm cleaning a plate, cinematic" \
+  -F "size=1024x1024" \
+  -F "num_inference_steps=50" \
+  -F "guidance_scale=6.0" \
+  -F "seed=42"
+```
+
+### Text-to-Video (T2V)
+
+```bash
+curl -s -o output.mp4 -X POST "http://localhost:8000/v1/videos/sync" \
+  -H "Accept: video/mp4" \
+  -F "model=/path/to/Cosmos3-Nano" \
+  -F "prompt=A robot arm cleaning a plate, cinematic shot" \
+  -F "negative_prompt=blurry" \
+  -F "size=832x480" \
+  -F "num_frames=29" \
+  -F "fps=24" \
+  -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" \
+  -F "flow_shift=10.0" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \
+  -F "seed=17"
+```
+
+> **Important:** When using DP multi-concurrency (AllGather mode), `num_inference_steps`
+> must be specified explicitly and must be the same for all concurrent requests.
+> `num_inference_steps=None` is not allowed because it may resolve differently
+> per request mode (e.g., T2V=35 vs action_mode=30), which would deadlock AllGather.
+
+### Multiple concurrent requests (DP multi-concurrency)
+
+With `--data-parallel-size 4`, send 4 concurrent requests to maximize throughput:
+
+```bash
+for i in 1 2 3 4; do
+  curl -s -o output_${i}.mp4 -X POST "http://localhost:8000/v1/videos/sync" \
+    -H "Accept: video/mp4" \
+    -F "model=/path/to/Cosmos3-Nano" \
+    -F "prompt=A robot arm cleaning a plate, shot ${i}" \
+    -F "negative_prompt=blurry" \
+    -F "size=832x480" \
+    -F "num_frames=29" \
+    -F "fps=24" \
+    -F "num_inference_steps=35" \
+    -F "guidance_scale=6.0" \
+    -F "max_sequence_length=4096" \
+    -F "flow_shift=10.0" \
+    -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \
+    -F "seed=17" &
+done
+wait
+```
+
+## CLI flags
+
+| Flag | Description | Default |
+|------|-------------|---------|
+| `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` |
+| `--data-parallel-size N` | Number of DP ranks (weight sharding + concurrent requests) | `1` |
+| `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` |
+| `--dlo-no-use-allgather` | Each rank loads full weights independently (no sharding, no AllGather) | `false` |
+| `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` |
+
+## Mode comparison
+
+| Mode | CLI flags | CPU/rank | HBM/card | Throughput | Use case |
+|------|----------|---------|---------|-----------|----------|
+| DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences |
+| DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory |
+| DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) |
+| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Full model | 2 blocks | 0.4× HSDP | No AllGather overhead, high CPU |
+| HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM |
+
+## Memory expectations
+
+| Model | Config | Host RAM (cgroup) | HBM/card |
+|-------|--------|-------------------|----------|
+| Nano (33 GB) | DP2 + AllGather | ~38 GB | ~10 GB |
+| Nano (33 GB) | DP4 + AllGather | ~47 GB | ~10 GB |
+| Super (124 GB) | DP2 + AllGather | ~157 GB | ~15 GB |
+| Super (124 GB) | DP4 + AllGather | ~172 GB | ~10 GB |
+
+> Host RAM is dominated by page cache (1× model_size, shared across ranks).
+> HBM only holds 2 transformer blocks + framework + activations, independent of model size.
+
+## Limitations
+
+- `--dlo-use-allgather` (default) requires all concurrent requests to have the same
+  `num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
+- Online quantization (FP8) is incompatible with mmap loading — falls back to regular
+  `load_weights()` automatically
+- Tensor Parallel is not supported (DLO uses DP-based sharding, a different dimension)
+- HSDP + AllGather is rejected (would double-shard weights)

recipes/cosmos3/Cosmos3-Edge.md

@@ -17,7 +17,9 @@ The small Cosmos3 checkpoint, for constrained deployments. It runs on the shared
 
 - **Transformer**: `Cosmos3EdgeVFMTransformer` (Nemotron-based), auto-selected from
   `backbone_type: cosmos3_edge_nemotron_dense` — nothing to pass.
-- **Defaults** differ; omit `--height` / `--width` to get them:
+- **Defaults** differ. The offline examples auto-detect Edge (from `edge` in the model id /
+  transformer class) and apply the Edge column below, so you can omit `--height` / `--width` /
+  `--guidance-scale` / `--flow-shift` and still get native Edge output:
 
 | | Edge | Nano / Super |
 |---|---|---|
@@ -26,6 +28,10 @@ The small Cosmos3 checkpoint, for constrained deployments. It runs on the shared
 | Video `guidance_scale` | 5.0 | 6.0 |
 | Video `flow_shift` | 3.0 | 10.0 |
 
+> **Important:** these auto-defaults live in the example CLIs
+> (`text_to_video.py` / `image_to_video.py`). Do **not** pass the Nano/Super numbers
+> (720p / `guidance_scale` 6.0 / `flow_shift` 10.0) to Edge — they produce degenerate output.
+
 For the larger checkpoints see [Cosmos3-Nano](./Cosmos3-Nano.md) and [Cosmos3-Super](./Cosmos3-Super.md).
 
 ## Support notes

recipes/cosmos3/Cosmos3-Super.md

@@ -125,3 +125,134 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
   under the top-level `action` field. Verified on the 64B Super under
   `--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`)
   and the rollout video reliably.
+
+## NPU
+
+### 8× Ascend910 (A2, A3)
+
+#### Environment
+
+- OS: Linux
+- Python: 3.10+
+- Driver / runtime: Ascend NPU driver with CANN toolkit
+- Recommended operator library: **mindie-sd** (Ascend high-performance fused
+  operators — enables `adalayernorm` and other fused kernels automatically upon
+  installation)
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+A pre-built Docker image is available on
+[Docker Hub](https://hub.docker.com/r/vllm/vllm-omni) and
+[Quay.io](https://quay.io/ascend/vllm-omni). Ensure the image tag matches your
+vLLM-Omni checkout so that NPU-specific code is in sync with the container.
+
+#### Prerequisites
+
+Install the **mindie-sd** operator library to enable Ascend-optimized fused
+operators (`adalayernorm`, etc.):
+
+```bash
+git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD
+
+# Comment out the tik_ops build step (not needed for this use case)
+sed -i 's|^\(\s*\)source ${current_script_dir}/build_tik_ops.sh|\1# source ${current_script_dir}/build_tik_ops.sh|' build/build_ops.sh
+
+python setup.py bdist_wheel
+cd dist
+pip install mindiesd-*.whl
+```
+
+After installation, enable the Laser Attention kernel for significant
+long-sequence speedups:
+
+```bash
+export MINDIE_SD_FA_TYPE=ascend_laser_attention
+```
+
+#### Command
+
+```bash
+export MINDIE_SD_FA_TYPE=ascend_laser_attention
+
+vllm serve nvidia/Cosmos3-Super \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --tensor-parallel-size 8 \
+  --model-class-name Cosmos3OmniDiffusersPipeline \
+  --no-guardrails \
+  --init-timeout 1800
+```
+
+#### Verification
+
+Same requests as the GPU section above — all modes (T2I, T2V, I2V, V2V,
+T2VS, I2VS) work identically on NPU. Quick reference with
+`--no-guardrails`:
+
+```bash
+curl http://localhost:8000/v1/models
+
+# T2I (1024x1024, 50 steps)
+curl -sS -X POST http://localhost:8000/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "nvidia/Cosmos3-Super",
+    "prompt": "A photorealistic red sports car on a city street at golden hour, cinematic lighting.",
+    "size": "1024x1024", "n": 1, "response_format": "b64_json",
+    "num_inference_steps": 50, "guidance_scale": 7.0, "seed": 42
+  }' | python3 -c "import sys,json,base64; open('cosmos3_super_t2i.png','wb').write(base64.b64decode(json.load(sys.stdin)['data'][0]['b64_json']))"
+
+# T2V (1280×720, 189 frames, 35 steps — official params)
+curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \
+  -F "seed=17" -o cosmos3_super_t2v.mp4
+
+# I2V — add an uploaded reference image
+curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Super" -F "prompt=The scene comes to life with smooth, natural motion." \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \
+  -F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
+  -o cosmos3_super_i2v.mp4
+
+# V2V — add an uploaded reference video. condition_video_keep can be "first" or "last".
+curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Super" -F "prompt=Continue the same scene with smooth natural motion." \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first","guardrails":false}' \
+  -F "seed=2222" -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
+  -o cosmos3_super_v2v.mp4
+
+# T2V + sound — add generate_sound/sound_duration (output muxes AAC 48 kHz stereo)
+curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F "generate_sound=true" -F "sound_duration=7.875" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \
+  -F "seed=17" -o cosmos3_super_t2vs.mp4
+
+# I2V + sound — reference image with synchronized audio
+curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Super" -F "prompt=The scene comes to life with smooth, natural motion and ambient sound." \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F "generate_sound=true" -F "sound_duration=7.875" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \
+  -F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
+  -o cosmos3_super_i2vs.mp4
+```
+
+#### Notes
+
+- **Parallelism:** `--tensor-parallel-size 8` matches the model's 8 KV heads (`num_key_value_heads: 8`, `num_attention_heads: 64` → GQA). This uses 8 davinci devices (NPU 0–3, both cores each). The remaining NPU 4–7 stay idle.
+- **Model config:** 64 layers × 5120 hidden size × 64 attention heads, ~120 GB transformer on disk (27 shards). Each NPU device loads ~15 GB of model weights.
+- **Peak HBM:** ~22.7 GB per device at startup (bf16 weights). Additional memory is allocated per-generation for KV cache and activations — resolution and frame count drive the peak.
+- **Performance (verified on 8× Ascend910):** T2I 256² / 2 steps / guidance 1.0 → ~1.5 s.
+- **Guardrails** are disabled with `--no-guardrails` (guards are on by default). The gated `nvidia/Cosmos-1.0-Guardrail` model and `cosmos-guardrail` package are not shipped. Add `guardrails: false` in `extra_params` for per-request overrides when the server has guardrails enabled.
+- **Known limitations:** FP8 quantization not yet validated on Ascend NPU. `--enable-layerwise-offload` is available but untested on NPU.

tests/benchmarks/test_daily_omni_pack_modes.py

@@ -0,0 +1,249 @@
+"""Tests for Daily-Omni ``--daily-omni-pack-mode`` packing recipes.
+
+The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches
+OpenBMB's reported Daily-Omni score when frames and 1s audio segments arrive strictly
+interleaved, with ``max_slice_nums=1`` / ``use_image_id=False`` and an empty MiniCPM
+system message (not the Qwen system prompt). These tests pin that contract so a refactor
+cannot silently regress the score.
+
+vllm stubs are installed by tests/benchmarks/conftest.py before collection.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Load the data module directly (bypasses vllm_omni.__init__ heavy imports).
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_MODULE_PATH = _REPO_ROOT / "vllm_omni" / "benchmarks" / "data_modules" / "daily_omni_dataset.py"
+_MODULE_NAME = "vllm_omni.benchmarks.data_modules.daily_omni_dataset"
+
+if _MODULE_NAME not in sys.modules:
+    _spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
+    _mod = importlib.util.module_from_spec(_spec)
+    sys.modules[_MODULE_NAME] = _mod
+    _spec.loader.exec_module(_mod)
+
+from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (  # noqa: E402
+    DailyOmniDataset,
+    _uniform_sample_indices,
+)
+
+_VIDEO_ID = "vid001"
+_NUM_FRAMES = 5
+
+
+@pytest.fixture
+def video_dir(tmp_path: Path) -> Path:
+    """Official Daily-Omni ``Videos/`` layout with placeholder media files."""
+    per_video = tmp_path / "Videos" / _VIDEO_ID
+    per_video.mkdir(parents=True)
+    (per_video / f"{_VIDEO_ID}_video.mp4").write_bytes(b"fake-mp4")
+    (per_video / f"{_VIDEO_ID}_audio.wav").write_bytes(b"fake-wav")
+    return tmp_path / "Videos"
+
+
+@pytest.fixture
+def qa_json(tmp_path: Path) -> Path:
+    path = tmp_path / "qa.json"
+    path.write_text(
+        json.dumps(
+            [
+                {
+                    "video_id": _VIDEO_ID,
+                    "question": "What happens first?",
+                    "choice": {"A": "a", "B": "b", "C": "c", "D": "d"},
+                    "answer": "A",
+                    "task_type": "AV Event Alignment",
+                    "video_duration": "30s",
+                }
+            ]
+        ),
+        encoding="utf-8",
+    )
+    return path
+
+
+def _make_dataset(qa_json: Path, video_dir: Path, **kwargs) -> DailyOmniDataset:
+    return DailyOmniDataset(
+        qa_json_path=str(qa_json),
+        video_dir=str(video_dir),
+        disable_shuffle=True,
+        **kwargs,
+    )
+
+
+def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) -> None:
+    """Replace decord/ffmpeg extraction with deterministic in-memory frames + segments."""
+    from PIL import Image
+
+    def fake_extract(video_path, *, audio_path, include_audio, **_kw):
+        counter.append(1)
+        frames = [Image.new("RGB", (16, 16), color=(i, i, i)) for i in range(_NUM_FRAMES)]
+        segments = [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else []
+        return frames, segments
+
+    monkeypatch.setattr(dataset, "_extract_minicpm_frame_audio_segments", fake_extract)
+
+
+# ---------------------------------------------------------------------------
+# _uniform_sample_indices
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(("n", "k"), [(10, 4), (100, 64), (1000, 64), (7, 7), (3, 5), (1, 1)])
+def test_uniform_sample_indices_matches_openbmb_midpoint(n: int, k: int) -> None:
+    """Pin OpenBMB midpoint-bin sampling (not np.linspace endpoints)."""
+    got = _uniform_sample_indices(n, k)
+    if n <= k:
+        assert got == list(range(n))
+    else:
+        gap = n / k
+        expected = [int(i * gap + gap / 2) for i in range(k)]
+        assert got == expected
+        # Endpoints differ from linspace for typical Daily-Omni downsample sizes.
+        linspace = [int(i) for i in np.linspace(0, n - 1, k, dtype=int)]
+        if n == 1000 and k == 64:
+            assert got[0] == 7 and got[-1] == 992
+            assert linspace[0] == 0 and linspace[-1] == 999
+            assert got != linspace
+
+
+def test_uniform_sample_indices_non_positive_k() -> None:
+    assert _uniform_sample_indices(10, 0) == []
+    assert _uniform_sample_indices(10, -1) == []
+
+
+# ---------------------------------------------------------------------------
+# minicpm-interleave packing
+# ---------------------------------------------------------------------------
+
+
+def test_minicpm_interleave_alternates_image_and_audio(monkeypatch, qa_json, video_dir) -> None:
+    ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all")
+    _stub_extract(monkeypatch, ds, [])
+
+    parts, extra, position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    assert position == "first"
+    assert [p["type"] for p in parts] == ["image_url", "audio_url"] * _NUM_FRAMES
+    assert extra["mm_processor_kwargs"] == {
+        "use_audio_in_video": False,
+        "max_slice_nums": 1,
+        "use_image_id": False,
+    }
+
+
+def test_minicpm_interleave_visual_mode_emits_frames_only(monkeypatch, qa_json, video_dir) -> None:
+    ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="visual")
+    _stub_extract(monkeypatch, ds, [])
+
+    parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    assert [p["type"] for p in parts] == ["image_url"] * _NUM_FRAMES
+
+
+def test_minicpm_interleave_writes_file_urls_under_video_dir(monkeypatch, qa_json, video_dir) -> None:
+    ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all")
+    _stub_extract(monkeypatch, ds, [])
+
+    parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    segment_dir = video_dir / ".minicpm_daily_omni_interleave" / _VIDEO_ID
+    for part in parts:
+        url = part[part["type"]]["url"]
+        assert url.startswith("file://")
+        assert Path(url[len("file://") :]).parent == segment_dir.resolve()
+    assert sorted(p.name for p in segment_dir.glob("frame_*.jpg")) == [f"frame_{i:04d}.jpg" for i in range(_NUM_FRAMES)]
+
+
+def test_minicpm_interleave_inline_mode_emits_data_urls(monkeypatch, qa_json, video_dir) -> None:
+    ds = _make_dataset(
+        qa_json,
+        video_dir,
+        pack_mode="minicpm-interleave",
+        input_mode="all",
+        inline_local_video=True,
+    )
+    _stub_extract(monkeypatch, ds, [])
+
+    parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    assert parts[0]["image_url"]["url"].startswith("data:image/jpeg;base64,")
+    assert parts[1]["audio_url"]["url"].startswith("data:audio/wav;base64,")
+    assert not (video_dir / ".minicpm_daily_omni_interleave").exists()
+
+
+def test_minicpm_interleave_caches_file_parts_but_not_inline_payloads(monkeypatch, qa_json, video_dir) -> None:
+    """Inline base64 must not be retained: ~700 Daily-Omni videos would pin GBs of RSS."""
+    file_calls: list[int] = []
+    ds_file = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all")
+    _stub_extract(monkeypatch, ds_file, file_calls)
+    ds_file._compose_daily_omni_multimodal(_VIDEO_ID, None)
+    ds_file._compose_daily_omni_multimodal(_VIDEO_ID, None)
+    assert len(file_calls) == 1
+    assert ds_file._minicpm_interleave_cache
+
+    inline_calls: list[int] = []
+    ds_inline = _make_dataset(
+        qa_json,
+        video_dir,
+        pack_mode="minicpm-interleave",
+        input_mode="all",
+        inline_local_video=True,
+    )
+    _stub_extract(monkeypatch, ds_inline, inline_calls)
+    ds_inline._compose_daily_omni_multimodal(_VIDEO_ID, None)
+    ds_inline._compose_daily_omni_multimodal(_VIDEO_ID, None)
+    assert len(inline_calls) == 2
+    assert not ds_inline._minicpm_interleave_cache
+
+
+def test_minicpm_interleave_missing_local_video_returns_none(qa_json, tmp_path) -> None:
+    empty = tmp_path / "empty_videos"
+    empty.mkdir()
+    ds = _make_dataset(qa_json, empty, pack_mode="minicpm-interleave", input_mode="all")
+
+    parts, extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    assert parts is None
+    assert extra is None
+
+
+# ---------------------------------------------------------------------------
+# System prompt selection
+# ---------------------------------------------------------------------------
+
+
+def test_minicpm_interleave_keeps_empty_system_message(monkeypatch, qa_json, video_dir) -> None:
+    """Official MiniCPM omni mode still sends role=system with empty content."""
+    ds = _make_dataset(qa_json, video_dir, pack_mode="minicpm-interleave", input_mode="all")
+    _stub_extract(monkeypatch, ds, [])
+    parts, _extra, _position = ds._compose_daily_omni_multimodal(_VIDEO_ID, None)
+
+    messages = ds._build_daily_omni_openai_messages(parts, "q?", {"A": "a"})
+
+    assert [m["role"] for m in messages] == ["system", "user"]
+    assert messages[0]["content"][0]["text"] == ""
+
+
+def test_qwen_pack_mode_keeps_system_message(qa_json, video_dir) -> None:
+    ds = _make_dataset(qa_json, video_dir, pack_mode="qwen", input_mode="all")
+
+    messages = ds._build_daily_omni_openai_messages([], "q?", {"A": "a"})
+
+    assert [m["role"] for m in messages] == ["system", "user"]
+    assert "Qwen" in messages[0]["content"][0]["text"]
+
+
+def test_invalid_pack_mode_rejected(qa_json, video_dir) -> None:
+    with pytest.raises(ValueError, match="pack_mode must be"):
+        _make_dataset(qa_json, video_dir, pack_mode="bogus")

tests/config/test_config_factory.py

@@ -1145,7 +1145,7 @@ class TestResolveScheduler:
 
 class TestDeployConfigLoading:
     def test_load_minicpmo_duplex_deploy_config(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "minicpmo_4_5_duplex.yaml"
+        deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
 
         deploy = load_deploy_config(deploy_path)
         pipeline = resolve_pipeline_config("minicpmo_4_5")
@@ -1160,8 +1160,10 @@ class TestDeployConfigLoading:
         assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
         assert [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False, False]
         assert all("Async" not in (stage.scheduler_cls or "") for stage in stages)
+        assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"]
         assert deploy.stages[1].enforce_eager is False
         assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 2048
+        assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 0
         assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [1]
 
     @pytest.mark.parametrize(
@@ -1180,7 +1182,7 @@ class TestDeployConfigLoading:
         stage2_devices: str,
         stage1_replicas: int,
     ):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / filename
+        deploy_path = Path(get_deploy_config_path(filename))
 
         deploy = load_deploy_config(deploy_path)
         pipeline = resolve_pipeline_config("minicpmo_4_5")
@@ -1224,15 +1226,54 @@ stages:
         assert {s.yaml_engine_args.get("custom_voice_dir") for s in stages} == {str(custom_voice_dir)}
 
     def test_load_qwen3_omni_moe_deploy_config(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         deploy = load_deploy_config(deploy_path)
         assert len(deploy.stages) == 3
         assert deploy.async_chunk is True
         assert deploy.connectors is not None
         assert deploy.platforms is not None
 
+    def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self):
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml"))
+        with patch(
+            "vllm_omni.config.config_factory.get_config",
+            return_value=Q3_OMNI_THINKER_HF_CONFIG,
+        ):
+            pipeline = StageConfigFactory.get_pipeline_config(
+                "Qwen/Qwen3-Omni-30B-A3B-Thinking",
+                trust_remote_code=False,
+                deploy_config_path=str(deploy_path),
+            )
+
+        assert pipeline is not None
+        assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
+
+        deploy = load_deploy_config(deploy_path)
+        stages = merge_pipeline_deploy(pipeline, deploy)
+
+        assert len(stages) == 1
+        thinker = stages[0]
+        assert thinker.model_stage == "thinker"
+        assert thinker.final_output is True
+        assert thinker.final_output_type == "text"
+        assert thinker.worker_type == "ar"
+        assert thinker.scheduler_cls == "vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler"
+        assert thinker.hf_config_name == "thinker_config"
+        assert thinker.is_comprehension is True
+        assert thinker.yaml_runtime["devices"] == "0,1"
+        assert thinker.yaml_engine_args["engine_output_type"] == "text"
+        assert thinker.yaml_engine_args["tensor_parallel_size"] == 2
+        assert thinker.yaml_engine_args["max_num_seqs"] == 1
+        assert thinker.yaml_engine_args["gpu_memory_utilization"] == 0.9
+        assert thinker.yaml_engine_args["enforce_eager"] is True
+        assert thinker.yaml_engine_args["async_scheduling"] is False
+        assert thinker.yaml_engine_args["distributed_executor_backend"] == "mp"
+        assert thinker.yaml_engine_args["enable_prefix_caching"] is True
+        assert thinker.yaml_engine_args["async_chunk"] is False
+        assert thinker.yaml_extras["default_sampling_params"]["detokenize"] is True
+
     def test_load_voxtral_tts_deploy_config_schema_fields(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "voxtral_tts.yaml"
+        deploy_path = Path(get_deploy_config_path("voxtral_tts.yaml"))
         deploy = load_deploy_config(deploy_path)
         assert deploy.stages[0].config_format == "mistral"
         assert deploy.stages[0].load_format == "mistral"
@@ -1242,7 +1283,7 @@ stages:
         )
 
     def test_load_ming_flash_omni_deploy_config_schema_fields(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml"
+        deploy_path = Path(get_deploy_config_path("ming_flash_omni.yaml"))
         deploy = load_deploy_config(deploy_path)
         assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}}
         assert "compilation_config" not in deploy.stages[0].engine_extras
@@ -1382,7 +1423,7 @@ stages:
         assert not list(stage_config_dir.glob("*.yaml"))
 
     def test_merge_pipeline_deploy(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -1409,7 +1450,7 @@ stages:
         )
         assert isinstance(pipeline, PipelineConfig)
 
-        base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not base.exists():
             pytest.skip("Deploy config not found")
 
@@ -1829,7 +1870,7 @@ class TestQwen3TTSPipeline:
 
     def test_per_stage_model_arch_flows_through_merge(self, tmp_path):
         """Verify the new ps.model_arch override survives merge_pipeline_deploy."""
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_tts.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_tts.yaml"))
         if not deploy_path.exists():
             pytest.skip("qwen3_tts deploy yaml not found")
 
@@ -1953,7 +1994,7 @@ class TestMingFlashOmniPipeline:
 
     def test_full_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni.yaml parses and merges with the registered pipeline."""
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml"
+        deploy_path = Path(get_deploy_config_path("ming_flash_omni.yaml"))
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni deploy yaml not found")
 
@@ -1973,7 +2014,7 @@ class TestMingFlashOmniPipeline:
 
     def test_tts_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni_tts.yaml parses and routes to the TTS-only pipeline."""
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_tts.yaml"
+        deploy_path = Path(get_deploy_config_path("ming_flash_omni_tts.yaml"))
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni_tts deploy yaml not found")
 
@@ -2011,7 +2052,7 @@ class TestMingFlashOmniPipeline:
 
     def test_thinker_only_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni_thinker_only.yaml parses and routes to the thinker-only pipeline."""
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_thinker_only.yaml"
+        deploy_path = Path(get_deploy_config_path("ming_flash_omni_thinker_only.yaml"))
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni_thinker_only deploy yaml not found")
 
@@ -2071,7 +2112,7 @@ class TestMingFlashOmniPipeline:
 
     def test_image_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni_image.yaml parses and routes to the image pipeline."""
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_image.yaml"
+        deploy_path = Path(get_deploy_config_path("ming_flash_omni_image.yaml"))
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni_image deploy yaml not found")
 
@@ -2133,7 +2174,7 @@ class TestBaseConfigInheritance:
 
     def test_pure_inheritance_overlay(self, tmp_path):
         """An overlay with only ``base_config`` inherits everything."""
-        base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not base.exists():
             pytest.skip("Base deploy config not found")
 
@@ -2146,7 +2187,7 @@ class TestBaseConfigInheritance:
 
     def test_single_field_overlay(self, tmp_path):
         """An overlay overriding one stage field merges with the base."""
-        base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not base.exists():
             pytest.skip("Base deploy config not found")
 
@@ -2162,8 +2203,45 @@ class TestBaseConfigInheritance:
 class TestPlatformOverrides:
     """Test platform-specific deploy config overrides."""
 
+    def test_qwen3_tts_rocm_disables_code2wav_outer_cudagraph(self):
+        deploy_path = Path(get_deploy_config_path("qwen3_tts.yaml"))
+
+        base = load_deploy_config(deploy_path)
+        assert base.stages[0].enforce_eager is None
+        assert base.stages[1].enforce_eager is False
+
+        rocm = _apply_platform_overrides(base, platform="rocm")
+        assert rocm.stages[0].enforce_eager is None
+        assert rocm.stages[1].enforce_eager is True
+
+    def test_minicpmo_4_5_cuda_caps_talker_kv_cache(self):
+        pipeline = resolve_pipeline_config("minicpmo_4_5")
+        assert isinstance(pipeline, PipelineConfig)
+        deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml"))
+
+        cuda = _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda")
+        cuda_stages = merge_pipeline_deploy(pipeline, cuda)
+        assert cuda_stages[1].yaml_engine_args["kv_cache_memory_bytes"] == 2 * 1024**3
+
+        # The CUDA memory budget must not leak into the existing NPU profile.
+        npu = _apply_platform_overrides(load_deploy_config(deploy_path), platform="npu")
+        npu_stages = merge_pipeline_deploy(pipeline, npu)
+        assert "kv_cache_memory_bytes" not in npu_stages[1].yaml_engine_args
+
+        # Multi-GPU replica profiles place Talker on dedicated devices and
+        # preserve their previous automatic KV-cache sizing.
+        for filename in (
+            "minicpmo_4_5_3gpu_stage1_replicas.yaml",
+            "minicpmo_4_5_4gpu_stage1_replicas.yaml",
+        ):
+            replica = _apply_platform_overrides(
+                load_deploy_config(Path(get_deploy_config_path(filename))), platform="cuda"
+            )
+            replica_stages = merge_pipeline_deploy(pipeline, replica)
+            assert replica_stages[1].yaml_engine_args["kv_cache_memory_bytes"] is None
+
     def test_npu_overrides(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -2177,7 +2255,7 @@ class TestPlatformOverrides:
         assert deploy.stages[2].enforce_eager is False
 
     def test_xpu_overrides(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -2189,7 +2267,7 @@ class TestPlatformOverrides:
         assert deploy.stages[0].engine_extras.get("max_cudagraph_capture_size") == 0
 
     def test_unknown_platform_noop(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -2234,7 +2312,7 @@ class TestCLIOverrideFlow:
     """Test --stage-overrides JSON merge into StageConfig."""
 
     def test_stage_overrides_merge(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -2252,7 +2330,7 @@ class TestCLIOverrideFlow:
         assert stages[0].runtime_overrides["gpu_memory_utilization"] == 0.5
 
     def test_global_override_applies_to_all(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -2272,7 +2350,7 @@ class TestCLIOverrideFlow:
 
 class TestAuraOmniDeploy:
     def test_aura_omni_deploy_forces_pipeline_override(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"
+        deploy_path = Path(get_deploy_config_path("aura_omni.yaml"))
         deploy = load_deploy_config(deploy_path)
 
         assert deploy.pipeline == "aura_omni"
@@ -2283,7 +2361,7 @@ class TestAuraOmniDeploy:
         stages, _ = StageConfigFactory._create_legacy_from_registry(
             pipeline_cfg,
             cli_overrides={},
-            deploy_config_path=str(Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"),
+            deploy_config_path=get_deploy_config_path("aura_omni.yaml"),
         )
 
         assert [stage.model_stage for stage in stages] == [
@@ -2597,7 +2675,12 @@ class TestSentinelDefaultPrecedence:
         from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy
 
         pipeline_path = (
-            Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "models" / "qwen3_omni" / "pipeline.py"
+            Path(__file__).parent.parent.parent
+            / "vllm_omni"
+            / "model_executor"
+            / "models"
+            / "qwen3_omni"
+            / "pipeline.py"
         )
         pipeline = runpy.run_path(str(pipeline_path))["QWEN3_OMNI_PIPELINE"]
 
@@ -2667,7 +2750,7 @@ class TestSamplingConstraintsPrecedence:
     """Test that pipeline sampling_constraints override deploy defaults."""
 
     def test_constraints_win(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 

tests/config/test_omni_config.py

@@ -577,9 +577,11 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degr
     assert stage.parallel_config.world_size == 2
 
 
-def test_diffusion_parallel_config_rejects_cfg_parallel_size_outside_current_bound():
-    with pytest.raises(ValidationError):
-        OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
+def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism():
+    cfg = OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
+
+    assert cfg.cfg_parallel_size == 4
+    assert cfg.world_size == 4
 
 
 def test_diffusion_parallel_config_rejects_allgather_with_ulysses_or_ring():

tests/conftest.py

@@ -3,6 +3,7 @@ Root pytest entrypoint for the vLLM-Omni test suite.
 
 - `tests/conftest.py` stays thin: plugin registration + compatibility re-exports.
 - Importable utilities live under `tests/helpers/`.
+- Helper unit tests live under `tests/helpers/tests/`.
 - Fixtures live under `tests/helpers/fixtures/` and are loaded via `pytest_plugins`.
 """
 

tests/core/sched/test_omni_ar_scheduler_free_request_cleanup.py

@@ -0,0 +1,74 @@
+"""Regression tests for vllm-project/vllm-omni#5349 (P1): normal completion
+goes through _free_request(), not finish_requests() (the external
+abort/cancel entry point). Without a cleanup_receiver() call there,
+_active_streams never releases on ordinary completion, so the bounded-K
+window fills with stale entries after K completions.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request before
+# Request / RequestStatus are bound in this module. Ruff isort would reorder them.
+# isort: off
+import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_scheduler(*, chunk_transfer_adapter=None) -> OmniARScheduler:
+    """Minimal OmniARScheduler exercising _free_request()'s no-KV-transfer
+    happy path."""
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched._omits_kv_transfer_cache = {}
+    sched._connector_finished = lambda request: (False, None)
+    sched.encoder_cache_manager = MagicMock()
+    sched.finished_req_ids = set()
+    sched._new_prompt_len_snapshot = {}
+    sched.finished_req_ids_dict = None
+    sched._should_transfer_kv_for_request = lambda req_id: False
+    sched._free_blocks = MagicMock()
+    sched._free_input_coordinator_request = MagicMock()
+    sched.chunk_transfer_adapter = chunk_transfer_adapter
+    return sched
+
+
+class _FakeFinishedRequest:
+    """Not a SimpleNamespace: SimpleNamespace defines __eq__, which makes it
+    unhashable, and _free_request() puts the request in a set."""
+
+    def __init__(self, request_id: str) -> None:
+        self.request_id = request_id
+
+    def is_finished(self) -> bool:
+        return True
+
+
+def _make_finished_request(request_id: str = "req-1"):
+    return _FakeFinishedRequest(request_id)
+
+
+def test_free_request_releases_chunk_transfer_adapter_receiver_state():
+    adapter = MagicMock()
+    sched = _make_scheduler(chunk_transfer_adapter=adapter)
+    request = _make_finished_request("req-1")
+
+    sched._free_request(request)
+
+    adapter.cleanup_receiver.assert_called_once_with("req-1")
+
+
+def test_free_request_is_safe_without_a_chunk_transfer_adapter():
+    """Most stages have no chunk transfer adapter configured."""
+    sched = _make_scheduler(chunk_transfer_adapter=None)
+    request = _make_finished_request("req-1")
+
+    sched._free_request(request)  # must not raise
+
+    sched._free_input_coordinator_request.assert_called_once_with("req-1")

tests/core/sched/test_omni_ar_scheduler_kv_transfer.py

@@ -0,0 +1,57 @@
+from types import SimpleNamespace
+
+import pytest
+from vllm import SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.engine.async_engine_utils import apply_omni_final_stage_metadata
+from vllm_omni.engine.serialization import deserialize_additional_information
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_engine_request() -> EngineCoreRequest:
+    return EngineCoreRequest(
+        request_id="req",
+        prompt_token_ids=[1],
+        mm_features=None,
+        sampling_params=SamplingParams(max_tokens=1),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+    )
+
+
+def _request_omits_kv_transfer(*, force_kv_transfer: bool) -> tuple[bool, dict]:
+    tagged = apply_omni_final_stage_metadata(
+        _make_engine_request(),
+        final_stage_id=0,
+        force_kv_transfer=force_kv_transfer,
+    )
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler._omits_kv_transfer_cache = {}
+    request = SimpleNamespace(
+        request_id="req",
+        additional_information=tagged.additional_information,
+    )
+    result = scheduler._request_omits_kv_transfer_to_next_stage(request)
+    metadata = deserialize_additional_information(tagged.additional_information)
+    return result, metadata
+
+
+def test_stage_zero_request_omits_kv_transfer():
+    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=False)
+
+    assert omits_transfer
+    assert "omni_force_kv_transfer" not in metadata
+
+
+def test_cfg_companion_forces_kv_transfer_without_downstream_payload():
+    omits_transfer, metadata = _request_omits_kv_transfer(force_kv_transfer=True)
+
+    assert not omits_transfer
+    assert metadata["omni_final_stage_id"] == 0
+    assert metadata["omni_force_kv_transfer"] is True

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -76,6 +76,7 @@ class _Request:
         self.status = RequestStatus.RUNNING
         self.sampling_params = SimpleNamespace(num_logprobs=1)
         self.num_computed_tokens = 0
+        self.num_in_flight_tokens = 0
         self.num_output_placeholders = 0
         self.has_encoder_inputs = False
         self.pooling_params = None
@@ -104,7 +105,87 @@ class _RequestQueue(list):
                 self.remove(request)
 
 
-@pytest.mark.skip(reason="#issue 5484")
+def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
+    """Spec decode can sample tokens past a stop; the trailing tokens are
+    trimmed by _update_request_with_output, and the emitted logprob rows must
+    be trimmed with them (upstream slices logprobs after the trim)."""
+    request = _Request("req")
+    request.num_computed_tokens = 8
+    scheduler = SimpleNamespace(
+        perf_metrics=None,
+        connector=None,
+        chunk_transfer_adapter=None,
+        requests={request.request_id: request},
+        running=[request],
+        waiting=_RequestQueue(),
+        skipped_waiting=_RequestQueue(),
+        structured_output_manager=SimpleNamespace(should_advance=lambda _request: False),
+        transfer_triggered_requests=set(),
+        active_kv_transfers=set(),
+        pending_stop_after_extraction=set(),
+        waiting_for_transfer_free=set(),
+        _new_prompt_len_snapshot={},
+        finished_req_ids=set(),
+        finished_req_ids_dict=defaultdict(set),
+        kv_cache_manager=SimpleNamespace(take_events=lambda: None),
+        kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
+    )
+
+    def update_request_trimming(req, token_ids):
+        # Mirror vllm Scheduler._update_request_with_output: the first sampled
+        # token hits EOS, so the spec-accepted tokens behind it are trimmed
+        # in place and the request stops.
+        del token_ids[1:]
+        req.status = RequestStatus.FINISHED_STOPPED
+        return token_ids, True
+
+    def free_request(req):
+        scheduler.requests.pop(req.request_id, None)
+        scheduler.finished_req_ids.add(req.request_id)
+        scheduler.finished_req_ids_dict[req.client_index].add(req.request_id)
+        return None, None
+
+    scheduler._update_request_with_output = update_request_trimming
+    scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
+    scheduler._handle_stopped_request = lambda _request: True
+    scheduler._free_request = free_request
+    scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
+    scheduler.make_stats = lambda *args, **kwargs: None
+    scheduler._capture_omni_connector_output = lambda _output: None
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"req": 3},
+        scheduled_spec_decode_tokens={"req": [8, 9]},
+        num_invalid_spec_tokens=0,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=[[7, 8, 9]],
+        logprobs=_logprobs(
+            [[7, 1], [8, 2], [9, 3]],
+            [[-0.7, -1.0], [-0.8, -1.1], [-0.9, -1.2]],
+            cumulative_rows=[0],
+        ),
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={"req": 0},
+        routed_experts=None,
+    )
+
+    outputs = OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
+    (output,) = outputs[0].outputs
+
+    assert output.new_token_ids == [7]
+    assert output.finish_reason is FinishReason.STOP
+    token_rows = np.asarray(output.new_logprobs.logprob_token_ids)
+    value_rows = np.asarray(output.new_logprobs.logprobs)
+    assert token_rows.shape[0] == len(output.new_token_ids)
+    assert value_rows.shape[0] == len(output.new_token_ids)
+    np.testing.assert_array_equal(token_rows[:, 0], [7])
+
+
 def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
     """A bad runner response must not bring down unrelated batch requests."""
     bad = _Request("bad")
@@ -138,7 +219,8 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
         scheduler.requests.pop(request.request_id, None)
         scheduler.finished_req_ids.add(request.request_id)
         scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
-        return None
+        # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
+        return None, None
 
     scheduler._update_request_with_output = update_request
     scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -262,9 +262,52 @@ def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None:
         "tts_token_ids": [10, 20],
         "meta": {"replace_streaming_prompt": True},
     }
+    update.model_intermediate_buffer = {
+        "ids": {"tts": [41, 42, 99]},
+        "meta": {"turn_eos_token_id": 99},
+    }
 
     sched._update_request_as_session(session, update)
 
     assert session.prompt_token_ids == [10, 20]
     assert session.additional_information == update.additional_information
+    assert session.model_intermediate_buffer == {
+        "ids": {"tts": [41, 42, 99]},
+        "meta": {"turn_eos_token_id": 99},
+    }
+    assert session.status == RequestStatus.WAITING
+
+
+def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None:
+    sched = _make_scheduler(stage_id=1)
+    sched.chunk_transfer_adapter = SimpleNamespace(
+        receives_chunks=False,
+        segment_finished_requests=set(),
+    )
+    session = _make_request()
+    session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+    session.prompt_token_ids = [0] * 59
+    session._all_token_ids.clear()
+    session._all_token_ids.extend(session.prompt_token_ids)
+    session.num_prompt_tokens = 59
+    session.num_computed_tokens = 59
+    update = _make_update([0] * 10)
+    update.additional_information = None
+    update.model_intermediate_buffer = {
+        "ids": {"tts": list(range(8))},
+        "hidden_states": {"tts": [[0.0]] * 8},
+        "meta": {
+            "replace_streaming_prompt": True,
+            "next_stage_prompt_len": 10,
+        },
+    }
+
+    sched._update_request_as_session(session, update)
+
+    assert session.prompt_token_ids == [0] * 10
+    assert list(session._all_token_ids) == [0] * 10
+    assert session.num_prompt_tokens == 10
+    assert session.num_computed_tokens == 0
+    assert session.additional_information is None
+    assert session.model_intermediate_buffer == update.model_intermediate_buffer
     assert session.status == RequestStatus.WAITING

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -21,6 +21,7 @@ import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
 from vllm.sampling_params import SamplingParams
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.engine import EngineCoreEventType
+from vllm.v1.metrics.stats import PrefillStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
@@ -54,6 +55,18 @@ class _SchedulerStub(OmniGenerationScheduler):
         raise AssertionError("unexpected enqueue for skipped_waiting miss")
 
 
+class _AsyncChunkStopSchedulerStub(OmniGenerationScheduler):
+    def __init__(self) -> None:
+        self.num_waiting_for_streaming_input = 0
+        self.chunk_transfer_adapter = SimpleNamespace(receives_chunks=True)
+        self.enqueued: list[Request] = []
+        self.enqueued_statuses: list[RequestStatus] = []
+
+    def _enqueue_waiting_request(self, session: Request) -> None:
+        self.enqueued.append(session)
+        self.enqueued_statuses.append(session.status)
+
+
 def _make_request(**kwargs) -> Request:
     sp = SamplingParams(max_tokens=8)
     defaults = dict(
@@ -81,6 +94,17 @@ def _make_update(**kwargs) -> StreamingUpdate:
     return StreamingUpdate(**defaults)
 
 
+def test_generation_scheduler_records_prefill_stats_for_metrics() -> None:
+    request = _make_request(prompt_token_ids=[1, 2, 3, 4])
+    request.prefill_stats = PrefillStats()
+
+    OmniGenerationScheduler._record_prefill_stats(request)
+
+    assert request.prefill_stats.num_prompt_tokens == 4
+    assert request.prefill_stats.num_computed_tokens == 4
+    assert request.prefill_stats.num_cached_tokens == 0
+
+
 def test_resumable_generation_stop_marks_segment_boundary() -> None:
     session = _make_request(request_id="req-generation-segment")
     session.status = RequestStatus.RUNNING
@@ -133,6 +157,21 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
     assert output.is_segment_finished is True
 
 
+def test_async_chunk_resumable_stop_rearms_connector_polling() -> None:
+    sched = _AsyncChunkStopSchedulerStub()
+    session = _make_request(request_id="req-async-chunk-next-segment")
+    session.status = RequestStatus.FINISHED_STOPPED
+    session.resumable = True
+
+    finished = sched._handle_stopped_request(session)
+
+    assert finished is False
+    assert sched.enqueued == [session]
+    assert sched.enqueued_statuses == [RequestStatus.WAITING]
+    assert session.status == RequestStatus.WAITING
+    assert sched.num_waiting_for_streaming_input == 0
+
+
 class TestReplaceSessionWithStreamingUpdate:
     def test_resets_tokens_and_prompt_from_update(self) -> None:
         sched = _SchedulerStub()

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -75,6 +75,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
     coordinator = FakeInputCoordinator()
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler.input_coordinator = coordinator
+    scheduler.chunk_transfer_adapter = None
     scheduler._omits_kv_transfer_cache = {"req-free": True}
     scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
     scheduler.finished_req_ids = set()

tests/core/sched/test_omni_scheduler_mixin_timeouts.py

@@ -18,6 +18,8 @@ import pytest
 
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
 
 class _FakeCoordinator:
     def __init__(self, timed_out_ids):

tests/core/test_prefix_cache.py

@@ -3,6 +3,9 @@ import torch
 
 from vllm_omni.core.prefix_cache import OmniTensorPrefixCache
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
 DEFAULT_SEQ_LEN = 15
 NUM_BLOCKS = 10
 BLOCK_SIZE = 4

tests/dfx/perf/scripts/run_benchmark.py

@@ -1,6 +1,8 @@
 import json
 import os
 import threading
+from collections.abc import Callable
+from contextlib import AbstractContextManager, ExitStack, contextmanager
 from pathlib import Path
 from typing import Any
 
@@ -63,35 +65,78 @@ paired_benchmark_params = create_paired_omni_benchmark_pytest_params(BENCHMARK_C
 _omni_server_lock = threading.Lock()
 
 
+class _SingleActiveContext:
+    """Reuse one active context while its configuration key is unchanged."""
+
+    def __init__(self) -> None:
+        self._key: Any = None
+        self._stack: ExitStack | None = None
+        self._value: Any = None
+
+    def acquire(self, key: Any, factory: Callable[[], AbstractContextManager[Any]]) -> Any:
+        if self._stack is not None and key == self._key:
+            return self._value
+
+        self.close()
+        stack = ExitStack()
+        value = stack.enter_context(factory())
+        self._key = key
+        self._stack = stack
+        self._value = value
+        return value
+
+    def close(self) -> None:
+        stack = self._stack
+        self._key = None
+        self._stack = None
+        self._value = None
+        if stack is not None:
+            stack.close()
+
+
+@contextmanager
+def _start_omni_server(server_param):
+    test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = server_param
+
+    print(f"Starting OmniServer with test: {test_name}, model: {model}")
+
+    server_args: list[str] = []
+    if use_omni:
+        server_args += ["--stage-init-timeout", "600", "--init-timeout", "900"]
+    # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py):
+    # deploy-config sets the base; stage-overrides are applied on top. Both can be set.
+    if stage_config_path:
+        server_args = ["--deploy-config", stage_config_path] + server_args
+    if stage_overrides:
+        server_args = ["--stage-overrides", stage_overrides] + server_args
+    if extra_cli_args:
+        server_args = list(extra_cli_args) + server_args
+    with OmniServer(model, server_args, use_omni=use_omni) as server:
+        server.test_name = test_name
+        print("OmniServer started successfully")
+        yield server
+        print("OmniServer stopping...")
+
+    print("OmniServer stopped")
+
+
 @pytest.fixture(scope="module")
-def omni_server(request):
+def omni_server_context():
     """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.
+    Reuse it for adjacent benchmark cases with the same server configuration.
     Multi-stage initialization can take 10-20+ minutes.
     """
     with _omni_server_lock:
-        test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = request.param
+        active_context = _SingleActiveContext()
+        try:
+            yield active_context
+        finally:
+            active_context.close()
 
-        print(f"Starting OmniServer with test: {test_name}, model: {model}")
 
-        server_args: list[str] = []
-        if use_omni:
-            server_args += ["--stage-init-timeout", "600", "--init-timeout", "900"]
-        # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py):
-        # deploy-config sets the base; stage-overrides are applied on top. Both can be set.
-        if stage_config_path:
-            server_args = ["--deploy-config", stage_config_path] + server_args
-        if stage_overrides:
-            server_args = ["--stage-overrides", stage_overrides] + server_args
-        if extra_cli_args:
-            server_args = list(extra_cli_args) + server_args
-        with OmniServer(model, server_args, use_omni=use_omni) as server:
-            server.test_name = test_name
-            print("OmniServer started successfully")
-            yield server
-            print("OmniServer stopping...")
-
-        print("OmniServer stopped")
+@pytest.fixture
+def omni_server(request, omni_server_context):
+    return omni_server_context.acquire(request.param, lambda: _start_omni_server(request.param))
 
 
 @pytest.fixture

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -94,6 +94,7 @@
                   "vae_patch_parallel_size": 2,
                   "enable_expert_parallel": true
                 },
+                "vae_use_tiling": true,
                 "input_connectors": {
                   "from_stage_0": "rdma_connector"
                 },

tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json

@@ -21,6 +21,8 @@
         "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
         "tensor-parallel-size": 2,
         "cfg-parallel-size": 2,
+        "vae-patch-parallel-size": 4,
+        "vae-use-tiling": true,
         "enable-diffusion-pipeline-profiler": true
       }
     },

tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json

@@ -21,6 +21,8 @@
         "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
         "tensor-parallel-size": 2,
         "ulysses-degree": 2,
+        "vae-patch-parallel-size": 4,
+        "vae-use-tiling": true,
         "enable-diffusion-pipeline-profiler": true
       }
     },

tests/dfx/perf/tests/test_hunyuan_image_tp4.json

@@ -20,6 +20,8 @@
       "serve_args": {
         "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
         "tensor-parallel-size": 4,
+        "vae-patch-parallel-size": 4,
+        "vae-use-tiling": true,
         "enable-diffusion-pipeline-profiler": true
       }
     },

tests/dfx/perf/tests/test_runner_metadata.py

@@ -1,6 +1,8 @@
 """Tests for DFX runner metadata field exclusion."""
 
 import json
+from contextlib import contextmanager
+from types import SimpleNamespace
 
 import pytest
 
@@ -254,8 +256,9 @@ def test_benchmark_param_id_suffix_from_task_eval_phase():
     ]
 
 
-def test_create_paired_omni_benchmark_pytest_params(tmp_path):
+def test_paired_omni_benchmark_reuses_server_and_preserves_case_metadata(tmp_path, monkeypatch):
     from tests.dfx.conftest import create_paired_omni_benchmark_pytest_params
+    from tests.dfx.perf.scripts import run_benchmark
 
     configs = [
         {
@@ -285,3 +288,51 @@ def test_create_paired_omni_benchmark_pytest_params(tmp_path):
     assert bench_row == ("test_tts", 0)
     assert any(m.name == "tts" for m in by_id["test_tts-p0"].marks)
     assert not any(m.name == "tts" for m in by_id["test_omni-p0"].marks)
+
+    omni_p0_server = by_id["test_omni-p0"].values[0]
+    omni_p1_server = by_id["test_omni-p1"].values[0]
+    tts_server = by_id["test_tts-p0"].values[0]
+    assert omni_p0_server == omni_p1_server
+    assert omni_p0_server != tts_server
+
+    events = []
+
+    @contextmanager
+    def fake_start(server_param):
+        events.append(("start", server_param))
+        yield object()
+        events.append(("stop", server_param))
+
+    monkeypatch.setattr(run_benchmark, "_start_omni_server", fake_start)
+    active_context = run_benchmark._SingleActiveContext()
+    try:
+        first = run_benchmark.omni_server.__wrapped__(
+            SimpleNamespace(param=omni_p0_server),
+            active_context,
+        )
+        second = run_benchmark.omni_server.__wrapped__(
+            SimpleNamespace(param=omni_p1_server),
+            active_context,
+        )
+        assert first is second
+        assert events == [("start", omni_p0_server)]
+
+        third = run_benchmark.omni_server.__wrapped__(
+            SimpleNamespace(param=tts_server),
+            active_context,
+        )
+        assert third is not first
+        assert events == [
+            ("start", omni_p0_server),
+            ("stop", omni_p0_server),
+            ("start", tts_server),
+        ]
+    finally:
+        active_context.close()
+
+    assert events == [
+        ("start", omni_p0_server),
+        ("stop", omni_p0_server),
+        ("start", tts_server),
+        ("stop", tts_server),
+    ]

tests/dfx/reliability/invalid_param_test/test_invalid_minicpmo_4_5_omni.py

@@ -14,17 +14,16 @@ from tests.helpers.stage_config import get_deploy_config_path
 
 pytestmark = [pytest.mark.slow, pytest.mark.omni, pytest.mark.full_model]
 
-_MINICPMO_DEPLOY = get_deploy_config_path("minicpmo_4_5_batching.yaml")
+_MINICPMO_DEPLOY = get_deploy_config_path("minicpmo_4_5.yaml")
 
 _MINICPMO_SERVER = [
     pytest.param(
         OmniServerParams(
             model="openbmb/MiniCPM-o-4_5",
             stage_config_path=_MINICPMO_DEPLOY,
-            use_stage_cli=True,
+            use_stage_cli=False,
             server_args=[
                 "--trust-remote-code",
-                "--no-async-chunk",
             ],
         ),
         id="minicpmo_4_5",
@@ -51,7 +50,7 @@ def _prompt() -> str:
     return "What is the capital of China? Answer in 20 words."
 
 
-@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2)
+@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", _MINICPMO_SERVER, indirect=True)
 def test_invalid_audio_format_rejected(omni_server: OmniServerParams, openai_client: OpenAIClientHandler) -> None:
     """

tests/diffusion/ar_diffusion/test_engine_routing.py

@@ -21,6 +21,8 @@ from vllm_omni.experimental.ar_diffusion.engine import (
     ARDiffusionEngine,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def _cfg(backend):
     return SimpleNamespace(engine_backend=backend)

tests/diffusion/ar_diffusion/test_kv_cache.py

@@ -25,6 +25,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
 )
 from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 BLOCK = 16
 
 

tests/diffusion/ar_diffusion/test_kv_orchestration.py

@@ -7,6 +7,7 @@ import torch
 from vllm_omni.experimental.ar_diffusion.capability import ARDiffusionKVBranchSpec
 from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 DIMS = dict(
     num_layers=2,
     num_kv_heads=4,

tests/diffusion/ar_diffusion/test_paged_attention.py

@@ -22,6 +22,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
 )
 from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 BLOCK = 16
 N_HEADS = 4
 HEAD_DIM = 64
@@ -385,17 +388,22 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
         return out
 
     torch._dynamo.reset()
-    from torch._dynamo.testing import CompileCounter
+    try:
+        from torch._dynamo.testing import CompileCounter
 
-    counter = CompileCounter()
+        counter = CompileCounter()
 
-    def fn(inputs, q, k, v):
-        return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
+        def fn(inputs, q, k, v):
+            return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0
 
-    compiled = torch.compile(fn, backend=counter, fullgraph=True)
+        compiled = torch.compile(fn, backend=counter, fullgraph=True)
 
-    run_one_forward(commit=True)  # history grows between calls ->
-    run_one_forward(commit=True)  # block-table VALUES change, shapes don't
-    run_one_forward(commit=False)
+        run_one_forward(commit=True)  # history grows between calls ->
+        run_one_forward(commit=True)  # block-table VALUES change, shapes don't
+        run_one_forward(commit=False)
 
-    assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"
+        assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"
+    finally:
+        # Leave a clean dynamo state for later suites in the same pytest process
+        # (e.g. model_executor transformers models).
+        torch._dynamo.reset()

tests/diffusion/ar_diffusion/test_pipeline_bridge.py

@@ -12,6 +12,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
 )
 from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 BLOCK = 16
 N_HEADS = 4
 HEAD_DIM = 64

tests/diffusion/ar_diffusion/test_slot_mapping.py

@@ -13,6 +13,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
     resident_block_ids,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 BLOCK = 16
 
 

tests/diffusion/attention/test_attention_config.py

@@ -52,6 +52,44 @@ class TestAttentionSpec:
         with pytest.raises(ValueError, match="only supported by the TRTLLM_ATTN"):
             AttentionSpec(backend="TORCH_SDPA", skip_softmax={"target_sparsity": 0.5})
 
+    def test_quant_serialized_with_defaults_and_overrides(self):
+        assert AttentionSpec(backend="TRTLLM_ATTN", quant={"dtype_qk": "fp8_e4m3"}).backend_kwargs()["quant"] == {
+            "dtype_qk": "fp8_e4m3",
+            "q_block_size": 1,
+            "k_block_size": 16,
+        }
+        assert AttentionSpec(
+            backend="TRTLLM_ATTN", quant={"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16}
+        ).backend_kwargs()["quant"] == {"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16}
+
+    def test_quant_superset_fields_passed_through(self):
+        # FLASHINFER_ATTN-style config (dtype_vo / flashinfer_backend) round-trips on the shared spec.
+        bk = AttentionSpec(
+            backend="FLASHINFER_ATTN",
+            quant={"dtype_qk": "bfloat16", "dtype_vo": "fp8_e4m3", "flashinfer_backend": "trtllm-gen"},
+        ).backend_kwargs()["quant"]
+        assert bk["dtype_qk"] == "bfloat16"
+        assert bk["dtype_vo"] == "fp8_e4m3"
+        assert bk["flashinfer_backend"] == "trtllm-gen"
+
+    def test_quant_and_skip_softmax_coexist(self):
+        bk = AttentionSpec(
+            backend="TRTLLM_ATTN", quant={"dtype_qk": "int8"}, skip_softmax={"target_sparsity": 0.5}
+        ).backend_kwargs()
+        assert bk["target_sparsity"] == 0.5 and bk["quant"]["dtype_qk"] == "int8"
+
+    @pytest.mark.parametrize(
+        "spec, match",
+        [
+            ({"backend": "TORCH_SDPA", "quant": {"dtype_qk": "int8"}}, "only supported by the TRTLLM_ATTN"),
+            ({"backend": "TRTLLM_ATTN", "quant": {"dtype_qk": "int4"}}, "quant.dtype_qk"),
+            ({"backend": "TRTLLM_ATTN", "quant": {"dtype_qk": "int8", "k_block_size": 8}}, "quant.k_block_size"),
+        ],
+    )
+    def test_quant_validation_rejects(self, spec, match):
+        with pytest.raises(ValueError, match=match):
+            AttentionSpec(**spec)
+
 
 class TestAttentionConfig:
     def test_empty_config(self):
@@ -288,6 +326,13 @@ class TestBuildAttentionConfig:
 class TestOmniDiffusionConfigAttentionParsing:
     """Test OmniDiffusionConfig attention shorthand and structured config."""
 
+    @pytest.fixture(autouse=True)
+    def _clear_diffusion_attention_backend_env(self, monkeypatch):
+        # OmniDiffusionConfig.__post_init__ applies DIFFUSION_ATTENTION_BACKEND via
+        # build_attention_config(); clear it so these tests assert config defaults,
+        # not whatever the process inherited from CI / sibling tests.
+        monkeypatch.delenv("DIFFUSION_ATTENTION_BACKEND", raising=False)
+
     def test_diffusion_attention_backend_sets_default(self):
         config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="SAGE_ATTN")
         assert isinstance(config.diffusion_attention_config, AttentionConfig)

tests/diffusion/attention/test_attention_sp.py

@@ -49,6 +49,8 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.forward_context import set_forward_context
 from vllm_omni.platforms import current_omni_platform
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def update_environment_variables(envs_dict: dict[str, str]):
     """Update multiple environment variables with logging."""

tests/diffusion/attention/test_flash_attn.py

@@ -18,6 +18,9 @@ from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.diffusion.attention.backends.utils import fa  # noqa: E402
 from vllm_omni.platforms import current_omni_platform
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu()
 HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN
 flash_attn_func = fa.flash_attn_func  # noqa: N813

tests/diffusion/attention/test_piecewise_attn.py

@@ -27,6 +27,9 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 DEVICE = torch.device("cpu")
 
 

tests/diffusion/attention/test_sage_attn3.py

@@ -9,6 +9,9 @@ import pytest
 import torch
 from vllm.platforms import current_platform
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 SAGE_ATTN3_MODULE = "vllm_omni.diffusion.attention.backends.sage_attn3"
 
 

tests/diffusion/attention/test_selector.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for diffusion attention backend selection."""
+
+from types import SimpleNamespace
+
+import pytest
+
+import vllm_omni.diffusion.attention.selector as selector
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _ConfiguredBackend:
+    @classmethod
+    def get_name(cls) -> str:
+        return "CONFIGURED"
+
+
+class _PlatformBackend:
+    @classmethod
+    def get_name(cls) -> str:
+        return "PLATFORM"
+
+
+@pytest.fixture(autouse=True)
+def clear_selector_caches():
+    selector._cached_get_backend_cls.cache_clear()
+    selector._log_backend_resolution.cache_clear()
+    yield
+    selector._cached_get_backend_cls.cache_clear()
+    selector._log_backend_resolution.cache_clear()
+
+
+@pytest.mark.parametrize(
+    ("role", "role_category", "expected_backend"),
+    [
+        ("ltx2.audio_to_video", "cross", "EXACT"),
+        ("ltx2.video_to_audio", "cross", "CATEGORY"),
+        ("joint", None, "DEFAULT"),
+    ],
+)
+def test_configured_selection_precedence(
+    monkeypatch,
+    role: str,
+    role_category: str | None,
+    expected_backend: str,
+):
+    config = AttentionConfig(
+        default=AttentionSpec(backend="DEFAULT"),
+        per_role={
+            "cross": AttentionSpec(backend="CATEGORY"),
+            "ltx2.audio_to_video": AttentionSpec(backend="EXACT"),
+        },
+    )
+    calls = []
+
+    def fake_get_backend(backend_name, head_size, allow_trtllm_default=True):
+        calls.append((backend_name, head_size, allow_trtllm_default))
+        return _ConfiguredBackend
+
+    monkeypatch.setattr(selector, "_cached_get_backend_cls", fake_get_backend)
+
+    backend, spec = selector.get_attn_backend_for_role(
+        role=role,
+        role_category=role_category,
+        head_size=128,
+        attention_config=config,
+    )
+
+    assert backend is _ConfiguredBackend
+    assert spec is not None
+    assert spec.backend == expected_backend
+    assert calls == [(expected_backend, 128, True)]
+
+
+@pytest.mark.parametrize("attention_config", [None, AttentionConfig()])
+def test_platform_default_used_without_resolved_spec(monkeypatch, attention_config):
+    calls = []
+
+    def fake_get_backend(backend_name, head_size, allow_trtllm_default=True):
+        calls.append((backend_name, head_size, allow_trtllm_default))
+        return _PlatformBackend
+
+    monkeypatch.setattr(selector, "_cached_get_backend_cls", fake_get_backend)
+
+    backend, spec = selector.get_attn_backend_for_role(
+        role="self",
+        head_size=64,
+        attention_config=attention_config,
+        allow_trtllm_default=False,
+    )
+
+    assert backend is _PlatformBackend
+    assert spec is None
+    assert calls == [(None, 64, False)]
+
+
+def test_backend_class_resolution_is_cached(monkeypatch):
+    platform_calls = []
+    load_calls = []
+    fake_platform = SimpleNamespace(
+        get_diffusion_attn_backend_cls=lambda **kwargs: (platform_calls.append(kwargs) or "fake.module.Backend")
+    )
+
+    monkeypatch.setattr("vllm_omni.platforms.current_omni_platform", fake_platform)
+
+    def fake_load_backend(path):
+        load_calls.append(path)
+        return _ConfiguredBackend
+
+    monkeypatch.setattr(selector, "_load_backend_cls", fake_load_backend)
+
+    first = selector._cached_get_backend_cls("FLASH_ATTN", 128, False)
+    second = selector._cached_get_backend_cls("FLASH_ATTN", 128, False)
+
+    assert first is second is _ConfiguredBackend
+    assert platform_calls == [
+        {
+            "selected_backend": "FLASH_ATTN",
+            "head_size": 128,
+            "allow_trtllm_default": False,
+        }
+    ]
+    assert load_calls == ["fake.module.Backend"]
+
+
+def test_load_backend_cls_reports_missing_module():
+    with pytest.raises(ImportError, match="Failed to import module missing_attention_backend"):
+        selector._load_backend_cls("missing_attention_backend.Backend")
+
+
+def test_load_backend_cls_reports_missing_class(monkeypatch):
+    monkeypatch.setattr(selector.importlib, "import_module", lambda _: SimpleNamespace())
+
+    with pytest.raises(AttributeError, match="Class MissingBackend not found in module"):
+        selector._load_backend_cls("fake.module.MissingBackend")

tests/diffusion/attention/test_trtllm_attn.py

@@ -47,6 +47,44 @@ def test_skip_config_pure_resolution():
     assert cfg.resolve_factor(4096, timestep=None) == pytest.approx(0.01 * 4096)
 
 
+def test_quant_config_from_backend_kwargs():
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig
+
+    assert not QuantConfig.from_backend_kwargs(None).enabled
+    assert not QuantConfig.from_backend_kwargs({}).enabled
+    assert QuantConfig.from_backend_kwargs({"quant": {"dtype_qk": "fp8_e4m3"}}) == QuantConfig("fp8_e4m3", 1, 16)
+    qc = QuantConfig.from_backend_kwargs({"quant": {"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16}})
+    assert qc.enabled and (qc.dtype_qk, qc.q_block_size, qc.k_block_size) == ("int8", 4, 16)
+
+
+def test_quant_rejects_non_sage_dtype():
+    with pytest.raises(RuntimeError, match="supports dtype_qk"):
+        _impl(quant={"dtype_qk": "bfloat16"})
+
+
+def test_quant_quantize_requires_flashinfer_routine(monkeypatch):
+    import vllm_omni.diffusion.attention.backends.trtllm_attn as mod
+
+    monkeypatch.setattr(mod, "_sage_kernel_available", lambda: True)
+    monkeypatch.setattr(mod, "_sage_quantize_fn", lambda: None)
+    with pytest.raises(RuntimeError, match="trtllm_sage_attention_quantize"):
+        _impl(quant={"dtype_qk": "int8"})
+
+
+def test_quant_quantize_calls_routine_and_shapes_sfs():
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig
+
+    captured = {}
+
+    def fake_quantize(q, k, v, q_block_size, k_block_size, qk_quant_dtype):
+        captured.update(q_block_size=q_block_size, k_block_size=k_block_size, qk_quant_dtype=qk_quant_dtype)
+        return "qq", "kq", "vq", "qsfs", "ksfs", "vsfs"
+
+    q_q, k_q, v_q, sfs, blk = QuantConfig(dtype_qk="int8").quantize(object(), object(), object(), fake_quantize)
+    assert (q_q, k_q, v_q, sfs, blk) == ("qq", "kq", "vq", ("qsfs", "ksfs", None, "vsfs"), (1, 16, 0, 1))
+    assert captured == {"q_block_size": 1, "k_block_size": 16, "qk_quant_dtype": torch.int8}
+
+
 def test_skip_factor_none_without_curve():
     assert _impl(target_sparsity=0.5)._resolve_skip_factor(4096) is None
 
@@ -257,6 +295,32 @@ def test_bf16_dense_matches_sdpa():
     assert rel < 0.01, f"BF16 dense rel err {rel:.4f} too high"
 
 
+requires_sage = pytest.mark.skipif(
+    not (_has_trtllm_attn() and tg._sage_quantize_fn() is not None),
+    reason="requires Blackwell SM100+ GPU with flashinfer >= 0.6.16rc1 (trtllm_sage_attention_quantize)",
+)
+
+
+@requires_sage
+@pytest.mark.parametrize("dtype_qk", ["fp8_e4m3", "int8"])
+def test_sage_quant_matches_sdpa(dtype_qk):
+    torch.manual_seed(0)
+    B, S, H, D = 2, 256, 8, 128
+    scale = 1.0 / math.sqrt(D)
+    q, k, v = (torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    impl = _impl(quant={"dtype_qk": dtype_qk, "q_block_size": 1, "k_block_size": 16})
+    try:
+        out = impl.forward_cuda(q, k, v, None)
+    except RuntimeError as e:  # pragma: no cover - arch-dependent (int8 is SM100-only)
+        if "Missing TRTLLM-GEN kernel" in str(e):
+            pytest.skip(f"no trtllm-gen SAGE kernel for dtype_qk={dtype_qk} on this GPU arch")
+        raise
+    assert out.shape == (B, S, H, D) and torch.isfinite(out).all()
+    ref = _sdpa_ref(q, k, v, scale)
+    rel = (out.float() - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.1, f"SAGE({dtype_qk}) rel err {rel:.4f} too high"
+
+
 @requires_trtllm_attn
 def test_skip_tiny_threshold_matches_dense():
     torch.manual_seed(0)

tests/diffusion/attention/test_ulysses_uaa.py

@@ -22,6 +22,8 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.forward_context import get_forward_context, set_forward_context
 from vllm_omni.platforms import current_omni_platform
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def _find_free_port() -> int:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:

tests/diffusion/cache/test_magcache_hook.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for the MagCache hook forward numerics.
+
+The cached MagCache residual is the residual of the whole block stack
+(tail output - head input). On a skipped step the head hook applies it
+once; every other hooked block must pass its inputs through unchanged,
+matching the reference implementation in diffusers.hooks.mag_cache.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm_omni.diffusion.cache.magcache.config import MagCacheConfig
+from vllm_omni.diffusion.cache.magcache.hook import (
+    _MAG_CACHE_LEADER_BLOCK_HOOK,
+    apply_mag_cache_hook,
+)
+from vllm_omni.diffusion.hooks.base import HookRegistry
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _CountingBlock(nn.Module):
+    """Toy dual-stream block: adds a constant delta to hidden_states.
+
+    Returns (encoder_hidden_states, hidden_states) to match the default
+    TransformerBlockMetadata(return_hidden_states_index=1,
+    return_encoder_hidden_states_index=0) registered by the hook.
+    """
+
+    def __init__(self, delta: float):
+        super().__init__()
+        self.delta = delta
+        self.forward_calls = 0
+
+    def forward(self, hidden_states, encoder_hidden_states):
+        self.forward_calls += 1
+        return encoder_hidden_states, hidden_states + self.delta
+
+
+class _ToyTransformer(nn.Module):
+    def __init__(self, n_blocks: int, delta: float = 1.0):
+        super().__init__()
+        self.blocks = nn.ModuleList(_CountingBlock(delta) for _ in range(n_blocks))
+
+    def forward(self, hidden_states, encoder_hidden_states):
+        for block in self.blocks:
+            encoder_hidden_states, hidden_states = block(hidden_states, encoder_hidden_states)
+        return hidden_states
+
+    @property
+    def total_forward_calls(self) -> int:
+        return sum(block.forward_calls for block in self.blocks)
+
+
+def _make_config(num_inference_steps: int = 2) -> MagCacheConfig:
+    # retention_ratio=0.5 with 2 steps -> retention_step = 1: step 0 always
+    # computes, step 1 with mag_ratio 1.0 accumulates zero error, which is a
+    # guaranteed cache hit.
+    return MagCacheConfig(
+        transformer_type="Toy",
+        threshold=0.24,
+        max_skip_steps=5,
+        retention_ratio=0.5,
+        num_inference_steps=num_inference_steps,
+        mag_ratios=torch.ones(num_inference_steps),
+    )
+
+
+def test_skipped_step_applies_full_stack_residual_once():
+    n_blocks, delta = 3, 1.0
+    model = _ToyTransformer(n_blocks=n_blocks, delta=delta)
+    apply_mag_cache_hook(model, _make_config())
+
+    x = torch.zeros(1, 4, 8)
+    e = torch.zeros(1, 2, 8)
+
+    # Step 0 computes and caches residual = tail_output - head_input.
+    out_computed = model(x, e)
+    assert torch.allclose(out_computed, x + n_blocks * delta)
+    assert model.total_forward_calls == n_blocks
+
+    # Step 1 is a cache hit: no block computes, and the cached residual is
+    # applied exactly once, so the output matches the computed step.
+    out_skipped = model(x, e)
+    assert model.total_forward_calls == n_blocks, "step 1 was not skipped"
+    assert torch.allclose(out_skipped, out_computed), (
+        f"cached residual over-applied on skip: expected "
+        f"{out_computed.flatten()[0].item()}, got {out_skipped.flatten()[0].item()}"
+    )
+
+
+def test_single_block_advances_schedule_on_skip():
+    model = _ToyTransformer(n_blocks=1)
+    apply_mag_cache_hook(model, _make_config(num_inference_steps=2))
+
+    x = torch.zeros(1, 4, 8)
+    e = torch.zeros(1, 2, 8)
+
+    model(x, e)  # step 0: computed
+    assert model.total_forward_calls == 1
+    model(x, e)  # step 1: cache hit
+    assert model.total_forward_calls == 1, "step 1 was not skipped"
+
+    # The head hook is also the tail here, so a skipped step must still
+    # advance the schedule; step 1 was the last step, which resets the
+    # state for the next inference run.
+    registry = HookRegistry.check_if_exists_or_initialize(model.blocks[0])
+    hook = registry.get_hook(_MAG_CACHE_LEADER_BLOCK_HOOK)
+    state = hook.state_manager.get_state()
+    assert state.step_index == 0
+    assert state._is_first_step
+    assert state.previous_residual is None
+
+    # A new run therefore starts by computing, not by replaying the stale
+    # residual from the previous run.
+    model(x, e)
+    assert model.total_forward_calls == 2, "new run reused stale skip state"

tests/diffusion/distributed/test_comm.py

@@ -1,12 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for SeqAllToAll4D and SeqAllToAll5D communication primitives."""
+"""Tests for SeqAllToAll4D, SeqAllToAll5D, and RingComm communication primitives.
+
+CPU tests use gloo on CPU tensors (no GPU). Nightly parity tests run the same
+checks on real multi-GPU NCCL collectives.
+"""
+
+from __future__ import annotations
 
 import os
+from typing import Literal
 
 import pytest
 import torch
 
+from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -16,56 +24,51 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.platforms import current_omni_platform
 
+_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
+_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
 
-def update_environment_variables(envs_dict: dict[str, str]):
-    """Update multiple environment variables with logging."""
-    for k, v in envs_dict.items():
-        os.environ[k] = v
+DeviceKind = Literal["cpu", "cuda"]
 
 
-@pytest.mark.parametrize("world_size", [2, 4])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("seq_len_per_rank", [8])
-@pytest.mark.parametrize("num_heads", [8])
-@pytest.mark.parametrize("head_size", [32])
-@pytest.mark.parametrize("use_sync", [False, True])
-def test_4d_identity(
+def _update_environment_variables(envs_dict: dict[str, str]) -> None:
+    for key, value in envs_dict.items():
+        os.environ[key] = value
+
+
+def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
+    if device_kind == "cpu":
+        return torch.device("cpu")
+    return torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+
+
+def _init_worker(
+    local_rank: int,
     world_size: int,
-    dtype: torch.dtype,
-    batch_size: int,
-    seq_len_per_rank: int,
-    num_heads: int,
-    head_size: int,
-    use_sync: bool,
-):
-    """Test that two consecutive all-to-all operations return the original input."""
-    # Skip if not enough GPUs available
-    available_gpus = current_omni_platform.get_device_count()
-    if available_gpus < world_size:
-        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
-
-    # Ensure num_heads is divisible by world_size
-    if num_heads % world_size != 0:
-        pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
-
-    # Run test with multiprocessing spawn
-    torch.multiprocessing.spawn(
-        _test_4d_identity_worker,
-        args=(
-            world_size,
-            dtype,
-            batch_size,
-            seq_len_per_rank,
-            num_heads,
-            head_size,
-            use_sync,
-        ),
-        nprocs=world_size,
+    master_port: int,
+    device_kind: DeviceKind,
+) -> None:
+    _update_environment_variables(
+        {
+            "RANK": str(local_rank),
+            "LOCAL_RANK": str(local_rank),
+            "WORLD_SIZE": str(world_size),
+            "MASTER_ADDR": "localhost",
+            "MASTER_PORT": str(master_port),
+        }
     )
+    backend = "gloo" if device_kind == "cpu" else None
+    init_distributed_environment(backend=backend)
 
 
-def _test_4d_identity_worker(
+def _close_tolerance(dtype: torch.dtype) -> tuple[float, float]:
+    if dtype == torch.bfloat16:
+        return 1e-3, 1e-3
+    if dtype == torch.float16:
+        return 1e-3, 1e-3
+    return 1e-5, 1e-5
+
+
+def _run_4d_identity(
     local_rank: int,
     world_size: int,
     dtype: torch.dtype,
@@ -74,29 +77,17 @@ def _test_4d_identity_worker(
     num_heads: int,
     head_size: int,
     use_sync: bool,
-):
-    """Worker function for test_4d_identity."""
-    # Set device
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    device = _worker_device(local_rank, device_kind)
+    if device_kind == "cuda":
+        current_omni_platform.set_device(device)
 
-    # Set environment variables for distributed training
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "29500",
-        }
-    )
+    _init_worker(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(ulysses_degree=world_size)
+    sp_group = get_sp_group().ulysses_group
 
-    # Initialize distributed environment
-    init_distributed_environment()
-    initialize_model_parallel(ulysses_degree=world_size)  # test ulysses sp by default
-    sp_group = get_sp_group().ulysses_group  # get ulysses sp group not ring sp group
-
-    # Create input tensor: (bs, seqlen/P, hc, hs)
     torch.manual_seed(42 + local_rank)
     input_tensor = torch.randn(
         batch_size,
@@ -106,100 +97,36 @@ def _test_4d_identity_worker(
         dtype=dtype,
         device=device,
     )
-
-    # Save original input for comparison
     original_input = input_tensor.clone()
 
-    # First all-to-all: (bs, seqlen/P, hc, hs) -> (bs, seqlen, hc/P, hs)
-    intermediate = SeqAllToAll4D.apply(
-        sp_group,
-        input_tensor,
-        2,  # scatter head dimension
-        1,  # gather sequence dimension
-        use_sync,
-    )
-
-    # Verify intermediate shape
-    expected_shape = (
+    intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, use_sync)
+    expected_intermediate_shape = (
         batch_size,
         seq_len_per_rank * world_size,
         num_heads // world_size,
         head_size,
     )
-    assert intermediate.shape == expected_shape, (
-        f"Intermediate shape mismatch: expected {expected_shape}, got {intermediate.shape}"
+    assert intermediate.shape == expected_intermediate_shape, (
+        f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
     )
 
-    # Second all-to-all: (bs, seqlen, hc/P, hs) -> (bs, seqlen/P, hc, hs)
-    output = SeqAllToAll4D.apply(
-        sp_group,
-        intermediate,
-        1,  # scatter sequence dimension
-        2,  # gather head dimension
-        use_sync,
-    )
-
-    # Verify output shape matches input
+    output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, use_sync)
     assert output.shape == original_input.shape, (
         f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
     )
 
-    # Verify output matches original input
+    rtol, atol = _close_tolerance(dtype)
     torch.testing.assert_close(
         output,
         original_input,
-        rtol=1e-5,
-        atol=1e-5,
+        rtol=rtol,
+        atol=atol,
         msg="Output does not match original input after two all-to-all operations",
     )
-
-    # Cleanup distributed environment
     destroy_distributed_env()
 
 
-@pytest.mark.parametrize("world_size", [2, 4])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("seq_len_per_rank", [8])
-@pytest.mark.parametrize("num_heads", [8])
-@pytest.mark.parametrize("head_size", [32])
-@pytest.mark.parametrize("use_sync", [False, True])
-def test_5d_identity(
-    world_size: int,
-    dtype: torch.dtype,
-    batch_size: int,
-    seq_len_per_rank: int,
-    num_heads: int,
-    head_size: int,
-    use_sync: bool,
-):
-    """Test that two consecutive all-to-all operations return the original input."""
-    # Skip if not enough GPUs available
-    available_gpus = current_omni_platform.get_device_count()
-    if available_gpus < world_size:
-        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
-
-    # Ensure num_heads is divisible by world_size
-    if num_heads % world_size != 0:
-        pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
-
-    # Run test with multiprocessing spawn
-    torch.multiprocessing.spawn(
-        _test_5d_identity_worker,
-        args=(
-            world_size,
-            dtype,
-            batch_size,
-            seq_len_per_rank,
-            num_heads,
-            head_size,
-            use_sync,
-        ),
-        nprocs=world_size,
-    )
-
-
-def _test_5d_identity_worker(
+def _run_5d_identity(
     local_rank: int,
     world_size: int,
     dtype: torch.dtype,
@@ -208,199 +135,329 @@ def _test_5d_identity_worker(
     num_heads: int,
     head_size: int,
     use_sync: bool,
-):
-    """Worker function for test_5d_identity."""
-    # Set device
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    device = _worker_device(local_rank, device_kind)
+    if device_kind == "cuda":
+        current_omni_platform.set_device(device)
 
-    # Set environment variables for distributed training
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "29500",
-        }
-    )
+    _init_worker(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(ulysses_degree=world_size)
+    sp_group = get_sp_group().ulysses_group
 
-    # Initialize distributed environment
-    init_distributed_environment()
-    initialize_model_parallel(ulysses_degree=world_size)  # test ulysses sp by default
-    sp_group = get_sp_group().ulysses_group  # get ulysses sp group not ring sp group
-
-    # Create input tensor: (bs, seqlen/P, 3, hc, hs)
-    # The '3' dimension is for Q, K, V
     torch.manual_seed(42 + local_rank)
     input_tensor = torch.randn(
         batch_size,
         seq_len_per_rank,
-        3,  # Q, K, V
+        3,
         num_heads,
         head_size,
         dtype=dtype,
         device=device,
     )
-
-    # Save original input for comparison
     original_input = input_tensor.clone()
 
-    # First all-to-all: (bs, seqlen/P, 3, hc, hs) -> (bs, seqlen, 3, hc/P, hs)
-    intermediate = SeqAllToAll5D.apply(
-        sp_group,
-        input_tensor,
-        3,  # scatter head dimension
-        1,  # gather sequence dimension
-        use_sync,
-    )
-
-    # Verify intermediate shape
-    expected_shape = (
+    intermediate = SeqAllToAll5D.apply(sp_group, input_tensor, 3, 1, use_sync)
+    expected_intermediate_shape = (
         batch_size,
         seq_len_per_rank * world_size,
         3,
         num_heads // world_size,
         head_size,
     )
-    assert intermediate.shape == expected_shape, (
-        f"Intermediate shape mismatch: expected {expected_shape}, got {intermediate.shape}"
+    assert intermediate.shape == expected_intermediate_shape, (
+        f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
     )
 
-    # Second all-to-all: (bs, seqlen, 3, hc/P, hs) -> (bs, seqlen/P, 3, hc, hs)
-    output = SeqAllToAll5D.apply(
-        sp_group,
-        intermediate,
-        1,  # scatter sequence dimension
-        3,  # gather head dimension
-        use_sync,
-    )
-
-    # Verify output shape matches input
+    output = SeqAllToAll5D.apply(sp_group, intermediate, 1, 3, use_sync)
     assert output.shape == original_input.shape, (
         f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
     )
 
-    # Verify output matches original input
+    rtol, atol = _close_tolerance(dtype)
     torch.testing.assert_close(
         output,
         original_input,
-        rtol=1e-5,
-        atol=1e-5,
+        rtol=rtol,
+        atol=atol,
         msg="Output does not match original input after two all-to-all operations",
     )
-
-    # Cleanup distributed environment
     destroy_distributed_env()
 
 
-@pytest.mark.parametrize("world_size", [2, 4])
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("batch_size", [2])
-@pytest.mark.parametrize("num_heads", [8])
-@pytest.mark.parametrize("head_size", [128])
-def test_ring_p2p(
-    world_size: int,
-    dtype: torch.dtype,
-    batch_size: int,
-    num_heads: int,
-    head_size: int,
-):
-    """Test Ring P2P communication (send_recv)."""
-    # Skip if not enough GPUs available
-    available_gpus = current_omni_platform.get_device_count()
-    if available_gpus < world_size:
-        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
-
-    torch.multiprocessing.spawn(
-        _test_ring_p2p_worker,
-        args=(world_size, dtype, batch_size, num_heads, head_size),
-        nprocs=world_size,
-    )
-
-
-def _test_ring_p2p_worker(
+def _run_ring_p2p(
     local_rank: int,
     world_size: int,
     dtype: torch.dtype,
     batch_size: int,
     num_heads: int,
     head_size: int,
-):
-    """Worker for Ring P2P test."""
-    import sys
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    device = _worker_device(local_rank, device_kind)
+    if device_kind == "cuda":
+        current_omni_platform.set_device(device)
 
-    # Set device
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
-
-    # Set env vars
-    # Use a different port to avoid conflict with other tests if run in parallel
-    update_environment_variables(
-        {
-            "RANK": str(local_rank),
-            "LOCAL_RANK": str(local_rank),
-            "WORLD_SIZE": str(world_size),
-            "MASTER_ADDR": "localhost",
-            "MASTER_PORT": "29501",
-        }
-    )
-
-    # Init distributed
     try:
-        init_distributed_environment()
-        # Ring degree = world_size to test ring group
+        _init_worker(local_rank, world_size, master_port, device_kind)
         initialize_model_parallel(ring_degree=world_size)
         sp_group = get_sp_group()
-
-        print(f"[Rank {local_rank}] Initialized. Ring group size: {sp_group.ring_group.size()}")
-        sys.stdout.flush()
-
-        # Create RingComm
         comm = RingComm(sp_group.ring_group)
 
-        # Create tensor: rank-specific data
-        # (batch, num_heads, head_size)
-        # Fill with rank value + 1 to avoid 0 and make verification easy
         input_tensor = torch.full(
-            (batch_size, num_heads, head_size), fill_value=float(local_rank + 1), dtype=dtype, device=device
+            (batch_size, num_heads, head_size),
+            fill_value=float(local_rank + 1),
+            dtype=dtype,
+            device=device,
         )
-
-        print(f"[Rank {local_rank}] Input sum: {input_tensor.sum().item()}")
-        sys.stdout.flush()
-
-        # Send input, receive from prev
-        # RingComm.send_recv sends to next, receives from prev
-        t0 = __import__("time").time()
         recv_tensor = comm.send_recv(input_tensor)
         comm.commit()
         comm.wait()
-        t1 = __import__("time").time()
 
-        print(f"[Rank {local_rank}] Communication done in {t1 - t0:.4f}s")
-
-        # Verify
-        # Expected value: from (rank - 1) % world_size
         prev_rank = (local_rank - 1 + world_size) % world_size
         expected_value = float(prev_rank + 1)
-
-        recv_sum = recv_tensor.sum().item()
-        print(f"[Rank {local_rank}] Received sum: {recv_sum}, Expected value: {expected_value}")
-        sys.stdout.flush()
-
         expected_tensor = torch.full_like(recv_tensor, fill_value=expected_value)
-
-        # Use a slightly loose tolerance for bfloat16
+        rtol, atol = _close_tolerance(dtype)
         torch.testing.assert_close(
-            recv_tensor, expected_tensor, rtol=1e-3, atol=1e-3, msg=f"[Rank {local_rank}] Data mismatch!"
+            recv_tensor,
+            expected_tensor,
+            rtol=rtol,
+            atol=atol,
+            msg=f"[Rank {local_rank}] Ring P2P data mismatch",
         )
-        print(f"[Rank {local_rank}] Verification PASSED")
-
-    except Exception as e:
-        print(f"[Rank {local_rank}] FAILED with error: {e}")
-        import traceback
-
-        traceback.print_exc()
-        raise e
     finally:
         destroy_distributed_env()
+
+
+def _spawn_4d_identity(
+    *,
+    world_size: int,
+    dtype: torch.dtype,
+    batch_size: int,
+    seq_len_per_rank: int,
+    num_heads: int,
+    head_size: int,
+    use_sync: bool,
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    torch.multiprocessing.spawn(
+        _run_4d_identity,
+        args=(
+            world_size,
+            dtype,
+            batch_size,
+            seq_len_per_rank,
+            num_heads,
+            head_size,
+            use_sync,
+            device_kind,
+            master_port,
+        ),
+        nprocs=world_size,
+    )
+
+
+def _spawn_5d_identity(
+    *,
+    world_size: int,
+    dtype: torch.dtype,
+    batch_size: int,
+    seq_len_per_rank: int,
+    num_heads: int,
+    head_size: int,
+    use_sync: bool,
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    torch.multiprocessing.spawn(
+        _run_5d_identity,
+        args=(
+            world_size,
+            dtype,
+            batch_size,
+            seq_len_per_rank,
+            num_heads,
+            head_size,
+            use_sync,
+            device_kind,
+            master_port,
+        ),
+        nprocs=world_size,
+    )
+
+
+def _spawn_ring_p2p(
+    *,
+    world_size: int,
+    dtype: torch.dtype,
+    batch_size: int,
+    num_heads: int,
+    head_size: int,
+    device_kind: DeviceKind,
+    master_port: int,
+) -> None:
+    torch.multiprocessing.spawn(
+        _run_ring_p2p,
+        args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
+        nprocs=world_size,
+    )
+
+
+def _require_heads_divisible(num_heads: int, world_size: int) -> None:
+    if num_heads % world_size != 0:
+        pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})")
+
+
+def _require_gpus(world_size: int) -> None:
+    available_gpus = current_omni_platform.get_device_count()
+    if available_gpus < world_size:
+        pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available")
+
+
+# ---------------------------------------------------------------------------
+# CPU: gloo collectives on CPU tensors (no GPU)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_4d_identity(world_size: int, use_sync: bool):
+    _require_heads_divisible(8, world_size)
+    _spawn_4d_identity(
+        world_size=world_size,
+        dtype=torch.float32,
+        batch_size=2,
+        seq_len_per_rank=8,
+        num_heads=8,
+        head_size=32,
+        use_sync=use_sync,
+        device_kind="cpu",
+        master_port=29610,
+    )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_5d_identity(world_size: int, use_sync: bool):
+    _require_heads_divisible(8, world_size)
+    _spawn_5d_identity(
+        world_size=world_size,
+        dtype=torch.float32,
+        batch_size=2,
+        seq_len_per_rank=8,
+        num_heads=8,
+        head_size=32,
+        use_sync=use_sync,
+        device_kind="cpu",
+        master_port=29611,
+    )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("world_size", [2, 4])
+def test_ring_p2p(world_size: int):
+    _spawn_ring_p2p(
+        world_size=world_size,
+        dtype=torch.float32,
+        batch_size=2,
+        num_heads=8,
+        head_size=128,
+        device_kind="cpu",
+        master_port=29612,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Nightly: same checks on real multi-GPU NCCL collectives
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+    "world_size",
+    [
+        pytest.param(2, marks=_L4_TWO_GPU),
+        pytest.param(4, marks=_L4_FOUR_GPU),
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_4d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
+    _require_gpus(world_size)
+    _require_heads_divisible(8, world_size)
+    _spawn_4d_identity(
+        world_size=world_size,
+        dtype=dtype,
+        batch_size=2,
+        seq_len_per_rank=8,
+        num_heads=8,
+        head_size=32,
+        use_sync=use_sync,
+        device_kind="cuda",
+        master_port=29500,
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+    "world_size",
+    [
+        pytest.param(2, marks=_L4_TWO_GPU),
+        pytest.param(4, marks=_L4_FOUR_GPU),
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("use_sync", [False, True])
+def test_5d_identity_parity(world_size: int, dtype: torch.dtype, use_sync: bool):
+    _require_gpus(world_size)
+    _require_heads_divisible(8, world_size)
+    _spawn_5d_identity(
+        world_size=world_size,
+        dtype=dtype,
+        batch_size=2,
+        seq_len_per_rank=8,
+        num_heads=8,
+        head_size=32,
+        use_sync=use_sync,
+        device_kind="cuda",
+        master_port=29502,
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+    "world_size",
+    [
+        pytest.param(2, marks=_L4_TWO_GPU),
+        pytest.param(4, marks=_L4_FOUR_GPU),
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_ring_p2p_parity(world_size: int, dtype: torch.dtype):
+    _require_gpus(world_size)
+    _spawn_ring_p2p(
+        world_size=world_size,
+        dtype=dtype,
+        batch_size=2,
+        num_heads=8,
+        head_size=128,
+        device_kind="cuda",
+        master_port=29501,
+    )

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -1,8 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for PipelineParallelMixin and related PP communication helpers.
+
+CPU unit tests exercise mixin logic without a process group. Distributed CPU
+tests use gloo on CPU tensors; nightly parity tests run the same checks on
+real multi-GPU NCCL collectives.
+"""
+
+from __future__ import annotations
 
 import os
 import socket
+from typing import Literal
 
 import pytest
 import torch
@@ -11,6 +20,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
 
 import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
+from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
     destroy_distributed_env,
@@ -22,7 +32,12 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
 from vllm_omni.platforms import current_omni_platform
 
-pytestmark = [pytest.mark.parallel]
+_L4_TWO_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=2)
+_L4_THREE_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=3)
+_L4_FOUR_GPU = hardware_marks(res={"cuda": "L4"}, num_cards=4)
+
+DeviceKind = Literal["cpu", "cuda"]
+_UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
 def _find_free_port() -> str:
@@ -36,6 +51,22 @@ def update_environment_variables(envs_dict: dict[str, str]) -> None:
         os.environ[k] = v
 
 
+def _cleanup_distributed() -> None:
+    """Destroy omni distributed state and clear test-process leftovers.
+
+    ``destroy_distributed_env`` leaves ``vllm.distributed.parallel_state._PP``
+    aliased to the destroyed group. Parent-process baselines must clear that
+    dangling reference (and RANK/MASTER_* env) so later unit tests that call
+    ``vllm.distributed.initialize_model_parallel`` do not see stale state.
+    """
+    destroy_distributed_env()
+    import vllm.distributed.parallel_state as vllm_parallel_state
+
+    vllm_parallel_state._PP = None
+    for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
+        os.environ.pop(key, None)
+
+
 # ---------------------------------------------------------------------------
 # Shared stubs used by both unit and distributed tests
 # ---------------------------------------------------------------------------
@@ -143,7 +174,7 @@ class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
 class TestAsyncLatents:
     """Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""
 
-    pytestmark = [pytest.mark.cpu]
+    pytestmark = _UNIT_MARKS
 
     def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
         return AsyncLatents({"latents": tensor}, handles or [], postproc or [])
@@ -214,7 +245,7 @@ class TestAsyncLatents:
 class TestSyncPPSend:
     """Verifies PipelineParallelMixin's internal PP-send flush."""
 
-    pytestmark = [pytest.mark.cpu]
+    pytestmark = _UNIT_MARKS
 
     @staticmethod
     def _make_pipeline() -> PipelineParallelMixin:
@@ -247,7 +278,7 @@ class TestSyncPPSend:
 class TestDiffuseWrapper:
     """Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""
 
-    pytestmark = [pytest.mark.cpu]
+    pytestmark = _UNIT_MARKS
 
     def test_diffuse_flushes_pending_sends_on_success(self):
         work = FakeWork()
@@ -289,7 +320,7 @@ class TestDiffuseWrapper:
 
 
 class TestVaeDecodeGuard:
-    pytestmark = [pytest.mark.cpu]
+    pytestmark = _UNIT_MARKS
 
     @staticmethod
     def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
@@ -351,6 +382,8 @@ class TestVaeDecodeGuard:
         assert pipeline.vae.decode.__doc__ == "Original decode docstring."
 
 
+@pytest.mark.core_model
+@pytest.mark.diffusion
 @pytest.mark.cpu
 def test_pipeline_parallel_requires_cfg_mixin():
     with pytest.raises(TypeError, match="inherits PipelineParallelMixin but not CFGParallelMixin"):
@@ -359,6 +392,8 @@ def test_pipeline_parallel_requires_cfg_mixin():
             pass
 
 
+@pytest.mark.core_model
+@pytest.mark.diffusion
 @pytest.mark.cpu
 def test_pipeline_parallel_requires_mro_before_cfg_mixin():
     with pytest.raises(TypeError, match="must inherit PipelineParallelMixin before CFGParallelMixin"):
@@ -372,10 +407,22 @@ def test_pipeline_parallel_requires_mro_before_cfg_mixin():
 # ---------------------------------------------------------------------------
 
 
-def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
+def _worker_device(local_rank: int, device_kind: DeviceKind) -> torch.device:
+    if device_kind == "cpu":
+        return torch.device("cpu")
+    return torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+
+
+def init_dist(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    device_kind: DeviceKind,
+) -> torch.device:
     """Initialise the distributed environment for a spawned worker."""
-    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
-    current_omni_platform.set_device(device)
+    device = _worker_device(local_rank, device_kind)
+    if device_kind == "cuda":
+        current_omni_platform.set_device(device)
     update_environment_variables(
         {
             "RANK": str(local_rank),
@@ -385,13 +432,31 @@ def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.devic
             "MASTER_PORT": master_port,
         }
     )
-    init_distributed_environment()
+    backend = "gloo" if device_kind == "cpu" else None
+    init_distributed_environment(backend=backend)
     return device
 
 
+def _mp_backend(device_kind: DeviceKind) -> str | None:
+    """Backend for model-parallel process groups.
+
+    ``GroupCoordinator.all_gather`` always uses ``device_group``. CPU tensors
+    therefore require gloo device groups; otherwise a CUDA host's default NCCL
+    backend breaks PP+CFG CPU tests at CFG all_gather (pp2-cfg2).
+    isend/irecv of CPU tensors can already fall back to ``cpu_group``.
+    """
+    return "gloo" if device_kind == "cpu" else None
+
+
+def _seed_device(input_seed: int, device: torch.device) -> None:
+    torch.manual_seed(input_seed)
+    if device.type == "cuda":
+        torch.cuda.manual_seed_all(input_seed)
+
+
 def make_pipeline_and_inputs(
     test_config: dict, dtype: torch.dtype, device: torch.device, do_true_cfg: bool = False
-) -> tuple["MockPipelineParallel", dict, dict | None]:
+) -> tuple[MockPipelineParallel, dict, dict | None]:
     """Create a MockPipelineParallel and seeded inputs from a test_config dict.
 
     Must be called after ``initialize_model_parallel`` so that ``make_layers``
@@ -408,16 +473,12 @@ def make_pipeline_and_inputs(
         dtype=dtype,
     )
 
-    torch.manual_seed(test_config["input_seed"])
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(test_config["input_seed"])
+    _seed_device(test_config["input_seed"], device)
     pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)}
 
     negative_kwargs = None
     if do_true_cfg:
-        torch.manual_seed(test_config["input_seed"] + 1)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
+        _seed_device(test_config["input_seed"] + 1, device)
         neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
         negative_kwargs = {"x": neg_x}
 
@@ -425,19 +486,23 @@ def make_pipeline_and_inputs(
 
 
 # ---------------------------------------------------------------------------
-# 3.  isend_tensor_dict / irecv_tensor_dict  (2 GPUs)
+# 3.  isend_tensor_dict / irecv_tensor_dict
 # ---------------------------------------------------------------------------
 
 
-def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, result_queue):
-    device = init_dist(local_rank, world_size, master_port)
-    initialize_model_parallel(pipeline_parallel_size=world_size)
+def isend_irecv_worker(
+    local_rank: int,
+    world_size: int,
+    master_port: str,
+    device_kind: DeviceKind,
+    result_queue,
+):
+    device = init_dist(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(pipeline_parallel_size=world_size, backend=_mp_backend(device_kind))
     pp_group = get_pp_group()
 
     if pp_group.is_first_rank:
-        torch.manual_seed(77)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(77)
+        _seed_device(77, device)
         tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
         handles = pp_group.isend_tensor_dict({"t": tensor})
         for h in handles:
@@ -449,27 +514,42 @@ def isend_irecv_worker(local_rank: int, world_size: int, master_port: str, resul
 
     if torch.distributed.is_initialized():
         torch.distributed.barrier()
-    destroy_distributed_env()
+    _cleanup_distributed()
 
 
-@pytest.mark.gpu
-@pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs")
-@pytest.mark.parametrize("pp_size", [2])
-def test_isend_irecv_tensor_dict(pp_size: int):
-    """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
+def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) -> None:
     mp_context = torch.multiprocessing.get_context("spawn")
     manager = mp_context.Manager()
     q = manager.Queue()
-
-    port = _find_free_port()
-    torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, q), nprocs=pp_size)
-
+    torch.multiprocessing.spawn(
+        isend_irecv_worker,
+        args=(pp_size, master_port, device_kind, q),
+        nprocs=pp_size,
+    )
     results = {label: tensor for label, tensor in [q.get(), q.get()]}
     torch.testing.assert_close(
         results["received"], results["sent"], rtol=0, atol=0, msg="isend/irecv transferred tensor incorrectly"
     )
 
 
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize("pp_size", [2])
+def test_isend_irecv_tensor_dict(pp_size: int):
+    """isend_tensor_dict / irecv_tensor_dict transfer a tensor dict without loss."""
+    _run_isend_irecv(pp_size, device_kind="cpu", master_port=_find_free_port())
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize("pp_size", [pytest.param(2, marks=_L4_TWO_GPU)])
+def test_isend_irecv_tensor_dict_parity(pp_size: int):
+    """Nightly: isend/irecv on real multi-GPU NCCL."""
+    _run_isend_irecv(pp_size, device_kind="cuda", master_port=_find_free_port())
+
+
 # ---------------------------------------------------------------------------
 # 4.  predict_noise_maybe_with_cfg
 # ---------------------------------------------------------------------------
@@ -477,12 +557,13 @@ def test_isend_irecv_tensor_dict(pp_size: int):
 _baseline_cache: dict[tuple, torch.Tensor] = {}
 
 
-def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_cfg: bool) -> torch.Tensor:
-    """Compute expected single-GPU output using the same MockPipelineParallel.
-
-    Initializes a trivial distributed env (world_size=1) so that ``make_layers`` and the PP/CFG mixins work normally.
-    Results are cached so identical configs are only computed once.
-    """
+def compute_single_gpu_baseline(
+    test_config: dict,
+    dtype: torch.dtype,
+    do_true_cfg: bool,
+    device_kind: DeviceKind,
+) -> torch.Tensor:
+    """Compute expected single-process output using the same MockPipelineParallel."""
     key = (
         test_config["num_layers"],
         test_config["dim"],
@@ -492,27 +573,29 @@ def compute_single_gpu_baseline(test_config: dict, dtype: torch.dtype, do_true_c
         test_config["cfg_scale"],
         dtype,
         do_true_cfg,
+        device_kind,
     )
     if key in _baseline_cache:
         return _baseline_cache[key]
 
-    device = init_dist(0, 1, _find_free_port())
-    initialize_model_parallel(pipeline_parallel_size=1)
+    device = init_dist(0, 1, _find_free_port(), device_kind)
+    try:
+        initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
 
-    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
-        test_config, dtype, device, do_true_cfg=do_true_cfg
-    )
-
-    with torch.inference_mode():
-        noise_pred = pipeline.predict_noise_maybe_with_cfg(
-            do_true_cfg=do_true_cfg,
-            true_cfg_scale=test_config["cfg_scale"],
-            positive_kwargs=positive_kwargs,
-            negative_kwargs=negative_kwargs,
-            cfg_normalize=False,
+        pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+            test_config, dtype, device, do_true_cfg=do_true_cfg
         )
 
-    destroy_distributed_env()
+        with torch.inference_mode():
+            noise_pred = pipeline.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg,
+                true_cfg_scale=test_config["cfg_scale"],
+                positive_kwargs=positive_kwargs,
+                negative_kwargs=negative_kwargs,
+                cfg_normalize=False,
+            )
+    finally:
+        _cleanup_distributed()
 
     _baseline_cache[key] = noise_pred.cpu()
     return _baseline_cache[key]
@@ -527,11 +610,16 @@ def predict_noise_worker(
     do_true_cfg: bool,
     dtype: torch.dtype,
     test_config: dict,
+    device_kind: DeviceKind,
     result_queue,
 ):
     """Generic predict-noise worker parameterized by PP and CFG topology."""
-    device = init_dist(local_rank, world_size, master_port)
-    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+    device = init_dist(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(
+        pipeline_parallel_size=pp_size,
+        cfg_parallel_size=cfg_size,
+        backend=_mp_backend(device_kind),
+    )
 
     pp_group = get_pp_group()
     cfg_rank = get_classifier_free_guidance_rank()
@@ -548,8 +636,6 @@ def predict_noise_worker(
             negative_kwargs=negative_kwargs,
             cfg_normalize=False,
         )
-    # This worker exercises predict_noise_maybe_with_cfg directly, bypassing diffuse().
-    # Flush the non-last PP rank's async send before barrier / process teardown.
     pipeline._sync_pp_send()
 
     if pp_group.is_last_rank:
@@ -561,77 +647,22 @@ def predict_noise_worker(
 
     if torch.distributed.is_initialized():
         torch.distributed.barrier()
-    destroy_distributed_env()
+    _cleanup_distributed()
 
 
-@pytest.mark.gpu
-@pytest.mark.parametrize(
-    "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
-    [
-        pytest.param(
-            2,
-            1,
-            False,
-            torch.float32,
-            4,
-            100,
-            1e-5,
-            1e-5,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
-            id="pp2-no_cfg-float32",
-        ),
-        pytest.param(
-            2,
-            1,
-            False,
-            torch.bfloat16,
-            4,
-            100,
-            1e-2,
-            1e-2,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
-            id="pp2-no_cfg-bfloat16",
-        ),
-        pytest.param(
-            2,
-            1,
-            True,
-            torch.bfloat16,
-            4,
-            100,
-            1e-2,
-            1e-2,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
-            id="pp2-seq_cfg-bfloat16",
-        ),
-        pytest.param(
-            2,
-            2,
-            True,
-            torch.bfloat16,
-            4,
-            100,
-            1e-2,
-            1e-2,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
-            id="pp2-cfg2-bfloat16",
-        ),
-        pytest.param(
-            3,
-            1,
-            False,
-            torch.bfloat16,
-            6,
-            100,
-            1e-2,
-            1e-2,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 3, reason="Need at least 3 GPUs"),
-            id="pp3-no_cfg-bfloat16",
-        ),
-    ],
-)
-def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
-    """predict_noise_maybe_with_cfg output matches the single-GPU baseline across PP / CFG topologies."""
+def _run_predict_noise(
+    *,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    dtype: torch.dtype,
+    num_layers: int,
+    input_seed: int,
+    rtol: float,
+    atol: float,
+    device_kind: DeviceKind,
+    master_port: str,
+) -> None:
     test_config = {
         "num_layers": num_layers,
         "dim": 64,
@@ -641,22 +672,20 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_
         "input_seed": input_seed,
     }
 
-    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)
+    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)
 
     mp_context = torch.multiprocessing.get_context("spawn")
     manager = mp_context.Manager()
     pp_q = manager.Queue()
 
     world_size = pp_size * cfg_size
-    port = _find_free_port()
     torch.multiprocessing.spawn(
         predict_noise_worker,
-        args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q),
+        args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
         nprocs=world_size,
     )
 
     pp_out = pp_q.get()
-
     assert baseline_out.shape == pp_out.shape
     torch.testing.assert_close(
         pp_out,
@@ -667,35 +696,157 @@ def test_predict_noise(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_
     )
 
 
+_CPU_PREDICT_NOISE_CASES = [
+    pytest.param(2, 1, False, 4, 100, 1e-5, 1e-5, id="pp2-no_cfg"),
+    pytest.param(2, 1, True, 4, 100, 1e-5, 1e-5, id="pp2-seq_cfg"),
+    pytest.param(2, 2, True, 4, 100, 1e-5, 1e-5, id="pp2-cfg2"),
+    pytest.param(3, 1, False, 6, 100, 1e-5, 1e-5, id="pp3-no_cfg"),
+]
+
+_GPU_PREDICT_NOISE_CASES = [
+    pytest.param(
+        2,
+        1,
+        False,
+        torch.float32,
+        4,
+        100,
+        1e-5,
+        1e-5,
+        marks=_L4_TWO_GPU,
+        id="pp2-no_cfg-float32",
+    ),
+    pytest.param(
+        2,
+        1,
+        False,
+        torch.bfloat16,
+        4,
+        100,
+        1e-2,
+        1e-2,
+        marks=_L4_TWO_GPU,
+        id="pp2-no_cfg-bfloat16",
+    ),
+    pytest.param(
+        2,
+        1,
+        True,
+        torch.bfloat16,
+        4,
+        100,
+        1e-2,
+        1e-2,
+        marks=_L4_TWO_GPU,
+        id="pp2-seq_cfg-bfloat16",
+    ),
+    pytest.param(
+        2,
+        2,
+        True,
+        torch.bfloat16,
+        4,
+        100,
+        1e-2,
+        1e-2,
+        marks=_L4_FOUR_GPU,
+        id="pp2-cfg2-bfloat16",
+    ),
+    pytest.param(
+        3,
+        1,
+        False,
+        torch.bfloat16,
+        6,
+        100,
+        1e-2,
+        1e-2,
+        marks=_L4_THREE_GPU,
+        id="pp3-no_cfg-bfloat16",
+    ),
+]
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, num_layers, input_seed, rtol, atol",
+    _CPU_PREDICT_NOISE_CASES,
+)
+def test_predict_noise(pp_size, cfg_size, do_true_cfg, num_layers, input_seed, rtol, atol):
+    """predict_noise_maybe_with_cfg output matches baseline across PP / CFG topologies."""
+    _run_predict_noise(
+        pp_size=pp_size,
+        cfg_size=cfg_size,
+        do_true_cfg=do_true_cfg,
+        dtype=torch.float32,
+        num_layers=num_layers,
+        input_seed=input_seed,
+        rtol=rtol,
+        atol=atol,
+        device_kind="cpu",
+        master_port=_find_free_port(),
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol",
+    _GPU_PREDICT_NOISE_CASES,
+)
+def test_predict_noise_parity(pp_size, cfg_size, do_true_cfg, dtype, num_layers, input_seed, rtol, atol):
+    """Nightly: predict_noise parity on real multi-GPU NCCL."""
+    _run_predict_noise(
+        pp_size=pp_size,
+        cfg_size=cfg_size,
+        do_true_cfg=do_true_cfg,
+        dtype=dtype,
+        num_layers=num_layers,
+        input_seed=input_seed,
+        rtol=rtol,
+        atol=atol,
+        device_kind="cuda",
+        master_port=_find_free_port(),
+    )
+
+
 # ---------------------------------------------------------------------------
 # 5.  scheduler_step_maybe_with_cfg
 # ---------------------------------------------------------------------------
 
 
-def compute_scheduler_step_baseline(test_config: dict, do_true_cfg: bool) -> torch.Tensor:
-    """Single-GPU reference: predict_noise + scheduler_step."""
-    device = init_dist(0, 1, _find_free_port())
-    initialize_model_parallel(pipeline_parallel_size=1)
+def compute_scheduler_step_baseline(
+    test_config: dict,
+    do_true_cfg: bool,
+    device_kind: DeviceKind,
+) -> torch.Tensor:
+    """Single-process reference: predict_noise + scheduler_step."""
+    device = init_dist(0, 1, _find_free_port(), device_kind)
+    try:
+        initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
 
-    pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
-        test_config, torch.float32, device, do_true_cfg=do_true_cfg
-    )
-    latents = positive_kwargs["x"]
-    t = torch.tensor(500, device=device)
-
-    with torch.inference_mode():
-        noise_pred = pipeline.predict_noise_maybe_with_cfg(
-            do_true_cfg=do_true_cfg,
-            true_cfg_scale=test_config["cfg_scale"],
-            positive_kwargs=positive_kwargs,
-            negative_kwargs=negative_kwargs,
-            cfg_normalize=False,
-        )
-        result = pipeline.scheduler_step_maybe_with_cfg(
-            noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+        pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
+            test_config, torch.float32, device, do_true_cfg=do_true_cfg
         )
+        latents = positive_kwargs["x"]
+        t = torch.tensor(500, device=device)
 
-    destroy_distributed_env()
+        with torch.inference_mode():
+            noise_pred = pipeline.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg,
+                true_cfg_scale=test_config["cfg_scale"],
+                positive_kwargs=positive_kwargs,
+                negative_kwargs=negative_kwargs,
+                cfg_normalize=False,
+            )
+            result = pipeline.scheduler_step_maybe_with_cfg(
+                noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
+            )
+    finally:
+        _cleanup_distributed()
     return result.cpu()
 
 
@@ -707,10 +858,15 @@ def scheduler_step_worker(
     cfg_size: int,
     do_true_cfg: bool,
     test_config: dict,
+    device_kind: DeviceKind,
     result_queue,
 ):
-    device = init_dist(local_rank, world_size, master_port)
-    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
+    device = init_dist(local_rank, world_size, master_port, device_kind)
+    initialize_model_parallel(
+        pipeline_parallel_size=pp_size,
+        cfg_parallel_size=cfg_size,
+        backend=_mp_backend(device_kind),
+    )
 
     pp_group = get_pp_group()
     cfg_rank = get_classifier_free_guidance_rank()
@@ -732,43 +888,25 @@ def scheduler_step_worker(
         latents = pipeline.scheduler_step_maybe_with_cfg(
             noise_pred=noise_pred, t=t, latents=latents, do_true_cfg=do_true_cfg
         )
-    # This worker exercises scheduler_step_maybe_with_cfg directly, bypassing diffuse().
-    # Flush the last PP rank's async latent send before barrier / process teardown.
     pipeline._sync_pp_send()
 
     if pp_group.is_first_rank and cfg_rank == 0:
-        resolved = latents.contiguous()
-        result_queue.put(resolved.cpu())
+        result_queue.put(latents.contiguous().cpu())
 
     if torch.distributed.is_initialized():
         torch.distributed.barrier()
-    destroy_distributed_env()
+    _cleanup_distributed()
 
 
-@pytest.mark.gpu
-@pytest.mark.parametrize(
-    "pp_size, cfg_size, do_true_cfg, input_seed",
-    [
-        pytest.param(
-            2,
-            1,
-            False,
-            300,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 2, reason="Need at least 2 GPUs"),
-            id="pp2-no_cfg",
-        ),
-        pytest.param(
-            2,
-            2,
-            True,
-            600,
-            marks=pytest.mark.skipif(current_omni_platform.get_device_count() < 4, reason="Need at least 4 GPUs"),
-            id="pp2-cfg2-true_cfg",
-        ),
-    ],
-)
-def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
-    """Rank 0 latents after scheduler_step match the single-GPU baseline across PP / CFG topologies."""
+def _run_scheduler_step(
+    *,
+    pp_size: int,
+    cfg_size: int,
+    do_true_cfg: bool,
+    input_seed: int,
+    device_kind: DeviceKind,
+    master_port: str,
+) -> None:
     test_config = {
         "num_layers": 4,
         "dim": 64,
@@ -778,17 +916,16 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
         "input_seed": input_seed,
     }
 
-    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)
+    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)
 
     mp_context = torch.multiprocessing.get_context("spawn")
     manager = mp_context.Manager()
     q = manager.Queue()
 
-    port = _find_free_port()
     world_size = pp_size * cfg_size
     torch.multiprocessing.spawn(
         scheduler_step_worker,
-        args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q),
+        args=(world_size, master_port, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
         nprocs=world_size,
     )
 
@@ -800,3 +937,47 @@ def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
         atol=0,
         msg=f"PP={pp_size} CFG={cfg_size} scheduler step latents on rank 0 do not match single-GPU baseline",
     )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, input_seed",
+    [
+        pytest.param(2, 1, False, 300, id="pp2-no_cfg"),
+        pytest.param(2, 2, True, 600, id="pp2-cfg2-true_cfg"),
+    ],
+)
+def test_scheduler_step(pp_size, cfg_size, do_true_cfg, input_seed):
+    """Rank 0 latents after scheduler_step match baseline across PP / CFG topologies."""
+    _run_scheduler_step(
+        pp_size=pp_size,
+        cfg_size=cfg_size,
+        do_true_cfg=do_true_cfg,
+        input_seed=input_seed,
+        device_kind="cpu",
+        master_port=_find_free_port(),
+    )
+
+
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@pytest.mark.parametrize(
+    "pp_size, cfg_size, do_true_cfg, input_seed",
+    [
+        pytest.param(2, 1, False, 300, marks=_L4_TWO_GPU, id="pp2-no_cfg"),
+        pytest.param(2, 2, True, 600, marks=_L4_FOUR_GPU, id="pp2-cfg2-true_cfg"),
+    ],
+)
+def test_scheduler_step_parity(pp_size, cfg_size, do_true_cfg, input_seed):
+    """Nightly: scheduler_step parity on real multi-GPU NCCL."""
+    _run_scheduler_step(
+        pp_size=pp_size,
+        cfg_size=cfg_size,
+        do_true_cfg=do_true_cfg,
+        input_seed=input_seed,
+        device_kind="cuda",
+        master_port=_find_free_port(),
+    )

tests/diffusion/distributed/test_tensor_parallel.py

@@ -1,16 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Z-Image end-to-end tests for diffusion parallelism.
+"""DiT tensor-parallel parity and perf checks for diffusion models.
 
-This file currently covers:
-- DiT tensor parallelism (TP=2) vs TP=1.
-- VAE patch parallelism (vae_patch_parallel_size=2) vs baseline on TP=2.
-
-Note: CUDA-only (>=2 GPUs). We use `enforce_eager=False` (default) to enable
-`torch.compile`.
+Compares TP=2 against TP=1 on output correctness, latency, and peak memory.
+Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
 """
 
+from __future__ import annotations
+
 import os
 import time
 from pathlib import Path
@@ -31,22 +29,15 @@ from vllm_omni.platforms import current_omni_platform
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 PROMPT = "a photo of a cat sitting on a laptop keyboard"
-
-
-def _get_zimage_model() -> str:
-    # Allow overriding the model for local/offline environments.
-    # Can be either a HuggingFace repo id or a local path.
-    return os.environ.get("VLLM_TEST_ZIMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo")
+DEFAULT_MODEL = "Tongyi-MAI/Z-Image-Turbo"
 
 
 def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
-    """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3]."""
     arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
     return torch.from_numpy(arr)
 
 
 def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
-    """Return (mean_abs_diff, p99_abs_diff) over RGB pixels in [0, 1]."""
     ta = _pil_to_float_rgb_tensor(a)
     tb = _pil_to_float_rgb_tensor(b)
     assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
@@ -71,7 +62,7 @@ def _extract_single_image(outputs) -> Image.Image:
     return images[0]
 
 
-def _run_zimage_generate(
+def _run_generate(
     *,
     tp_size: int,
     height: int,
@@ -79,8 +70,6 @@ def _run_zimage_generate(
     num_inference_steps: int,
     seed: int,
     enforce_eager: bool,
-    vae_use_tiling: bool = False,
-    vae_patch_parallel_size: int = 1,
     num_requests: int = 4,
 ) -> tuple[Image.Image, float, float]:
     if num_requests < 2:
@@ -91,25 +80,11 @@ def _run_zimage_generate(
     monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
     try:
-        # Each run needs a distinct DiffusionParallelConfig; use OmniRunner per call (not the
-        # parametrized omni_runner fixture, which is fixed per module).
         with OmniRunner(
-            _get_zimage_model(),
-            parallel_config=DiffusionParallelConfig(
-                tensor_parallel_size=tp_size,
-                vae_patch_parallel_size=vae_patch_parallel_size,
-            ),
+            DEFAULT_MODEL,
+            parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size),
             enforce_eager=enforce_eager,
-            vae_use_tiling=vae_use_tiling,
         ) as runner:
-            # NOTE: Omni closes itself when a generate() call is exhausted.
-            # To avoid measuring teardown time (process shutdown, memory cleanup),
-            # we measure the latency to produce *subsequent* outputs within a single
-            # generator run.
-            #
-            # This also serves as a warmup: the first output may include extra
-            # compilation/caching overhead, while later outputs are closer to
-            # steady-state inference.
             gen = runner.omni.generate(
                 [PROMPT] * num_requests,
                 OmniDiffusionSamplingParams(
@@ -134,14 +109,11 @@ def _run_zimage_generate(
                 per_request_times_s.append(t_now - t_prev)
                 t_prev = t_now
 
-            # Ensure the generator is fully consumed so it can clean up.
             for _ in gen:
                 pass
 
             median_time_s = float(np.median(per_request_times_s))
-
             peak_memory_mb = monitor.peak_used_mb
-
             return _extract_single_image([last_output]), median_time_s, peak_memory_mb
     finally:
         monitor.stop()
@@ -151,20 +123,20 @@ def _run_zimage_generate(
 @pytest.mark.diffusion
 @pytest.mark.parallel
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
-def test_zimage_tensor_parallel_tp2(tmp_path: Path):
+def test_tensor_parallel_tp2(tmp_path: Path):
     if current_omni_platform.is_npu():
-        pytest.skip("Z-Image TP e2e test is only supported on CUDA and ROCm for now.")
+        pytest.skip("Tensor parallel parity test is only supported on CUDA and ROCm for now.")
     if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
-        pytest.skip("Z-Image TP=2 requires >= 2 devices.")
+        pytest.skip("Tensor parallel TP=2 requires >= 2 devices.")
 
+    model_name = DEFAULT_MODEL
     enforce_eager = False
-
     height = 512
     width = 512
     num_inference_steps = 2
     seed = 42
 
-    tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate(
+    tp1_img, tp1_time_s, tp1_peak_mem = _run_generate(
         tp_size=1,
         height=height,
         width=width,
@@ -172,7 +144,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
         seed=seed,
         enforce_eager=enforce_eager,
     )
-    tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate(
+    tp2_img, tp2_time_s, tp2_peak_mem = _run_generate(
         tp_size=2,
         height=height,
         width=width,
@@ -181,8 +153,8 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
         enforce_eager=enforce_eager,
     )
 
-    tp1_path = tmp_path / "zimage_tp1.png"
-    tp2_path = tmp_path / "zimage_tp2.png"
+    tp1_path = tmp_path / "tp1.png"
+    tp2_path = tmp_path / "tp2.png"
     tp1_img.save(tp1_path)
     tp2_img.save(tp2_path)
 
@@ -193,7 +165,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
     mean_threshold = 3e-2
     p99_threshold = 2.5e-1
     print(
-        "Z-Image TP image diff stats (TP=1 vs TP=2): "
+        f"{model_name} TP image diff (TP=1 vs TP=2): "
         f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
         f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
         f"tp1_img={tp1_path}, tp2_img={tp2_path}"
@@ -203,73 +175,11 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path):
         f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
     )
 
-    print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
-    # ROCm is not optimized TP2 can be slower than TP1
+    print(f"TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
     if not current_omni_platform.is_rocm():
         assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})"
 
-    print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
+    print(f"TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
     assert tp2_peak_mem < tp1_peak_mem, (
         f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})"
     )
-
-
-@pytest.mark.advanced_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
-def test_zimage_vae_patch_parallel_tp2(tmp_path: Path):
-    if current_omni_platform.is_npu():
-        pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA and ROCm for now.")
-    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
-        pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.")
-
-    enforce_eager = False
-
-    # Use a larger image to ensure there are multiple VAE tiles.
-    height = 1152
-    width = 1152
-    num_inference_steps = 2
-    seed = 42
-
-    baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate(
-        tp_size=2,
-        height=height,
-        width=width,
-        num_inference_steps=num_inference_steps,
-        seed=seed,
-        enforce_eager=enforce_eager,
-        vae_use_tiling=True,
-        vae_patch_parallel_size=1,
-        num_requests=2,
-    )
-    pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate(
-        tp_size=2,
-        height=height,
-        width=width,
-        num_inference_steps=num_inference_steps,
-        seed=seed,
-        enforce_eager=enforce_eager,
-        vae_use_tiling=True,
-        vae_patch_parallel_size=2,
-        num_requests=2,
-    )
-
-    baseline_path = tmp_path / "zimage_tp2_vae_pp1.png"
-    pp2_path = tmp_path / "zimage_tp2_vae_pp2.png"
-    baseline_img.save(baseline_path)
-    pp2_img.save(pp2_path)
-
-    mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img)
-    mean_threshold = 5e-3
-    p99_threshold = 1e-1
-    print(
-        "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): "
-        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
-        f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
-        f"pp1_img={baseline_path}, pp2_img={pp2_path}"
-    )
-    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
-        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
-        f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
-    )

tests/diffusion/distributed/test_vae_decode_parallelism.py

@@ -1,151 +1,181 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""VAE patch-parallel decode parity tests for diffusion models.
+
+Compares ``vae_patch_parallel_size=2`` against ``vae_patch_parallel_size=1`` on
+TP=2 with VAE tiling enabled. Covers one video model (Wan T2V) and one image
+model (Qwen-Image). CUDA/ROCm-only (>=2 devices).
+"""
+
+from __future__ import annotations
+
 import os
-import sys
 from pathlib import Path
+from typing import Any
 
 import numpy as np
 import pytest
 import torch
 from PIL import Image
-from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# ruff: noqa: E402
-REPO_ROOT = Path(__file__).resolve().parents[2]
-if str(REPO_ROOT) not in sys.path:
-    sys.path.insert(0, str(REPO_ROOT))
-
-import time
-
-from tests.helpers.runtime import OmniRunner
 from vllm_omni.platforms import current_omni_platform
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
-all_model_configs = [
-    {
-        "model_name": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
-        "out_width": 1280,
-        "out_height": 720,
-        "out_frames": 16,
-    },
-    {
-        "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-        "out_width": 1280,
-        "out_height": 720,
-        "out_frames": 16,
-    },
-    {
-        "model_name": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-        "out_width": 1280,
-        "out_height": 720,
-        "out_frames": 16,
-    },
-    {
-        "model_name": "Qwen/Qwen-Image",
-        "out_width": 3072,
-        "out_height": 3072,
-        "out_frames": 1,
-    },
-    # {
-    #     "model_name": "stepfun-ai/NextStep-1.1",
-    #     "out_width": 3072,
-    #     "out_height": 3072,
-    #     "out_frames": 1,
-    # },
+PROMPT = "A cat sitting on a table"
+
+MODEL_CASES = [
+    pytest.param(
+        {
+            "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+            "height": 512,
+            "width": 512,
+            "num_frames": 5,
+            "needs_image": False,
+            "is_moe": False,
+            "mean_threshold": 3e-2,
+            "p99_threshold": 3e-2,
+        },
+        id="wan22_t2v_a14b",
+    ),
+    pytest.param(
+        {
+            "model_name": "Qwen/Qwen-Image",
+            "height": 1152,
+            "width": 1152,
+            "num_frames": 1,
+            "needs_image": False,
+            "is_moe": False,
+            "mean_threshold": 5e-3,
+            "p99_threshold": 1e-1,
+        },
+        id="qwen_image",
+    ),
 ]
 
 
-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 _to_float_array(data: Any) -> np.ndarray:
+    if isinstance(data, Image.Image):
+        return np.asarray(data.convert("RGB"), dtype=np.float32) / 255.0
+    if isinstance(data, torch.Tensor):
+        arr = data.detach().cpu().numpy()
+    else:
+        arr = np.asarray(data)
+    if arr.dtype == np.uint8:
+        return arr.astype(np.float32) / 255.0
+    return arr.astype(np.float32)
 
 
-def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
+def _extract_output_array(output: Any) -> np.ndarray:
+    req_out = output.request_output
+    payload = req_out.images[0]
+    return _to_float_array(payload)
+
+
+def _diff_metrics(a: np.ndarray, b: np.ndarray) -> tuple[float, float]:
+    assert a.shape == b.shape, f"Output shapes differ: {a.shape} vs {b.shape}"
+    abs_diff = np.abs(a - b)
+    return float(abs_diff.mean()), float(np.quantile(abs_diff, 0.99))
+
+
+def _save_preview(arr: np.ndarray, path: Path) -> None:
+    """Save a single RGB frame preview; video tensors use the first frame."""
+    preview = np.asarray(arr)
+    # (B, T, H, W, C) video
+    if preview.ndim == 5:
+        preview = preview[0, 0]
+    # (T, H, W, C) video or (B, H, W, C) image batch
+    elif preview.ndim == 4:
+        preview = preview[0]
+    if preview.ndim != 3 or preview.shape[-1] not in (3, 4):
+        raise ValueError(f"Expected HWC RGB(A) preview, got shape {arr.shape} -> {preview.shape}")
+    preview = np.clip(preview[..., :3], 0.0, 1.0)
+    Image.fromarray((preview * 255.0).astype(np.uint8)).save(path)
+
+
+def _build_sampling_params(model_case: dict[str, Any], seed: int) -> OmniDiffusionSamplingParams:
+    kwargs: dict[str, Any] = {
+        "height": model_case["height"],
+        "width": model_case["width"],
+        "num_frames": model_case["num_frames"],
+        "num_inference_steps": 2,
+        "seed": seed,
+    }
+    if model_case["model_name"].startswith("Wan-AI/"):
+        kwargs["guidance_scale"] = 4.0
+    if model_case["is_moe"]:
+        kwargs["guidance_scale_2"] = 1.0
+        kwargs["boundary_ratio"] = 0.5
+    return OmniDiffusionSamplingParams(**kwargs)
+
+
+def _run_generate(
+    *,
+    model_case: dict[str, Any],
+    vae_patch_parallel_size: int,
+    seed: int,
+) -> np.ndarray:
+    current_omni_platform.empty_cache()
     parallel_config = DiffusionParallelConfig(
-        tensor_parallel_size=tp,
+        tensor_parallel_size=2,
         vae_patch_parallel_size=vae_patch_parallel_size,
     )
-
-    omni_kwargs = {
-        "vae_use_tiling": using_tile,
-        "parallel_config": parallel_config,
-    }
-    use_nextstep = is_nextstep_model(model_configs["model_name"])
-    if use_nextstep:
-        # NextStep-1.1 requires explicit pipeline class
-        omni_kwargs["model_class_name"] = "NextStep11Pipeline"
-
-    with OmniRunner(model_configs["model_name"], **omni_kwargs) as runner:
-        m = runner.omni
-        image = Image.new("RGB", (out_width, out_height), (0, 0, 0))
-        start = time.perf_counter()
-        outputs = m.generate(
-            {
-                "prompt": "A cat sitting on a table",
-                "multi_modal_data": {"image": image},
-            },
-            sampling_params_list=OmniDiffusionSamplingParams(
-                height=model_configs["out_height"],
-                width=model_configs["out_width"],
-                num_frames=model_configs["out_frames"],
-                num_inference_steps=2,
-                generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
-            ),
+    with OmniRunner(
+        model_case["model_name"],
+        parallel_config=parallel_config,
+        vae_use_tiling=True,
+        enforce_eager=False,
+    ) as runner:
+        request: dict[str, Any] = {"prompt": model_case.get("prompt", PROMPT)}
+        if model_case["needs_image"]:
+            request["multi_modal_data"] = {
+                "image": Image.new("RGB", (model_case["width"], model_case["height"]), (0, 0, 0))
+            }
+        outputs = runner.omni.generate(
+            request,
+            sampling_params_list=_build_sampling_params(model_case, seed),
         )
-        end = time.perf_counter()
-        first_output = outputs[0]
-        req_out = first_output.request_output
-        frames = req_out.images[0]
-        if isinstance(frames, torch.Tensor):
-            frames = frames.detach().cpu().numpy()
-        elif isinstance(frames, Image.Image):
-            frames = np.array(frames)
-        # frames shape: (batch, num_frames, height, width, channels)
-        cost = (end - start) * 1000
-        return frames, cost
-        cleanup_dist_env_and_memory()
+        return _extract_output_array(outputs[0])
 
 
-@pytest.mark.parametrize("model_configs", all_model_configs)
-def test_vae_parallel_model(model_configs: str):
-    parallel_result, parallel_time = model_run(
-        model_configs=model_configs,
-        tp=2,
-        out_width=1280,
-        out_height=704,
-        out_frames=5,
-        using_tile=True,
-        vae_patch_parallel_size=2,
-    )
-    non_parallel_result, non_parallel_time = model_run(
-        model_configs=model_configs,
-        tp=2,
-        out_width=1280,
-        out_height=704,
-        out_frames=5,
-        using_tile=True,
-        vae_patch_parallel_size=1,
-    )
-    result_diff = np.abs(non_parallel_result - parallel_result)
+@pytest.mark.full_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
+@pytest.mark.parametrize("model_case", MODEL_CASES)
+def test_vae_patch_parallel_tp2(model_case: dict[str, Any], tmp_path: Path):
+    if current_omni_platform.is_npu():
+        pytest.skip("VAE patch parallel test is only supported on CUDA and ROCm for now.")
+    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
+        pytest.skip("VAE patch parallel TP=2 requires >= 2 devices.")
 
-    mean_threshold = 3e-2
-    max_threshold = 3e-2  # they should be totally same
+    seed = 42
+    model_id = model_case["model_name"].split("/")[-1]
+
+    baseline = _run_generate(model_case=model_case, vae_patch_parallel_size=1, seed=seed)
+    parallel = _run_generate(model_case=model_case, vae_patch_parallel_size=2, seed=seed)
+
+    baseline_path = tmp_path / f"{model_id}_tp2_vae_pp1.png"
+    parallel_path = tmp_path / f"{model_id}_tp2_vae_pp2.png"
+    _save_preview(baseline, baseline_path)
+    _save_preview(parallel, parallel_path)
+
+    mean_abs_diff, p99_abs_diff = _diff_metrics(baseline, parallel)
+    mean_threshold = model_case["mean_threshold"]
+    p99_threshold = model_case["p99_threshold"]
     print(
-        f"{model_configs['model_name']} TP = 2 (tile + parallel vs tile): "
-        f"mean_abs_diff={result_diff.mean():.6e}, max_abs_diff={result_diff.max():.6e}; "
-        f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; "
-        f"parallel generate take time: {parallel_time:.2f} ms, non-parallel take time: {non_parallel_time:.2f} ms"
+        f"{model_case['model_name']} VAE patch parallel diff (TP=2, pp=1 vs pp=2): "
+        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
+        f"pp1_preview={baseline_path}, pp2_preview={parallel_path}"
+    )
+    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
+        f"Output diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, "
+        f"p99_abs_diff={p99_abs_diff:.6e} "
+        f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
     )
-    # assert non_parallel_time > parallel_time
-    assert result_diff.mean() < mean_threshold
-    assert result_diff.max() < max_threshold

tests/diffusion/hooks/test_hook_registry.py

@@ -12,6 +12,9 @@ from torch import nn
 
 from vllm_omni.diffusion.hooks.base import HookRegistry, ModelHook
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 DEFAULT_OUT = "ECHO"
 OVERRIDE_OUT = "OVERRIDE"
 INPUT_KWARG = "inp"

tests/diffusion/layers/test_rotary_emb_equivalence.py

@@ -13,6 +13,8 @@ from __future__ import annotations
 import pytest
 import torch
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def _apply_rotary_emb_helios_original(
     hidden_states: torch.Tensor,

tests/diffusion/models/bagel/test_negative_prompt.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+from pytest_mock import MockerFixture
+
+from vllm_omni.diffusion.models.bagel.pipeline_bagel import BagelPipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _NegativePromptPreparedError(Exception):
+    pass
+
+
+@pytest.mark.parametrize(
+    ("prompt_negative", "extra_negative", "expected"),
+    [
+        ("prompt negative", "legacy negative", "prompt negative"),
+        (None, "legacy negative", "legacy negative"),
+    ],
+    ids=["prompt-owner", "legacy-extra-args"],
+)
+def test_forward_reads_negative_prompt_from_prompt_with_legacy_fallback(
+    mocker: MockerFixture,
+    prompt_negative: str | None,
+    extra_negative: str,
+    expected: str,
+) -> None:
+    pipeline = object.__new__(BagelPipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+    pipeline.od_config = mocker.Mock(dtype=torch.float32)
+    pipeline.tokenizer = mocker.sentinel.tokenizer
+    pipeline.new_token_ids = {}
+    pipeline.language_model = mocker.Mock(vocab_size=10)
+    pipeline.image_processor = None
+    pipeline.vae = None
+
+    bagel = mocker.MagicMock()
+    bagel.max_latent_size = 32
+    bagel.latent_downsample = 8
+    bagel.config.llm_config.num_hidden_layers = 1
+    prepared_prompts: list[str] = []
+
+    def prepare_prompts(*, prompts, **_):
+        prepared_prompts.append(prompts[0])
+        if len(prepared_prompts) == 1:
+            return {"packed_text_ids": torch.tensor([1])}, [1], [1]
+        raise _NegativePromptPreparedError
+
+    bagel.prepare_prompts.side_effect = prepare_prompts
+    pipeline.bagel = bagel
+
+    prompt = {
+        "prompt": "positive",
+        "negative_prompt": prompt_negative,
+    }
+    request = DiffusionRequestBatch(
+        requests=[
+            OmniDiffusionRequest(
+                prompt=prompt,
+                sampling_params=OmniDiffusionSamplingParams(
+                    extra_args={"negative_prompt": extra_negative},
+                ),
+                request_id="test",
+            )
+        ]
+    )
+
+    with pytest.raises(_NegativePromptPreparedError):
+        pipeline.forward(request)
+
+    assert prepared_prompts == ["positive", expected]