Revert "[BugFix] Fix 'GPUARWorker' attribute error (#5478)"

This reverts commit f4640d9a2e.

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

@@ -89,7 +89,7 @@ steps:
       mirror_hardwares: [amdproduction]
       grade: Blocking
       commands:
-        - pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py
+        - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.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/features/custom_pipeline/ -m "core_model"
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
 
 
 - label: "Voxtral TTS CUDA Unit Test"

.buildkite/cuda/test-merge.yml

@@ -8,24 +8,14 @@ steps:
   - group: ":card_index_dividers: Simple Test"
     depends_on: upload-merge-pipeline
     steps:
-      - label: "Simple · Model Executor Test"
+      - label: "Simple · Diffusion & Model Executor Test"
         commands:
-          - "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"
+          - "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"
         mirror_hardwares: l4_1
 
       - label: "Simple · Other Test"
         commands:
-          - "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"
+          - "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"
         mirror_hardwares: l4_1
 
   - group: ":card_index_dividers: Diffusion Test"
@@ -37,6 +27,12 @@ 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
@@ -139,32 +135,39 @@ 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'
-        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
+        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
 
       - label: "TTS · VoxCPM2 Test"
         source_file_dependencies:
@@ -182,6 +185,22 @@ 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,13 +15,7 @@ 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 --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"
+          - 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
         mirror_hardwares: h100_2
 
       - label: ":full_moon: Omni · Doc Test with L4"
@@ -455,41 +449,36 @@ 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"
+    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 with L4"
-        timeout_in_minutes: 60
-        commands:
-          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4' --run-level "full_model"
-        mirror_hardwares: l4_4
-
-      - label: ":full_moon: Diffusion · Distributed Test with H100"
-        timeout_in_minutes: 60
-        commands:
-          - pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100' --run-level "full_model"
-        mirror_hardwares: h100_4
 
       - label: ":full_moon: Diffusion · Offloader Test"
-        timeout_in_minutes: 60
+        timeout_in_minutes: 40
         commands:
-          - pytest -sv tests/diffusion/offloader -m "full_model and cuda"
+          - timeout 40m pytest -s -v 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:
-          - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
+          - timeout 60m 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,26 +8,22 @@ steps:
   - group: ":card_index_dividers: Simple Test"
     depends_on: upload-ready-pipeline
     steps:
-      - label: "Simple · Model Executor Test"
+      - label: "Simple · Diffusion & Model Executor Test"
         commands:
-          - "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"
+          - "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"
         mirror_hardwares: l4_1
 
       - label: "Simple · Other Test"
         commands:
-          - "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"
+          - "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"
         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:
@@ -108,34 +104,11 @@ 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/features/rlhf_test/test_verl_omni_e2e.py
-          - tests/e2e/features/helpers/verl_omni_server.py
-          - tests/e2e/features/helpers/custom_pipeline.py
+          - tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
+          - tests/e2e/offline_inference/custom_pipeline/
           - vllm_omni/diffusion/
           - vllm_omni/engine/
           - vllm_omni/entrypoints/
@@ -144,7 +117,7 @@ steps:
         commands:
           - |
             timeout 20m bash -c "
-              pytest -s -v tests/e2e/features/rlhf_test/test_verl_omni_e2e.py
+              pytest -s -v tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
             "
         mirror_hardwares: l4_1
 
@@ -170,7 +143,34 @@ steps:
             timeout 20m bash -c "
               pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py -m 'core_model' --run-level 'core_model'
             "
-        mirror_hardwares: h100_2
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
 
       - label: "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'
             "
-        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
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
 
       - label: "TTS · VoxCPM2 Test"
         source_file_dependencies:
@@ -344,3 +344,13 @@ 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.26.0"
+          VLLM_VERSION: "v0.25.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: false

.buildkite/npu/test-npu-nightly.yml

@@ -95,80 +95,6 @@ steps:
             buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
             buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
             exit $$EXIT
-  - group: ":card_index_dividers: Diffusion I2I Model Test"
-    key: nightly-HunyuanImage3-x2i-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: HunyuanImage3-DIT · Accuracy Test"
-        timeout_in_minutes: 180
-        mirror_hardwares: a3_npu_4
-        env:
-          VLLM_WORKER_MULTIPROC_METHOD: spawn
-          HUNYUAN_IMAGE3_MODEL: "tencent/HunyuanImage-3.0-Instruct"
-          HUNYUAN_IMAGE3_DEVICES: "0,1,2,3"
-          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
-        commands:
-          - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and A3" --run-level full_model
-  - group: ":card_index_dividers: Diffusion Benchmark Perf Test"
-    key: nightly-diffusion-benchmark-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 Perf · HunyuanImage3 TP=2 CFGP=2"
-        timeout_in_minutes: 180
-        mirror_hardwares: a3_npu_4
-        env:
-          VLLM_WORKER_MULTIPROC_METHOD: spawn
-          HF_TOKEN: "${HF_TOKEN}"
-          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
-        commands:
-          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
-            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$EXIT
-      - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 SP=2"
-        timeout_in_minutes: 180
-        mirror_hardwares: a3_npu_4
-        env:
-          VLLM_WORKER_MULTIPROC_METHOD: spawn
-          HF_TOKEN: "${HF_TOKEN}"
-          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
-        commands:
-          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
-            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$EXIT
-      - label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=4"
-        timeout_in_minutes: 180
-        mirror_hardwares: a3_npu_4
-        env:
-          VLLM_WORKER_MULTIPROC_METHOD: spawn
-          HF_TOKEN: "${HF_TOKEN}"
-          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
-        commands:
-          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp4.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
-            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
-            exit $$EXIT
   - group: ":card_index_dividers: Diffusion X2V Model Test"
     key: nightly-diffusion-x2v-group
     depends_on: upload-nightly-pipeline
@@ -177,14 +103,6 @@ steps:
       build.pull_request.labels includes "nightly-test" ||
       build.pull_request.labels includes "diffusion-x2v-test"
     steps:
-      - label: ":full_moon: Diffusion X2V · Function Test"
-        mirror_hardwares: a2b3_npu_4
-        env:
-          VLLM_WORKER_MULTIPROC_METHOD: spawn
-          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
-          HF_TOKEN: "${HF_TOKEN}"
-        commands:
-          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"
       - label: ":full_moon: Diffusion X2V · Perf Test"
         key: nightly-diffusion-x2v-performance
         timeout_in_minutes: 180

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

@@ -1,8 +1,12 @@
 from collections.abc import Callable
 from enum import Enum, auto
-from typing import Any, Literal, TypeAlias
-
-from typing_extensions import NotRequired, TypedDict
+from typing import (
+    Any,
+    Literal,
+    NotRequired,
+    TypeAlias,
+    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.26.0
+ARG VLLM_ASCEND_TAG=v0.25.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm
@@ -26,10 +26,6 @@ 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.26.0-a3
+ARG VLLM_ASCEND_TAG=v0.25.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 # WORKDIR /vllm-workspace/vllm
@@ -26,10 +26,6 @@ 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.26.0
+ARG VLLM_ASCEND_TAG=v0.25.0
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
@@ -35,7 +35,6 @@ 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.26.0-a3
+ARG VLLM_ASCEND_TAG=v0.25.0-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
@@ -35,7 +35,6 @@ 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" \
@@ -45,7 +44,6 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
       "hyperpyyaml>=1.2.0" \
       "onnx>=1.16.0" \
       "librosa>=0.10.0" \
-      "torchmetrics>=1.9.0" \
       --index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
       --trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
       --default-timeout=60 && \

docs/contributing/ci/test_system_overview.md

@@ -56,11 +56,7 @@ 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>
-        <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>/tests/{component_name}/test_xxx</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
@@ -72,18 +68,16 @@ 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 and component tests that require launching instances</td>
+      <td>High-priority features (using random lightweight models)</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 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>
+        <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
       </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>
@@ -96,19 +90,18 @@ 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 and component tests that require launching instances</td>
+      <td>Medium-priority features (using random lightweight models)</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 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>
+        <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
       </td>
       <td>
         <a href="test_writing_guide.md#l3-level-testing-core-integration-performance-and-accuracy-verification">L3</a><br>
@@ -121,23 +114,26 @@ 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 and component tests that require launching instances</td>
+      <td>Low-priority features (using real weights)</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 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>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>Performance:</strong><br>
-        <code>tests/dfx/perf/tests/</code><br>
-        <strong>Doc examples:</strong><br>
-        <code>tests/examples/online_serving/</code>, <code>tests/examples/offline_inference/</code>
+        /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
       </td>
       <td>
         <a href="test_writing_guide.md#l4-level-testing-full-functionality-performance-and-documentation-testing">L4</a><br>
@@ -156,9 +152,11 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
       <td> Depends on reality </td>
       <td>
         <strong>Stability:</strong><br>
-        <code>tests/dfx/stability/tests/</code><br>
+        /tests/dfx/stability/tests/test_qwen3_omni.json<br>
+        /tests/dfx/stability/tests/test_wan22.json<br>
         <strong>Reliability:</strong><br>
-        <code>tests/dfx/reliability/test_reliability_{model}.py</code>
+        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>)
       </td>
       <td>
         <a href="test_writing_guide.md#l5-level-testing-stability-and-reliability-testing">L5</a><br>
@@ -170,9 +168,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
   </tbody>
 </table>
 
-**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).
+For per-level test authoring (directories, markers, examples), see [Test Writing Guide](./test_writing_guide.md).
 
 ## Common Specifications
 

docs/contributing/ci/test_writing_guide.md

@@ -4,20 +4,137 @@ This document describes how to author tests for each CI level (L1–L5). For the
 
 ## Test Directory Structure
 
-Canonical placement by CI level is in the **Test Dir** column of the level matrix in [Test System Overview](./test_system_overview.md).
+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)
 
-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)).
+# Migrated models (qwen3_omni_moe, qwen2_5_omni, qwen3_tts) live under
+# vllm_omni/deploy/ instead — see docs/configuration/stage_configs.md.
+```
 
 ## Markers for Tests
 
@@ -256,13 +373,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. See the **Test Dir** column in [Test System Overview](./test_system_overview.md) for the canonical placement rules.
+A clear directory structure is key to managing test cases efficiently.
 
--   ***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/`).
+-   ***L1 Test Directory***: `/tests/{component_name}/test_xxx.py`
+-   -   Here, `{component_name}` corresponds to modules in the source code, such as `distributed`, `entrypoints`, etc., and `test_xxx.py` is the specific test file.
+-   ***L2 Test Directory***:
+-   -   Online Serving: `/tests/e2e/online_serving/test_{model_name}.py`
+    -   Offline Inference: `/tests/e2e/offline_inference/test_{model_name}.py`
 
 #### 1.4 Execution Method and Example
 
@@ -368,12 +485,10 @@ L3 level testing executes after code is merged into the main branch. Its core pu
 
 #### 2.3 Test Directory and Execution Files
 
--   ***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).
+-   ***Functional Testing***:
+-   -   Online Serving: `/tests/e2e/online_serving/test_{model_name}_expansion.py`
+    -   Offline Inference: `/tests/e2e/offline_inference/test_{model_name}_expansion.py`
+    -   (Note: `_expansion.py` likely means it contains more comprehensive scenario cases compared to L2 tests).
 
 #### 2.4 Execution Method and Example
 

docs/design/feature/async_diffusion_output.md

@@ -1,174 +0,0 @@
-# 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/diffusion/distributed/test_tensor_parallel.py` | E2E testing | TP correctness and performance |
+| **TP Tests** | `tests/e2e/offline_inference/test_zimage_parallelism.py` | E2E testing | TP correctness and performance |
 | **Constraint Tests** | `tests/diffusion/models/z_image/test_zimage_tp_constraints.py` | Unit testing | Validation logic |
 
 ---

docs/design/metrics.md

@@ -148,13 +148,7 @@ 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. 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.
+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.
 
 ## 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.26.0
-# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.26.0-a3
+# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.25.0
+# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.25.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:v0.26.0-a3 bash
+    -it quay.io/atlas-ci/vllm-ascend:0.25.0-a3 bash
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.26.0 https://github.com/vllm-project/vllm-omni.git
+git clone -b v0.25.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.26.0
-# Atlas A3: quay.io/ascend/vllm-omni:v0.26.0-a3
+# Atlas A2: quay.io/ascend/vllm-omni:v0.25.0
+# Atlas A3: quay.io/ascend/vllm-omni:v0.25.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.26.0-a3 bash
+    -it quay.io/ascend/vllm-omni:v0.25.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/models/supported_models.md

@@ -44,7 +44,6 @@ th {
 | `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
 | `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
-| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, FL2VA, and Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
 | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |

docs/serving/speech_api.md

@@ -276,42 +276,20 @@ Client -> Server:
 
 | Message | Description |
 |---------|-------------|
-| `{"type": "session.config", ...}` | Session configuration (first message; may be resent between utterances to change it) |
+| `{"type": "session.config", ...}` | Session configuration (sent once, first message) |
 | `{"type": "input.text", "text": "..."}` | Text chunk |
-| `{"type": "input.done"}` | End of utterance: flushes the buffer and keeps the connection open |
-| `{"type": "session.close"}` | End of connection |
+| `{"type": "input.done"}` | End of input, flushes remaining buffer |
 
 Server -> Client:
 
 | Message | Description |
 |---------|-------------|
-| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
+| `{"type": "audio.start", "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", "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": "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": "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

@@ -15,7 +15,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`. |
-| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
+| `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`. |
 | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
@@ -167,48 +167,6 @@ 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`.
@@ -235,17 +193,9 @@ DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/te
     --model Lightricks/LTX-2 ...
 ```
 
-### FlashAttention-4 on Blackwell
+### FA4 not yet integrated
 
-Install the optional CUDA 13 extra to use the CuTe-based FA4 path:
-
-```bash
-pip install 'vllm-omni[fa4]'
-```
-
-`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2
-when it is unavailable. Version `4.0.0b18` is required; earlier beta
-wheels had known JIT failures on Blackwell.
+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
 
@@ -403,3 +353,64 @@ 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,134 +179,7 @@ class Flux2Transformer2DModel(nn.Module):
 - Support single GPU only for now
 
 
-## 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
+### Implementation Notes
 
 **Module Discovery**
 
@@ -335,10 +208,7 @@ 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
-│                          (single-GPU, full weights on host)
-└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
-                                         (multi-GPU, 1/dp_size sharded weights + AllGather)
+└── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
 ```
 
 Factory function `get_offload_backend()` selects the appropriate backend based on
@@ -351,22 +221,20 @@ reasoner/generator components inside the model forward pass.
 
 ## Supported Models
 
-| 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"` |
+| 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"` |
 
 **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<-- "vllm_omni/deploy/qwen3_omni_moe_thinking.yaml"
+    --8<-- "examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml"
     ``````
 ??? abstract "run_curl_multimodal_generation.sh"
     ``````sh

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -422,19 +422,12 @@ 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 and synthesizes the buffered text as one continuous request on `input.done`:
+The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence:
 ```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,4 +8,5 @@ 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,7 +334,6 @@ 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
@@ -348,19 +347,7 @@ 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_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:
+    if 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,13 +182,10 @@ 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(), sample_rate)
-                    print(f"Saved audio to {out_path} (sr={sample_rate})")
+                    sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
+                    print(f"Saved audio to {out_path}")
             else:
                 print("No multimodal output found.")
         except Exception as e:

examples/offline_inference/text_to_video/text_to_video.py

@@ -56,19 +56,6 @@ _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,
@@ -116,10 +103,6 @@ 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,18 +23,17 @@ 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%, 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
+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
 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_2gpu.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. |
+| `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_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. |
@@ -90,18 +89,17 @@ python openai_chat_completion_client_for_multimodal_generation.py \
     --prompt "Briefly introduce yourself."
 ```
 
-Streaming text + audio (use `--stream`):
+Streaming text + audio:
 
 ```bash
-python openai_chat_completion_client_for_multimodal_generation.py \
-    --query-type text \
-    --prompt "Briefly introduce yourself." \
-    --port 8099 \
-    --stream
+python streaming_chat_completion.py \
+    --base-url http://localhost:8099/v1 \
+    --output minicpmo_stream.wav
 ```
 
-The client prints text deltas as they arrive and saves streamed audio chunks
-to WAV files.
+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.
 
 Shared helpers also work if you pass MiniCPM defaults yourself:
 

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

@@ -6,7 +6,6 @@ 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);
@@ -32,7 +31,6 @@ 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();
@@ -44,7 +42,6 @@ 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;
@@ -79,7 +76,6 @@ 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;
@@ -120,7 +116,6 @@ 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) {
@@ -128,16 +123,6 @@ 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

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

@@ -0,0 +1,36 @@
+# 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
-DEPLOY_CONFIG=""
+STAGE_CONFIGS_PATH=""
 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
             ;;
-        --deploy-config)
-            DEPLOY_CONFIG="$2"
+        --stage-configs-path)
+            STAGE_CONFIGS_PATH="$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 "  --deploy-config PATH          Path to custom deploy config YAML file (optional)"
+            echo "  --stage-configs-path PATH     Path to custom stage configs 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 "$DEPLOY_CONFIG" ]; then
-    SERVER_CMD+=("--deploy-config" "$DEPLOY_CONFIG")
+if [ -n "$STAGE_CONFIGS_PATH" ]; then
+    SERVER_CMD+=("--stage-configs-path" "$STAGE_CONFIGS_PATH")
 fi
 
 # Function to cleanup on exit

examples/online_serving/text_to_speech/README.md

@@ -537,19 +537,12 @@ 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 and synthesizes the buffered text as one continuous request on `input.done`:
+The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence:
 ```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,12 +7,6 @@ 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 "打开电子枪，关闭扫描，切换到低倍模式。" \
@@ -70,27 +64,18 @@ except ImportError:
     raise SystemExit(1)
 
 
-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:
+def save_audio_file(output_dir: str, sentence_index: int, response_format: str, chunks: list[bytes]) -> str:
     audio_bytes = b"".join(chunks)
     if response_format == "pcm":
-        raw_path = os.path.join(output_dir, f"{basename}.pcm")
+        raw_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.pcm")
         with open(raw_path, "wb") as f:
             f.write(audio_bytes)
 
-        wav_path = os.path.join(output_dir, f"{basename}.wav")
+        wav_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.wav")
         save_pcm_wav(wav_path, audio_bytes)
         return wav_path
 
-    filename = os.path.join(output_dir, f"{basename}.{response_format}")
+    filename = os.path.join(output_dir, f"sentence_{sentence_index:03d}.{response_format}")
     with open(filename, "wb") as f:
         f.write(audio_bytes)
     return filename
@@ -110,207 +95,183 @@ 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,
-    texts: list[str],
+    text: 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.
-
-    Every text is synthesized over the same connection: input.done flushes an
-    utterance without closing, so only the first one pays the handshake.
-    """
+    """Connect to the streaming TTS endpoint and process audio responses."""
     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; it stays in effect for every utterance below.
+        # 1. Send session config
         config_msg = {"type": "session.config", **config}
         await ws.send(json.dumps(config_msg))
         print(f"Sent session config: {config}")
 
-        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 receive_utterance(
-                    ws,
-                    output_dir=output_dir,
-                    response_format=response_format,
-                    word_timestamps=word_timestamps,
-                    save_chunks=save_chunks,
+        # 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,
+                        }
+                    )
                 )
-            finally:
-                sender_task.cancel()
-                try:
-                    await sender_task
-                except asyncio.CancelledError:
-                    pass  # Task cancellation is expected during shutdown
+                print(f"Sent full text: {text!r}")
 
-        # 3. Release the connection once there is nothing left to say.
-        await ws.send(json.dumps({"type": "session.close"}))
-        print("Sent session.close")
+            # 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()
+            try:
+                await sender_task
+            except asyncio.CancelledError:
+                pass  # Task cancellation is expected during shutdown
 
     print(f"\nAudio files saved to: {output_dir}/")
 
@@ -325,8 +286,7 @@ def main():
     parser.add_argument(
         "--text",
         required=True,
-        action="append",
-        help="Text to synthesize; repeat to send several utterances over one connection",
+        help="Text to synthesize",
     )
     parser.add_argument(
         "--output-dir",
@@ -425,7 +385,7 @@ def main():
     asyncio.run(
         stream_tts(
             url=args.url,
-            texts=args.text,
+            text=args.text,
             config=config,
             output_dir=args.output_dir,
             simulate_stt=args.simulate_stt,

pyproject.toml

@@ -75,7 +75,6 @@ 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次" -> "一次")
@@ -96,12 +95,6 @@ quack = [
     "quack-kernels>=0.3.11",
 ]
 
-# Optional Blackwell-native FlashAttention-4 backend. Keep this opt-in because
-# the CuTe/CUTLASS dependency is CUDA 13 and Blackwell specific.
-fa4 = [
-    "flash-attn-4[cu13]==4.0.0b18",
-]
-
 docs = [
     "mkdocs>=1.5.0",
     "mkdocs-api-autonav",

realtime_video_prompt_interaction_protocol.md

@@ -0,0 +1,138 @@
+# 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/e2e/features/fullduplex   # framework + JoyVL unit tests
+pytest tests/fullduplex   # framework + JoyVL unit tests
 ```
 
 ## Testing with an RTSP stream (optional)

recipes/LTX/LTX-2.md

@@ -165,40 +165,6 @@ 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
@@ -307,8 +273,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` shards the complete LTX guidance plan, including STG,
-  modality guidance, and rescale-compatible prediction gathering.
+- `--cfg-parallel-size 2` supports only the CFG-only plan: disable STG and
+  modality guidance and set both rescale fields to `0.0`.
 
 ## Operational Notes
 

recipes/MiniMaxAI/MiniMax-H3-NPU.md

@@ -1,162 +0,0 @@
-# MiniMax H3 on Ascend NPU
-
-> Joint video and audio generation with text, first-frame, image/audio, or
-> multi-video conditions — Ascend NPU deployment guide
-
-## Summary
-
-- Vendor: MiniMaxAI
-- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
-- Tasks: T2VA, FL2VA, and Ref2VA
-- Mode: OpenAI-compatible `/v1/videos` HTTP serving
-- Hardware: Atlas 800I A3
-- Maintainer: Community
-
-This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for Ascend NPU
-environments. Differences from the GPU path:
-
-- Runs on a CPU-only (aarch64) PyTorch build with `torch_npu`; no CUDA
-  runtime is present.
-- Audio loading does **not** require TorchCodec (whose aarch64 wheels are
-  built against CUDA torch and fail to load on CPU-only builds). vLLM-Omni
-  automatically falls back to soundfile / ffmpeg for wav/mp3/m4a/mp4 audio
-  inputs.
-
-## Prerequisites
-
-### Checkpoint
-
-Same as the GPU recipe — Hugging Face access approval is required:
-
-```bash
-hf auth login
-export MODEL_ROOT=/path/to/MiniMax-H3
-hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
-```
-
-### Environment
-
-- Ascend driver & firmware: 25.5.1
-- CANN toolkit: 9.0.1
-- Python: 3.12
-- PyTorch: 2.10.0+cpu
-- torch_npu: 2.10.0.post2
-- Install vLLM-Omni from a checkout with MiniMax H3 support:
-
-```bash
-uv venv
-source .venv/bin/activate
-uv pip install -e .
-```
-
-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
-```
-
-- `ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
-  reference-video preparation and MP4 output.
-- Reference-video decoding uses `decord` when available and falls back to
-  PyAV otherwise.
-- Audio inputs do not require TorchCodec: wav/mp3/m4a/mp4 files are loaded
-  at native sample rate through the soundfile / ffmpeg fallback.
-
-## Start a server
-
-One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
-and FL2VA requests, or to `Ref2VA` for Ref2VA requests.
-
-### Multi-NPU: 768P validated configuration
-
-Validated on eight NPUs of an Atlas 800I A3 server with Ulysses sequence
-parallelism degree 8, text-encoder tensor parallelism degree 8, native tiled
-VAE patch parallelism degree 8, and layerwise offload:
-
-```bash
-export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
-export PORT=9098
-export MODEL="${MODEL_ROOT}/FL2VA"
-export VLLM_WORKER_MULTIPROC_METHOD=spawn
-export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800
-export PYTHONDONTWRITEBYTECODE=1
-export HF_HUB_OFFLINE=1
-export TRANSFORMERS_OFFLINE=1
-
-vllm serve "${MODEL}" \
-  --omni \
-  --host 0.0.0.0 \
-  --port "${PORT}" \
-  --trust-remote-code \
-  --num-gpus 8 \
-  --usp 8 \
-  --ring 1 \
-  --text-encoder-tp-size 8 \
-  --enable-layerwise-offload \
-  --vae-parallel-mode tile \
-  --vae-use-tiling \
-  --vae-patch-parallel-size 8 \
-  --diffusion-attention-backend FLASH_ATTN
-```
-
-Do not add `--enforce-eager`. The first request includes regional
-compilation; warm the server once before measuring steady-state latency.
-H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
-
-To serve Ref2VA, stop the FL2VA server and restart with:
-
-```bash
-export MODEL="${MODEL_ROOT}/Ref2VA"
-```
-
-## HTTP API examples
-
-The request format is identical to the GPU recipe; see
-[MiniMax-H3.md § HTTP API examples](MiniMax-H3.md#http-api-examples).
-Use the validated 768P shapes (e.g. `width=1344 height=768`) on NPU.
-
-## Key parameters
-
-Same as the GPU recipe; see
-[MiniMax-H3.md § Key parameters](MiniMax-H3.md#key-parameters).
-The validated resolution on NPU is 768P (e.g. 1344x768).
-
-## Validated NPU evidence
-
-Measured on an Atlas 800I A3 server (8x NPU) with CANN 9.0.1,
-PyTorch 2.10.0+cpu, and torch_npu 2.10.0.post2, using the multi-NPU
-configuration above:
-
-| Workload | Configuration |
-|----------|---------------|
-| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
-| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
-
-These measurements describe the validated shapes rather than a general
-throughput guarantee.
-
-## Known limitations
-
-- Each server process loads only one checkpoint partition.
-- H3 currently executes one generation request per diffusion batch.
-- The first regional-compile request is a warmup and should not be included
-  in steady-state performance measurements.
-- Image+audio Ref2VA accepts exactly one image and one audio reference.
-- Video Ref2VA accepts one or more video files, but not an additional
-  standalone audio reference.
-- VAE patch parallelism requires size 1 or the full DiT group size and
-  supports the H3 native `tile` mode only.
-
-## Additional resources
-
-- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
-- [Supported models](../../docs/models/supported_models.md)
-- [Video API](../../docs/serving/videos_api.md)

recipes/MiniMaxAI/MiniMax-H3.md

@@ -1,329 +0,0 @@
-# MiniMax H3
-
-> Joint video and audio generation with text, first-frame, image/audio, or
-> multi-video conditions
-
-## Summary
-
-- Vendor: MiniMaxAI
-- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
-- Tasks: T2VA, FL2VA, and Ref2VA
-- Mode: OpenAI-compatible `/v1/videos` HTTP serving
-- Maintainer: Community
-
-MiniMax H3 is a CFG-distilled joint video/audio diffusion transformer. Its
-checkpoint has two independently served partitions:
-
-- `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio
-  (`fl2va`)
-- `Ref2VA`: image+audio or one-or-more-video reference-to-video+audio
-  (`ref2va`)
-
-The generated MP4 contains H.264 video and synchronized stereo audio.
-
-## Prerequisites
-
-The checkpoint requires Hugging Face access approval. Download it and point
-`MODEL_ROOT` at the directory containing the `FL2VA` and `Ref2VA`
-subdirectories:
-
-```bash
-hf auth login
-export MODEL_ROOT=/path/to/MiniMax-H3
-hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
-```
-
-Install vLLM-Omni from the checkout containing MiniMax H3 support. On
-Blackwell, install the optional FlashAttention-4 dependency:
-
-```bash
-uv venv
-source .venv/bin/activate
-uv pip install -e '.[fa4]'
-```
-
-`ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
-reference-video preparation and MP4 output.
-
-## Start a server
-
-One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
-and FL2VA requests, or to `Ref2VA` for either Ref2VA request.
-
-### Single GPU: accuracy and memory first
-
-The single-GPU configuration uses model-level CPU offload.
-This matches the accuracy-qualified reference path and prevents the Qwen3-VL
-encoder and DiT from being resident on the GPU at the same time.
-
-```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
-export PORT=8091
-
-CUDA_VISIBLE_DEVICES=0 \
-VLLM_WORKER_MULTIPROC_METHOD=spawn \
-VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
-vllm serve "${MODEL}" \
-  --omni \
-  --host 0.0.0.0 \
-  --port "${PORT}" \
-  --trust-remote-code \
-  --num-gpus 1 \
-  --enable-cpu-offload \
-  --diffusion-attention-backend FLASH_ATTN
-```
-
-Use a GPU with enough memory for the active H3 component and enough system RAM
-for the offloaded components. CPU offload reduces GPU memory pressure but adds
-PCIe/NVLink transfer latency.
-
-### Four GPUs: measured best-practice throughput
-
-The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
-
-- no CPU or layerwise offload;
-- Ulysses sequence parallelism degree 4;
-- native tiled VAE patch parallelism degree 4;
-- regional `torch.compile` for the repeated DiT blocks;
-- FlashAttention, with Ring and TP left at 1.
-
-```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
-export PORT=8091
-
-CUDA_VISIBLE_DEVICES=0,1,2,3 \
-FLASHINFER_DISABLE_VERSION_CHECK=1 \
-VLLM_WORKER_MULTIPROC_METHOD=spawn \
-VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
-vllm serve "${MODEL}" \
-  --omni \
-  --host 0.0.0.0 \
-  --port "${PORT}" \
-  --trust-remote-code \
-  --num-gpus 4 \
-  --usp 4 \
-  --ring 1 \
-  --vae-patch-parallel-size 4 \
-  --vae-parallel-mode tile \
-  --vae-use-tiling \
-  --diffusion-attention-backend FLASH_ATTN
-```
-
-Do not add `--enforce-eager` to this performance configuration. The first
-request includes regional compilation; warm the server once before measuring
-steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
-1. The H3 VAE supports its native `tile` mode, not
-`spatial_shard_height` or `spatial_shard_width`.
-
-### Text encoder tensor parallelism
-
-The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by
-default fully resident on the DiT main rank.  On multi-GPU no-offload runs that
-rank becomes the peak-memory hotspot.  Add `--text-encoder-tp-size N` to shard
-the encoder across the first `N` DiT ranks (the encoder is implemented with
-vLLM-style tensor-parallel layers and runs with distributed collectives over
-its own encoder process group):
-
-```bash
-export MODEL="${MODEL_ROOT}/FL2VA"
-export PORT=8091
-
-CUDA_VISIBLE_DEVICES=0,1,2,3 \
-FLASHINFER_DISABLE_VERSION_CHECK=1 \
-VLLM_WORKER_MULTIPROC_METHOD=spawn \
-VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
-vllm serve "${MODEL}" \
-  --omni \
-  --host 0.0.0.0 \
-  --port "${PORT}" \
-  --trust-remote-code \
-  --num-gpus 4 \
-  --usp 4 \
-  --ring 1 \
-  --text-encoder-tp-size 4 \
-  --vae-patch-parallel-size 4 \
-  --vae-parallel-mode tile \
-  --vae-use-tiling \
-  --diffusion-attention-backend FLASH_ATTN
-```
-
-`N` must divide the Qwen3-VL head counts (64 attention heads / 8 KV heads), so
-valid values on a 4-GPU server are 1, 2, and 4 (1, 2, 4, 8 on 8 GPUs).  The
-encoder TP rank set is the first `N` DiT ranks; on 4 GPUs
-`--text-encoder-tp-size 4` shards the encoder 4-way, dropping the DiT main
-rank's no-offload peak by roughly `(N-1)/N` of the ~51.5 GB encoder while the
-other ranks each gain `~51.5/N` GB.  The encoder output remains identical to
-the reference path within bf16 rounding: every encoder rank all-reduces the
-row-parallel projections, so the full `[seq, 5120]` layer-50 hidden state is
-replicated on every rank.
-
-To serve Ref2VA, stop the FL2VA server and restart the same one- or four-GPU
-command with:
-
-```bash
-export MODEL="${MODEL_ROOT}/Ref2VA"
-```
-
-## HTTP API examples
-
-The following requests use the synchronous endpoint so the returned body can
-be saved directly as an MP4. The asynchronous `POST /v1/videos` endpoint can
-also be used when job polling is preferred.
-
-All four tasks use 24 FPS, 50 sigma points, seed values from the validated
-workloads, and the checkpoint-reference video/audio flow shifts of 12 and 3.
-Decimal durations are passed through `extra_params`.
-
-Set the endpoint once:
-
-```bash
-export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
-```
-
-### 1. T2VA: text to video and audio
-
-Run this request against the `FL2VA` partition:
-
-```bash
-curl -sS -X POST "${API_URL}" \
-  -F 'prompt=In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' \
-  -F 'width=1344' \
-  -F 'height=768' \
-  -F 'fps=24' \
-  -F 'num_inference_steps=50' \
-  -F 'flow_shift=12' \
-  -F 'seed=1101' \
-  -F 'extra_params={"task":"t2va","duration":8.7,"audio_flow_shift":3.0}' \
-  -o t2va.mp4
-```
-
-### 2. FL2VA: first frame to video and audio
-
-Run this request against the `FL2VA` partition. When width and height are
-omitted, H3 preserves the first-frame aspect ratio and uses a 768-pixel short
-edge.
-
-```bash
-export FIRST_FRAME=/path/to/fl2va_first_frame.png
-
-curl -sS -X POST "${API_URL}" \
-  -F 'prompt=A man stands beside a yellow car at night. The car drives away; he follows it with his eyes and begins singing sadly, with synchronized voice and city ambience.' \
-  -F 'fps=24' \
-  -F 'num_inference_steps=50' \
-  -F 'flow_shift=12' \
-  -F 'seed=2101' \
-  -F 'extra_params={"task":"fl2va","duration":8.7,"audio_flow_shift":3.0}' \
-  -F "input_reference=@${FIRST_FRAME};type=image/png" \
-  -o fl2va.mp4
-```
-
-### 3. Ref2VA: image and audio reference
-
-Run this request against the `Ref2VA` partition. `audio_reference` accepts an
-HTTP(S) URL or a `data:` URL. In one terminal, expose the local reference
-assets to the serving host:
-
-```bash
-python -m http.server 8092 \
-  --bind 127.0.0.1 \
-  --directory /path/to/reference_assets
-```
-
-Then submit the image and audio request from another terminal:
-
-```bash
-export REF_IMAGE=/path/to/reference_assets/ref2va_image.png
-export AUDIO_URL=http://127.0.0.1:8092/ref2va_audio.mp3
-
-curl -sS -X POST "${API_URL}" \
-  -F 'prompt=A white cat with black mustache and eyebrow markings sits on a beige couch, lip-syncing precisely to the complete reference audio before shifting from confusion to deadpan speechlessness.' \
-  -F 'width=1344' \
-  -F 'height=768' \
-  -F 'fps=24' \
-  -F 'num_inference_steps=50' \
-  -F 'flow_shift=12' \
-  -F 'seed=3101' \
-  -F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \
-  -F "input_reference=@${REF_IMAGE};type=image/png" \
-  -F "audio_reference={\"audio_url\":\"${AUDIO_URL}\"}" \
-  -o ref2va_image_audio.mp4
-```
-
-The requested duration should cover the complete audio. If `duration` is
-shorter, the reference soundtrack is truncated to the generated clip.
-
-### 4. Ref2VA: two video references
-
-Run this request against the `Ref2VA` partition. Repeat the
-`input_references` multipart field once per source video. H3 consumes the
-videos in form order and preserves their original soundtracks during
-conditioning.
-
-```bash
-export SUBJECT_VIDEO=/path/to/green_screen_subject.mp4
-export BACKGROUND_VIDEO=/path/to/fairytale_background.mov
-
-curl -sS -X POST "${API_URL}" \
-  -F 'prompt=Remove the green screen background of Video 1 and replace it with the fairytale environment from Video 2. Match the background motion to the character actions and relight the character to fit the scene.' \
-  -F 'width=1344' \
-  -F 'height=768' \
-  -F 'fps=24' \
-  -F 'num_inference_steps=50' \
-  -F 'flow_shift=12' \
-  -F 'seed=3101' \
-  -F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \
-  -F "input_references=@${SUBJECT_VIDEO};type=video/mp4" \
-  -F "input_references=@${BACKGROUND_VIDEO};type=video/quicktime" \
-  -o ref2va_video_video.mp4
-```
-
-The server stores uploaded multi-video references only for the lifetime of the
-request and deletes the temporary files after generation. Video Ref2VA uses
-the source-video soundtracks and does not accept a separate
-`audio_reference`.
-
-## Key parameters
-
-| Parameter | Recommended value | Notes |
-|-----------|-------------------|-------|
-| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; must match the served partition |
-| `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count |
-| `fps` | `24` | H3 output FPS is fixed |
-| `num_inference_steps` | `50` | Matches the reference accuracy workloads |
-| `flow_shift` | `12` | Video sigma shift |
-| `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` |
-| `seed` | Task-specific | Use a fixed value for reproducibility |
-| `width`, `height` | Multiples of 32 | Aspect ratio must be between 1:4 and 4:1 |
-
-## Validated four-GPU evidence
-
-The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one
-excluded warmup followed by three requests.
-
-| Workload | Configuration | Observed result |
-|----------|---------------|-----------------|
-| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency |
-| Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean |
-
-These measurements describe the validated shapes rather than a general
-throughput guarantee. Multi-video Ref2VA is much slower because the two
-reference videos expand both the Qwen3-VL vision sequence and the packed DiT
-attention sequence.
-
-## Known limitations
-
-- Each server process loads only one checkpoint partition.
-- H3 currently executes one generation request per diffusion batch.
-- The first regional-compile request is a warmup and should not be included in
-  steady-state performance measurements.
-- Image+audio Ref2VA accepts exactly one image and one audio reference.
-- Video Ref2VA accepts one or more video files, but not an additional standalone
-  audio reference.
-- VAE patch parallelism requires size 1 or the full DiT group size and supports
-  the H3 native `tile` mode only.
-
-## Additional resources
-
-- [Supported models](../../docs/models/supported_models.md)
-- [Video API](../../docs/serving/videos_api.md)
-- [Diffusion parallelism](../../docs/user_guide/diffusion/parallelism/overview.md)

recipes/NVIDIA/GR00T-N1.7.md

@@ -69,13 +69,8 @@ Notes:
 ## Verification
 
 ```python
-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"
+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))
 ```
 
 Or run the e2e test suite:

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

@@ -30,6 +30,8 @@ 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)
@@ -66,7 +68,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_2gpu.yaml` is the recommended
+single-GPU entry point; `minicpmo_4_5_batching.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.
 
@@ -77,9 +79,7 @@ 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.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
+`gpu_memory_utilization` budgets are 0.55, 0.22, and 0.22. 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_2gpu.yaml
+--deploy-config vllm_omni/deploy/minicpmo_4_5_batching.yaml
 ```
 
 #### Performance comparison
@@ -175,18 +175,16 @@ 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** (use `--stream`):
+**Streaming text + speech**:
 
 ```bash
-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
+python examples/online_serving/minicpmo/streaming_chat_completion.py \
+    --base-url http://localhost:8099/v1 \
+    --output minicpmo_stream.wav
 ```
 
-The client prints text deltas as they arrive and saves streamed audio chunks
-to WAV files.
+The client prints text deltas as they arrive and reconstructs one valid WAV
+from the independently encoded audio deltas.
 
 **Gradio demo (text + image + audio + video UI)**:
 
@@ -204,10 +202,9 @@ speech output (TTS)"** checkbox on / off.
 
 #### Notes
 
-- 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.
+- 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.
 - `--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/README.md

@@ -45,7 +45,6 @@ recipes/
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
 | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
 | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
-| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
 | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
 | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |

recipes/cosmos3/Cosmos3-DistOffload.md

@@ -1,191 +0,0 @@
-# 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,9 +17,7 @@ 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. 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:
+- **Defaults** differ; omit `--height` / `--width` to get them:
 
 | | Edge | Nano / Super |
 |---|---|---|
@@ -28,10 +26,6 @@ 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,134 +125,3 @@ 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/ar_diffusion/test_engine_routing.py

@@ -21,8 +21,6 @@ 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/ar_diffusion/test_kv_cache.py

@@ -25,9 +25,6 @@ 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/ar_diffusion/test_kv_orchestration.py

@@ -7,7 +7,6 @@ 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/ar_diffusion/test_paged_attention.py

@@ -22,9 +22,6 @@ 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
@@ -388,22 +385,17 @@ def test_custom_op_compiles_fullgraph_without_recompile_on_value_change():
         return out
 
     torch._dynamo.reset()
-    try:
-        from torch._dynamo.testing import CompileCounter
+    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}"
-    finally:
-        # Leave a clean dynamo state for later suites in the same pytest process
-        # (e.g. model_executor transformers models).
-        torch._dynamo.reset()
+    assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}"

tests/ar_diffusion/test_pipeline_bridge.py

@@ -12,9 +12,6 @@ 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/ar_diffusion/test_slot_mapping.py

@@ -13,9 +13,6 @@ 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/benchmarks/test_daily_omni_pack_modes.py

@@ -1,249 +0,0 @@
-"""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/comfyui/conftest.py

@@ -25,7 +25,7 @@ def _setup_comfyui_test_environment():
     """Set up the test environment for ComfyUI plugin testing."""
     # === Add ComfyUI plugin path to allow importing comfyui_vllm_omni ===
     _COMFYUI_PLUGIN_PATH = os.path.abspath(
-        os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "apps", "ComfyUI-vLLM-Omni")
+        os.path.join(os.path.dirname(__file__), "..", "..", "apps", "ComfyUI-vLLM-Omni")
     )
     if not os.path.isdir(_COMFYUI_PLUGIN_PATH):
         raise FileNotFoundError(

tests/comfyui/test_comfyui_integration.py

@@ -6,13 +6,11 @@ It ensures that
 2. The sampling parameters are correctly passed from the node to AsyncOmni through the API layer.
 """
 
-from __future__ import annotations
-
 import multiprocessing
 import time
 import traceback
 from collections.abc import Iterable, Sequence
-from enum import Enum
+from enum import StrEnum, auto
 from types import SimpleNamespace
 from typing import Any, NamedTuple
 
@@ -35,7 +33,7 @@ from vllm.outputs import CompletionOutput, RequestOutput
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams
+from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 
@@ -54,20 +52,20 @@ class ServerCase(NamedTuple):
 class SamplingCase(NamedTuple):
     """Parametrizing the input sampling parameters."""
 
-    kind: SamplingKind
+    kind: "SamplingKind"
     sampling_params: dict | list[dict] | None
     lora: dict | None = None
 
 
-class SamplingKind(str, Enum):
-    IMAGE_NONE = "image_none"
-    IMAGE_DIFFUSION_SINGLE = "image_diffusion_single"
-    UNDERSTANDING_NONE = "understanding_none"
-    UNDERSTANDING_AR_LIST = "understanding_ar_list"
-    TTS_NONE = "tts_none"
-    TTS_DIFFUSION_SINGLE = "tts_diffusion_single"
-    VIDEO_NONE = "video_none"
-    VIDEO_DIFFUSION_SINGLE = "video_diffusion_single"
+class SamplingKind(StrEnum):
+    IMAGE_NONE = auto()
+    IMAGE_DIFFUSION_SINGLE = auto()
+    UNDERSTANDING_NONE = auto()
+    UNDERSTANDING_AR_LIST = auto()
+    TTS_NONE = auto()
+    TTS_DIFFUSION_SINGLE = auto()
+    VIDEO_NONE = auto()
+    VIDEO_DIFFUSION_SINGLE = auto()
 
 
 # Pre-defined arguments to be used in function calls during the tests
@@ -417,7 +415,7 @@ def mock_async_omni(
     mock_instance.stage_configs = server_case.stage_configs
     mock_instance.output_modalities = _build_output_modalities(server_case.stage_configs)
     mock_instance.default_sampling_params_list = [
-        OmniDiffusionSamplingParams() if _stage_type(stage) == "diffusion" else SamplingParams()
+        SamplingParams() if _stage_type(stage) != "diffusion" else mocker.MagicMock()
         for stage in server_case.stage_configs
     ]
     mock_instance.errored = False

tests/config/test_omni_config.py

@@ -440,7 +440,6 @@ def test_sub_config_fields_match_rfc_scopes():
         "ulysses_mode",
         "cfg_parallel_size",
         "vae_patch_parallel_size",
-        "text_encoder_tp_size",
         "vae_parallel_mode",
         "use_hsdp",
         "mask_sp_padding",
@@ -578,11 +577,9 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degr
     assert stage.parallel_config.world_size == 2
 
 
-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_cfg_parallel_size_outside_current_bound():
+    with pytest.raises(ValidationError):
+        OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
 
 
 def test_diffusion_parallel_config_rejects_allgather_with_ulysses_or_ring():

tests/conftest.py

@@ -3,7 +3,6 @@ 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

@@ -1,74 +0,0 @@
-"""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

@@ -1,57 +0,0 @@
-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,7 +76,6 @@ 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
@@ -105,87 +104,7 @@ class _RequestQueue(list):
                 self.remove(request)
 
 
-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])
-
-
+@pytest.mark.skip(reason="#issue 5484")
 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")
@@ -219,8 +138,7 @@ 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)
-        # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
-        return None, None
+        return 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,52 +262,9 @@ 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,7 +21,6 @@ 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
@@ -55,18 +54,6 @@ 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(
@@ -94,17 +81,6 @@ 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
@@ -157,21 +133,6 @@ 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,7 +75,6 @@ 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,8 +18,6 @@ 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,9 +3,6 @@ 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,8 +1,6 @@
 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
 
@@ -65,78 +63,35 @@ 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_context():
+def omni_server(request):
     """Start vLLM-Omni server as a subprocess with actual model weights.
-    Reuse it for adjacent benchmark cases with the same server configuration.
+    Uses session scope so the server starts only once for the entire test session.
     Multi-stage initialization can take 10-20+ minutes.
     """
     with _omni_server_lock:
-        active_context = _SingleActiveContext()
-        try:
-            yield active_context
-        finally:
-            active_context.close()
+        test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = request.param
 
+        print(f"Starting OmniServer with test: {test_name}, model: {model}")
 
-@pytest.fixture
-def omni_server(request, omni_server_context):
-    return omni_server_context.acquire(request.param, lambda: _start_omni_server(request.param))
+        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

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -5,8 +5,7 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100",
-            "npu": "A3"
+            "cuda": "H100"
           },
           "num_cards": 4
         }
@@ -28,9 +27,16 @@
             "async_chunk": false,
             "trust_remote_code": true,
             "connectors": {
-              "shared_mem_connector": {
-                "name": "SharedMemoryConnector",
-                "extra": {}
+              "rdma_connector": {
+                "name": "MooncakeTransferEngineConnector",
+                "extra": {
+                  "host": "auto",
+                  "zmq_port": 50051,
+                  "protocol": "rdma",
+                  "device_name": "",
+                  "memory_pool_size": 17179869184,
+                  "memory_pool_device": "cpu"
+                }
               }
             },
             "stages": [
@@ -60,7 +66,7 @@
                   "need_send_cache": true
                 },
                 "output_connectors": {
-                  "to_stage_1": "shared_mem_connector"
+                  "to_stage_1": "rdma_connector"
                 },
                 "default_sampling_params": {
                   "temperature": 0.0,
@@ -88,9 +94,8 @@
                   "vae_patch_parallel_size": 2,
                   "enable_expert_parallel": true
                 },
-                "vae_use_tiling": true,
                 "input_connectors": {
-                  "from_stage_0": "shared_mem_connector"
+                  "from_stage_0": "rdma_connector"
                 },
                 "default_sampling_params": {
                   "num_inference_steps": 50,

tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json

@@ -5,8 +5,7 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100",
-            "npu": "A3"
+            "cuda": "H100"
           },
           "num_cards": 4
         }
@@ -22,8 +21,6 @@
         "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

@@ -5,8 +5,7 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100",
-            "npu": "A3"
+            "cuda": "H100"
           },
           "num_cards": 4
         }
@@ -22,8 +21,6 @@
         "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

@@ -5,8 +5,7 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100",
-            "npu": "A3"
+            "cuda": "H100"
           },
           "num_cards": 4
         }
@@ -21,8 +20,6 @@
       "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,8 +1,6 @@
 """Tests for DFX runner metadata field exclusion."""
 
 import json
-from contextlib import contextmanager
-from types import SimpleNamespace
 
 import pytest
 
@@ -256,9 +254,8 @@ def test_benchmark_param_id_suffix_from_task_eval_phase():
     ]
 
 
-def test_paired_omni_benchmark_reuses_server_and_preserves_case_metadata(tmp_path, monkeypatch):
+def test_create_paired_omni_benchmark_pytest_params(tmp_path):
     from tests.dfx.conftest import create_paired_omni_benchmark_pytest_params
-    from tests.dfx.perf.scripts import run_benchmark
 
     configs = [
         {
@@ -288,51 +285,3 @@ def test_paired_omni_benchmark_reuses_server_and_preserves_case_metadata(tmp_pat
     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,16 +14,17 @@ 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.yaml")
+_MINICPMO_DEPLOY = get_deploy_config_path("minicpmo_4_5_batching.yaml")
 
 _MINICPMO_SERVER = [
     pytest.param(
         OmniServerParams(
             model="openbmb/MiniCPM-o-4_5",
             stage_config_path=_MINICPMO_DEPLOY,
-            use_stage_cli=False,
+            use_stage_cli=True,
             server_args=[
                 "--trust-remote-code",
+                "--no-async-chunk",
             ],
         ),
         id="minicpmo_4_5",
@@ -50,7 +51,7 @@ def _prompt() -> str:
     return "What is the capital of China? Answer in 20 words."
 
 
-@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=1)
+@hardware_test(res={"cuda": "H100", "npu": "A2"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", _MINICPMO_SERVER, indirect=True)
 def test_invalid_audio_format_rejected(omni_server: OmniServerParams, openai_client: OpenAIClientHandler) -> None:
     """

tests/diffusion/attention/test_attention_config.py

@@ -52,44 +52,6 @@ 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):
@@ -326,13 +288,6 @@ 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,8 +49,6 @@ 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,8 +18,6 @@ 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
@@ -329,48 +327,6 @@ def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
     assert calls[0]["softmax_scale"] == 0.5
 
 
-def test_packed_varlen_metadata_bypasses_mask_unpadding(monkeypatch):
-    calls = []
-
-    def fake_varlen_func(q, k, v, **kwargs):
-        calls.append((q.shape, k.shape, v.shape, kwargs))
-        return q
-
-    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
-    monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
-
-    impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False)
-    query = torch.randn(1, 8, 2, 4)
-    cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
-    metadata = AttentionMetadata(
-        attn_mask=torch.tensor([[True] * 6 + [False] * 2]),
-        extra={
-            "cu_seqlens_q": cu_seqlens,
-            "cu_seqlens_k": cu_seqlens,
-            "max_seqlen_q": 6,
-            "max_seqlen_k": 6,
-        },
-    )
-
-    output = impl.forward_cuda(query, query, query, metadata)
-
-    assert torch.equal(output, query)
-    assert len(calls) == 1
-    assert calls[0][0] == torch.Size([8, 2, 4])
-    assert calls[0][3]["cu_seqlens_q"] is cu_seqlens
-    assert calls[0][3]["max_seqlen_q"] == 6
-
-
-def test_packed_varlen_metadata_must_be_complete(monkeypatch):
-    monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
-    impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False)
-    query = torch.randn(1, 8, 2, 4)
-    metadata = AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, 8], dtype=torch.int32)})
-
-    with pytest.raises(ValueError, match="Incomplete packed FlashAttention metadata"):
-        impl.forward_cuda(query, query, query, metadata)
-
-
 if __name__ == "__main__":
     print("Running FlashAttention Padding Tests...")
     print("=" * 60)

tests/diffusion/attention/test_piecewise_attn.py

@@ -27,9 +27,6 @@ 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_ring_fa4.py

@@ -1,74 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.attention.backends.ring import ring_kernels
-from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
-from vllm_omni.diffusion.attention.parallel import ring as ring_parallel
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
-
-
-def test_fa4_wrapper_returns_lse_and_normalizes_unbounded_window(monkeypatch):
-    expected_out = object()
-    expected_lse = object()
-    captured = {}
-
-    def fake_fa4(q, k, v, **kwargs):
-        captured.update(q=q, k=k, v=v, **kwargs)
-        return expected_out, expected_lse
-
-    monkeypatch.setattr(ring_kernels, "HAS_FA4", True)
-    monkeypatch.setattr(ring_kernels, "fa4_attn_func", fake_fa4)
-
-    out, lse = ring_kernels.flash_attn4_func_forward(
-        "q",
-        "k",
-        "v",
-        softmax_scale=0.125,
-        causal=True,
-        window_size=(-1, -1),
-    )
-
-    assert (out, lse) == (expected_out, expected_lse)
-    assert captured == {
-        "q": "q",
-        "k": "k",
-        "v": "v",
-        "softmax_scale": 0.125,
-        "causal": True,
-        "window_size": (None, None),
-        "softcap": 0.0,
-        "return_lse": True,
-    }
-
-
-def test_ring_prefers_fa4_on_blackwell(monkeypatch):
-    captured = {}
-
-    def fake_ring_flash_attn_func(*args, **kwargs):
-        captured.update(kwargs)
-        return "output"
-
-    monkeypatch.setattr(ring_parallel, "HAS_FA4", True)
-    monkeypatch.setattr(ring_parallel, "HAS_FA3", True)
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3))
-
-    from vllm_omni.diffusion.attention.backends import ring_flash_attn
-
-    monkeypatch.setattr(ring_flash_attn, "ring_flash_attn_func", fake_ring_flash_attn_func)
-
-    query = SimpleNamespace(
-        dtype=torch.bfloat16,
-        shape=(1, 16, 4, 128),
-        is_cuda=True,
-        device=torch.device("cuda:0"),
-    )
-    strategy = ring_parallel.RingParallelAttention(SimpleNamespace(ring_group=object()))
-
-    assert strategy.run_attention(query, "key", "value", None) == "output"
-    assert captured["attn_type"] is AttnType.FA4

tests/diffusion/attention/test_sage_attn3.py

@@ -9,9 +9,6 @@ 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

@@ -1,138 +0,0 @@
-# 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,44 +47,6 @@ 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
 
@@ -295,32 +257,6 @@ 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,8 +22,6 @@ 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

@@ -1,122 +0,0 @@
-# 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,20 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""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
+"""Unit tests for SeqAllToAll4D and SeqAllToAll5D communication primitives."""
 
 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,
@@ -24,70 +16,87 @@ 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)
 
-DeviceKind = Literal["cpu", "cuda"]
+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
 
 
-def _update_environment_variables(envs_dict: dict[str, str]) -> None:
-    for key, value in envs_dict.items():
-        os.environ[key] = value
+@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(
+    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,
+    )
 
 
-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(
+def _test_4d_identity_worker(
     local_rank: int,
     world_size: int,
-    master_port: int,
-    device_kind: DeviceKind,
-) -> None:
-    _update_environment_variables(
+    dtype: torch.dtype,
+    batch_size: int,
+    seq_len_per_rank: int,
+    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)
+
+    # 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": str(master_port),
+            "MASTER_PORT": "29500",
         }
     )
-    backend = "gloo" if device_kind == "cpu" else None
-    init_distributed_environment(backend=backend)
 
+    # 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
 
-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,
-    batch_size: int,
-    seq_len_per_rank: int,
-    num_heads: int,
-    head_size: int,
-    use_sync: bool,
-    device_kind: DeviceKind,
-    master_port: int,
-) -> None:
-    device = _worker_device(local_rank, device_kind)
-    if device_kind == "cuda":
-        current_omni_platform.set_device(device)
-
-    _init_worker(local_rank, world_size, master_port, device_kind)
-    initialize_model_parallel(ulysses_degree=world_size)
-    sp_group = get_sp_group().ulysses_group
-
+    # Create input tensor: (bs, seqlen/P, hc, hs)
     torch.manual_seed(42 + local_rank)
     input_tensor = torch.randn(
         batch_size,
@@ -97,37 +106,65 @@ def _run_4d_identity(
         dtype=dtype,
         device=device,
     )
+
+    # Save original input for comparison
     original_input = input_tensor.clone()
 
-    intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, use_sync)
-    expected_intermediate_shape = (
+    # 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 = (
         batch_size,
         seq_len_per_rank * world_size,
         num_heads // world_size,
         head_size,
     )
-    assert intermediate.shape == expected_intermediate_shape, (
-        f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
+    assert intermediate.shape == expected_shape, (
+        f"Intermediate shape mismatch: expected {expected_shape}, got {intermediate.shape}"
     )
 
-    output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, use_sync)
+    # 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
     assert output.shape == original_input.shape, (
         f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
     )
 
-    rtol, atol = _close_tolerance(dtype)
+    # Verify output matches original input
     torch.testing.assert_close(
         output,
         original_input,
-        rtol=rtol,
-        atol=atol,
+        rtol=1e-5,
+        atol=1e-5,
         msg="Output does not match original input after two all-to-all operations",
     )
+
+    # Cleanup distributed environment
     destroy_distributed_env()
 
 
-def _run_5d_identity(
-    local_rank: int,
+@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,
@@ -135,329 +172,235 @@ def _run_5d_identity(
     num_heads: int,
     head_size: int,
     use_sync: bool,
-    device_kind: DeviceKind,
-    master_port: int,
-) -> None:
-    device = _worker_device(local_rank, device_kind)
-    if device_kind == "cuda":
-        current_omni_platform.set_device(device)
-
-    _init_worker(local_rank, world_size, master_port, device_kind)
-    initialize_model_parallel(ulysses_degree=world_size)
-    sp_group = get_sp_group().ulysses_group
-
-    torch.manual_seed(42 + local_rank)
-    input_tensor = torch.randn(
-        batch_size,
-        seq_len_per_rank,
-        3,
-        num_heads,
-        head_size,
-        dtype=dtype,
-        device=device,
-    )
-    original_input = input_tensor.clone()
-
-    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_intermediate_shape, (
-        f"Intermediate shape mismatch: expected {expected_intermediate_shape}, got {intermediate.shape}"
-    )
-
-    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}"
-    )
-
-    rtol, atol = _close_tolerance(dtype)
-    torch.testing.assert_close(
-        output,
-        original_input,
-        rtol=rtol,
-        atol=atol,
-        msg="Output does not match original input after two all-to-all operations",
-    )
-    destroy_distributed_env()
-
-
-def _run_ring_p2p(
-    local_rank: int,
-    world_size: int,
-    dtype: torch.dtype,
-    batch_size: int,
-    num_heads: int,
-    head_size: int,
-    device_kind: DeviceKind,
-    master_port: int,
-) -> None:
-    device = _worker_device(local_rank, device_kind)
-    if device_kind == "cuda":
-        current_omni_platform.set_device(device)
-
-    try:
-        _init_worker(local_rank, world_size, master_port, device_kind)
-        initialize_model_parallel(ring_degree=world_size)
-        sp_group = get_sp_group()
-        comm = RingComm(sp_group.ring_group)
-
-        input_tensor = torch.full(
-            (batch_size, num_heads, head_size),
-            fill_value=float(local_rank + 1),
-            dtype=dtype,
-            device=device,
-        )
-        recv_tensor = comm.send_recv(input_tensor)
-        comm.commit()
-        comm.wait()
-
-        prev_rank = (local_rank - 1 + world_size) % world_size
-        expected_value = float(prev_rank + 1)
-        expected_tensor = torch.full_like(recv_tensor, fill_value=expected_value)
-        rtol, atol = _close_tolerance(dtype)
-        torch.testing.assert_close(
-            recv_tensor,
-            expected_tensor,
-            rtol=rtol,
-            atol=atol,
-            msg=f"[Rank {local_rank}] Ring P2P data mismatch",
-        )
-    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:
+):
+    """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})")
 
-# ---------------------------------------------------------------------------
-# 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,
+    # 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,
     )
 
 
-@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,
+def _test_5d_identity_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,
+):
+    """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)
+
+    # 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",
+        }
     )
 
+    # 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
 
-@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,
+    # 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
+        num_heads,
+        head_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,
+        device=device,
     )
 
+    # Save original input for comparison
+    original_input = input_tensor.clone()
 
-@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),
-    ],
-)
+    # 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 = (
+        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}"
+    )
+
+    # 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
+    assert output.shape == original_input.shape, (
+        f"Output shape mismatch: expected {original_input.shape}, got {output.shape}"
+    )
+
+    # Verify output matches original input
+    torch.testing.assert_close(
+        output,
+        original_input,
+        rtol=1e-5,
+        atol=1e-5,
+        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("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.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,
     )
 
 
-@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,
+def _test_ring_p2p_worker(
+    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
+
+    # 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
+        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
+        )
+
+        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
+        torch.testing.assert_close(
+            recv_tensor, expected_tensor, rtol=1e-3, atol=1e-3, msg=f"[Rank {local_rank}] 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()

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -1,17 +1,8 @@
 # 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
@@ -20,7 +11,6 @@ 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,
@@ -32,12 +22,7 @@ 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
 
-_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]
+pytestmark = [pytest.mark.parallel]
 
 
 def _find_free_port() -> str:
@@ -51,22 +36,6 @@ 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
 # ---------------------------------------------------------------------------
@@ -174,7 +143,7 @@ class MockPipelineParallel(PipelineParallelMixin, CFGParallelMixin):
 class TestAsyncLatents:
     """Verifies the lazy-resolution behaviour of AsyncLatents without a real process group."""
 
-    pytestmark = _UNIT_MARKS
+    pytestmark = [pytest.mark.cpu]
 
     def _make(self, tensor: torch.Tensor, handles: list | None = None, postproc: list | None = None) -> AsyncLatents:
         return AsyncLatents({"latents": tensor}, handles or [], postproc or [])
@@ -245,7 +214,7 @@ class TestAsyncLatents:
 class TestSyncPPSend:
     """Verifies PipelineParallelMixin's internal PP-send flush."""
 
-    pytestmark = _UNIT_MARKS
+    pytestmark = [pytest.mark.cpu]
 
     @staticmethod
     def _make_pipeline() -> PipelineParallelMixin:
@@ -278,7 +247,7 @@ class TestSyncPPSend:
 class TestDiffuseWrapper:
     """Verifies that PipelineParallelMixin flushes pending sends when diffuse() exits."""
 
-    pytestmark = _UNIT_MARKS
+    pytestmark = [pytest.mark.cpu]
 
     def test_diffuse_flushes_pending_sends_on_success(self):
         work = FakeWork()
@@ -320,7 +289,7 @@ class TestDiffuseWrapper:
 
 
 class TestVaeDecodeGuard:
-    pytestmark = _UNIT_MARKS
+    pytestmark = [pytest.mark.cpu]
 
     @staticmethod
     def _make_pipeline(distributed_enabled: bool = False) -> PipelineParallelMixin:
@@ -382,8 +351,6 @@ 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"):
@@ -392,8 +359,6 @@ 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"):
@@ -407,22 +372,10 @@ def test_pipeline_parallel_requires_mro_before_cfg_mixin():
 # ---------------------------------------------------------------------------
 
 
-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:
+def init_dist(local_rank: int, world_size: int, master_port: str) -> torch.device:
     """Initialise the distributed environment for a spawned worker."""
-    device = _worker_device(local_rank, device_kind)
-    if device_kind == "cuda":
-        current_omni_platform.set_device(device)
+    device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
+    current_omni_platform.set_device(device)
     update_environment_variables(
         {
             "RANK": str(local_rank),
@@ -432,31 +385,13 @@ def init_dist(
             "MASTER_PORT": master_port,
         }
     )
-    backend = "gloo" if device_kind == "cpu" else None
-    init_distributed_environment(backend=backend)
+    init_distributed_environment()
     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``
@@ -473,12 +408,16 @@ def make_pipeline_and_inputs(
         dtype=dtype,
     )
 
-    _seed_device(test_config["input_seed"], device)
+    torch.manual_seed(test_config["input_seed"])
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(test_config["input_seed"])
     pos_x = {"x": torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)}
 
     negative_kwargs = None
     if do_true_cfg:
-        _seed_device(test_config["input_seed"] + 1, device)
+        torch.manual_seed(test_config["input_seed"] + 1)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(test_config["input_seed"] + 1)
         neg_x = torch.randn(test_config["batch_size"], test_config["dim"], dtype=dtype, device=device)
         negative_kwargs = {"x": neg_x}
 
@@ -486,23 +425,19 @@ def make_pipeline_and_inputs(
 
 
 # ---------------------------------------------------------------------------
-# 3.  isend_tensor_dict / irecv_tensor_dict
+# 3.  isend_tensor_dict / irecv_tensor_dict  (2 GPUs)
 # ---------------------------------------------------------------------------
 
 
-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))
+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)
     pp_group = get_pp_group()
 
     if pp_group.is_first_rank:
-        _seed_device(77, device)
+        torch.manual_seed(77)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(77)
         tensor = torch.randn(3, 5, dtype=torch.float32, device=device)
         handles = pp_group.isend_tensor_dict({"t": tensor})
         for h in handles:
@@ -514,42 +449,27 @@ def isend_irecv_worker(
 
     if torch.distributed.is_initialized():
         torch.distributed.barrier()
-    _cleanup_distributed()
+    destroy_distributed_env()
 
 
-def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, master_port: str) -> None:
+@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."""
     mp_context = torch.multiprocessing.get_context("spawn")
     manager = mp_context.Manager()
     q = manager.Queue()
-    torch.multiprocessing.spawn(
-        isend_irecv_worker,
-        args=(pp_size, master_port, device_kind, q),
-        nprocs=pp_size,
-    )
+
+    port = _find_free_port()
+    torch.multiprocessing.spawn(isend_irecv_worker, args=(pp_size, port, 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
 # ---------------------------------------------------------------------------
@@ -557,13 +477,12 @@ def test_isend_irecv_tensor_dict_parity(pp_size: int):
 _baseline_cache: dict[tuple, torch.Tensor] = {}
 
 
-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."""
+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.
+    """
     key = (
         test_config["num_layers"],
         test_config["dim"],
@@ -573,29 +492,27 @@ def compute_single_gpu_baseline(
         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(), device_kind)
-    try:
-        initialize_model_parallel(pipeline_parallel_size=1, backend=_mp_backend(device_kind))
+    device = init_dist(0, 1, _find_free_port())
+    initialize_model_parallel(pipeline_parallel_size=1)
 
-        pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
-            test_config, dtype, device, do_true_cfg=do_true_cfg
+    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,
         )
 
-        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()
+    destroy_distributed_env()
 
     _baseline_cache[key] = noise_pred.cpu()
     return _baseline_cache[key]
@@ -610,16 +527,11 @@ 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, device_kind)
-    initialize_model_parallel(
-        pipeline_parallel_size=pp_size,
-        cfg_parallel_size=cfg_size,
-        backend=_mp_backend(device_kind),
-    )
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
 
     pp_group = get_pp_group()
     cfg_rank = get_classifier_free_guidance_rank()
@@ -636,6 +548,8 @@ 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:
@@ -647,22 +561,77 @@ def predict_noise_worker(
 
     if torch.distributed.is_initialized():
         torch.distributed.barrier()
-    _cleanup_distributed()
+    destroy_distributed_env()
 
 
-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:
+@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."""
     test_config = {
         "num_layers": num_layers,
         "dim": 64,
@@ -672,20 +641,22 @@ def _run_predict_noise(
         "input_seed": input_seed,
     }
 
-    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)
+    baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg)
 
     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, master_port, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
+        args=(world_size, port, pp_size, cfg_size, do_true_cfg, dtype, test_config, pp_q),
         nprocs=world_size,
     )
 
     pp_out = pp_q.get()
+
     assert baseline_out.shape == pp_out.shape
     torch.testing.assert_close(
         pp_out,
@@ -696,157 +667,35 @@ def _run_predict_noise(
     )
 
 
-_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,
-    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))
+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)
 
-        pipeline, positive_kwargs, negative_kwargs = make_pipeline_and_inputs(
-            test_config, torch.float32, device, 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)
+
+    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
         )
-        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
-            )
-    finally:
-        _cleanup_distributed()
+    destroy_distributed_env()
     return result.cpu()
 
 
@@ -858,15 +707,10 @@ 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, device_kind)
-    initialize_model_parallel(
-        pipeline_parallel_size=pp_size,
-        cfg_parallel_size=cfg_size,
-        backend=_mp_backend(device_kind),
-    )
+    device = init_dist(local_rank, world_size, master_port)
+    initialize_model_parallel(pipeline_parallel_size=pp_size, cfg_parallel_size=cfg_size)
 
     pp_group = get_pp_group()
     cfg_rank = get_classifier_free_guidance_rank()
@@ -888,25 +732,43 @@ 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:
-        result_queue.put(latents.contiguous().cpu())
+        resolved = latents.contiguous()
+        result_queue.put(resolved.cpu())
 
     if torch.distributed.is_initialized():
         torch.distributed.barrier()
-    _cleanup_distributed()
+    destroy_distributed_env()
 
 
-def _run_scheduler_step(
-    *,
-    pp_size: int,
-    cfg_size: int,
-    do_true_cfg: bool,
-    input_seed: int,
-    device_kind: DeviceKind,
-    master_port: str,
-) -> None:
+@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."""
     test_config = {
         "num_layers": 4,
         "dim": 64,
@@ -916,16 +778,17 @@ def _run_scheduler_step(
         "input_seed": input_seed,
     }
 
-    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)
+    baseline = compute_scheduler_step_baseline(test_config, do_true_cfg)
 
     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, master_port, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
+        args=(world_size, port, pp_size, cfg_size, do_true_cfg, test_config, q),
         nprocs=world_size,
     )
 
@@ -937,47 +800,3 @@ def _run_scheduler_step(
         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_vae_decode_parallelism.py

@@ -1,181 +1,151 @@
-# 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"
 
-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": 4e-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",
-    ),
+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,
+    # },
 ]
 
 
-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 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 _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()
+def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
     parallel_config = DiffusionParallelConfig(
-        tensor_parallel_size=2,
+        tensor_parallel_size=tp,
         vae_patch_parallel_size=vae_patch_parallel_size,
     )
-    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),
+
+    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),
+            ),
         )
-        return _extract_output_array(outputs[0])
+        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()
 
 
-@pytest.mark.full_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 4})
-@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.")
+@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)
 
-    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"]
+    mean_threshold = 3e-2
+    max_threshold = 3e-2  # they should be totally same
     print(
-        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})"
+        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"
     )
+    # assert non_parallel_time > parallel_time
+    assert result_diff.mean() < mean_threshold
+    assert result_diff.max() < max_threshold

tests/diffusion/distributed/test_zimage_parallelism.py

@@ -1,14 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""DiT tensor-parallel parity and perf checks for diffusion models.
+"""Z-Image end-to-end tests for diffusion parallelism.
 
-Compares TP=2 against TP=1 on output correctness, latency, and peak memory.
-Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
+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`.
 """
 
-from __future__ import annotations
-
 import os
 import time
 from pathlib import Path
@@ -29,15 +31,22 @@ 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"
-DEFAULT_MODEL = "Tongyi-MAI/Z-Image-Turbo"
+
+
+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")
 
 
 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}"
@@ -62,7 +71,7 @@ def _extract_single_image(outputs) -> Image.Image:
     return images[0]
 
 
-def _run_generate(
+def _run_zimage_generate(
     *,
     tp_size: int,
     height: int,
@@ -70,6 +79,8 @@ def _run_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:
@@ -80,11 +91,25 @@ def _run_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(
-            DEFAULT_MODEL,
-            parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size),
+            _get_zimage_model(),
+            parallel_config=DiffusionParallelConfig(
+                tensor_parallel_size=tp_size,
+                vae_patch_parallel_size=vae_patch_parallel_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(
@@ -109,11 +134,14 @@ def _run_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()
@@ -123,20 +151,20 @@ def _run_generate(
 @pytest.mark.diffusion
 @pytest.mark.parallel
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
-def test_tensor_parallel_tp2(tmp_path: Path):
+def test_zimage_tensor_parallel_tp2(tmp_path: Path):
     if current_omni_platform.is_npu():
-        pytest.skip("Tensor parallel parity test is only supported on CUDA and ROCm for now.")
+        pytest.skip("Z-Image TP 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("Tensor parallel TP=2 requires >= 2 devices.")
+        pytest.skip("Z-Image 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_generate(
+    tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate(
         tp_size=1,
         height=height,
         width=width,
@@ -144,7 +172,7 @@ def test_tensor_parallel_tp2(tmp_path: Path):
         seed=seed,
         enforce_eager=enforce_eager,
     )
-    tp2_img, tp2_time_s, tp2_peak_mem = _run_generate(
+    tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate(
         tp_size=2,
         height=height,
         width=width,
@@ -153,8 +181,8 @@ def test_tensor_parallel_tp2(tmp_path: Path):
         enforce_eager=enforce_eager,
     )
 
-    tp1_path = tmp_path / "tp1.png"
-    tp2_path = tmp_path / "tp2.png"
+    tp1_path = tmp_path / "zimage_tp1.png"
+    tp2_path = tmp_path / "zimage_tp2.png"
     tp1_img.save(tp1_path)
     tp2_img.save(tp2_path)
 
@@ -165,7 +193,7 @@ def test_tensor_parallel_tp2(tmp_path: Path):
     mean_threshold = 3e-2
     p99_threshold = 2.5e-1
     print(
-        f"{model_name} TP image diff (TP=1 vs TP=2): "
+        "Z-Image TP image diff stats (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}"
@@ -175,11 +203,73 @@ def test_tensor_parallel_tp2(tmp_path: Path):
         f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
     )
 
-    print(f"TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
+    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
     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"TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
+    print(f"Z-Image 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/hooks/test_hook_registry.py

@@ -12,9 +12,6 @@ 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,8 +13,6 @@ 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

@@ -1,78 +0,0 @@
-# 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]

tests/diffusion/models/ltx2/test_ltx2_conditioning.py

@@ -110,7 +110,7 @@ class TestLTXImageToVideoForwardStages:
 
         kwargs = pipe._denoise_timestep_kwargs(
             ts,
-            SimpleNamespace(guidance_parallel_ready=False),
+            SimpleNamespace(cfg_parallel_ready=False),
             denoise_ctx,
             video_token_count=2,
             audio_token_count=1,

tests/diffusion/models/ltx2/test_ltx2_parallel.py

@@ -3,7 +3,6 @@
 
 """Unit tests for LTX guidance and forward-parallel behavior."""
 
-from contextlib import nullcontext
 from types import SimpleNamespace
 
 import pytest
@@ -51,8 +50,7 @@ class TestCFGParallelHelpers:
         assert torch.allclose(video_combined, expected_video)
         assert torch.allclose(audio_combined, expected_audio)
 
-    @pytest.mark.parametrize("guidance_world_size", [1, 2, 3, 4])
-    def test_guidance_parallel_accepts_official_guidance(self, guidance_world_size):
+    def test_cfg_parallel_rejects_non_cfg_guidance(self):
         from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
             LTX_GUIDANCE_EXECUTOR,
             LTXGuidancePlan,
@@ -61,180 +59,26 @@ class TestCFGParallelHelpers:
 
         plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
 
-        LTX_GUIDANCE_EXECUTOR.validate_guidance_world_size(plan, guidance_world_size)
+        with pytest.raises(ValueError, match="CFG-only guidance"):
+            LTX_GUIDANCE_EXECUTOR.validate_cfg_world_size(plan, 2)
 
-    def test_guidance_parallel_dummy_warms_full_guidance_plan(self, monkeypatch):
+    def test_cfg_parallel_dummy_warms_supported_cfg_only_plan(self, monkeypatch):
         from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
         from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
-        monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2)
+        monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
         pipe = object.__new__(LTX2Pipeline)
         req = SimpleNamespace(is_dummy_run=lambda: True)
         request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance)
 
-        guidance_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)
+        cfg_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)
 
-        assert guidance_parallel_ready is True
-        assert pipe._guidance_plan.names == ("cond", "uncond", "ptb", "mod")
+        assert cfg_parallel_ready is True
+        assert pipe._guidance_plan.names == ("cond", "uncond")
         assert pipe._guidance_plan.spec.video.cfg_scale == 3.0
         assert pipe._guidance_plan.spec.audio.cfg_scale == 7.0
 
-    def test_guidance_parallel_warns_for_imbalanced_passes(self, monkeypatch):
-        from vllm_omni.diffusion.models.ltx2 import ltx2_guidance
-
-        spec = ltx2_guidance.LTXGuidanceSpec(
-            video=ltx2_guidance.LTXModalityGuidance(
-                cfg_scale=3.0,
-                stg_scale=1.0,
-                stg_blocks=(1,),
-            )
-        )
-        plan = ltx2_guidance.LTXGuidancePlan.build(spec)
-        captured = []
-        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: 0)
-        monkeypatch.setattr(ltx2_guidance.logger, "warning_once", lambda *args: captured.append(args))
-
-        ltx2_guidance.LTX_GUIDANCE_EXECUTOR.warn_if_imbalanced(plan, 4, "hq_generate")
-
-        assert plan.names == ("cond", "uncond", "ptb")
-        assert len(captured) == 1
-        assert captured[0][1:] == ("hq_generate", 3, 4, 1, 75.0)
-
-    @pytest.mark.parametrize("guidance_rank", [0, 1])
-    def test_default_guidance_is_sharded_and_combined(self, monkeypatch, guidance_rank):
-        from vllm_omni.diffusion.models.ltx2 import ltx2_guidance
-        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
-        from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
-        from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
-
-        executor = ltx2_guidance.LTX_GUIDANCE_EXECUTOR
-        plan = ltx2_guidance.LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
-        video_predictions = {
-            "cond": torch.tensor([[[0.2, -0.3]]]),
-            "uncond": torch.tensor([[[-0.4, 0.1]]]),
-            "ptb": torch.tensor([[[0.5, -0.2]]]),
-            "mod": torch.tensor([[[-0.1, 0.6]]]),
-        }
-        audio_predictions = {
-            "cond": torch.tensor([[[0.7, -0.2]]]),
-            "uncond": torch.tensor([[[0.1, 0.4]]]),
-            "ptb": torch.tensor([[[-0.3, 0.2]]]),
-            "mod": torch.tensor([[[0.4, -0.5]]]),
-        }
-        state = LTXAVState(
-            video=torch.tensor([[[1.0, -2.0]]]),
-            audio=torch.tensor([[[0.5, 3.0]]]),
-        )
-        denoise_ctx = LTXDenoiseContext(
-            latents=state.video,
-            audio_latents=state.audio,
-            video_coords=torch.zeros(1, 1, 3),
-            audio_coords=torch.zeros(1, 1, 1),
-        )
-        executor.prepare_denoise_context(plan, True, 2, denoise_ctx)
-
-        class FakeGuidanceGroup:
-            def __init__(self):
-                self.call_index = 0
-
-            def all_gather(self, tensor, separate_tensors=True):
-                assert separate_tensors
-                gathered = (
-                    [video_predictions["cond"], video_predictions["uncond"]],
-                    [video_predictions["ptb"], video_predictions["mod"]],
-                    [audio_predictions["cond"], audio_predictions["uncond"]],
-                    [audio_predictions["ptb"], audio_predictions["mod"]],
-                )[self.call_index]
-                torch.testing.assert_close(tensor, gathered[guidance_rank])
-                self.call_index += 1
-                return gathered
-
-        group = FakeGuidanceGroup()
-        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_world_size", lambda: 2)
-        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: guidance_rank)
-        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_group", lambda: group)
-
-        local_names = ("cond", "ptb") if guidance_rank == 0 else ("uncond", "mod")
-        captured = {}
-
-        class FakeTransformer:
-            def __call__(self, **kwargs):
-                captured.update(kwargs)
-                return (
-                    torch.cat([video_predictions[name] for name in local_names]),
-                    torch.cat([audio_predictions[name] for name in local_names]),
-                )
-
-        video_sigma = torch.tensor(0.25)
-        audio_sigma = torch.tensor(0.5)
-
-        class FakePipeline:
-            scheduler = SimpleNamespace(sigmas=torch.stack([video_sigma]))
-            transformer = FakeTransformer()
-
-            @staticmethod
-            def _transformer_cache_context(name):
-                assert name == "guided"
-                return nullcontext()
-
-            @staticmethod
-            def _build_transformer_kwargs(forward_ctx, denoise_ctx, **kwargs):
-                del forward_ctx, denoise_ctx
-                return kwargs
-
-            @staticmethod
-            def _video_guidance_model_sigma(sigma, denoise_ctx):
-                del denoise_ctx
-                return sigma
-
-        positive_context = torch.ones(1, 1, 1)
-        negative_context = torch.zeros(1, 1, 1)
-        forward_ctx = SimpleNamespace(
-            prompt_context=SimpleNamespace(
-                positive_connector_prompt_embeds=positive_context,
-                positive_connector_audio_prompt_embeds=positive_context,
-                negative_connector_prompt_embeds=negative_context,
-                negative_connector_audio_prompt_embeds=negative_context,
-            ),
-            attention_kwargs=None,
-            audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])),
-        )
-
-        actual_video, actual_audio = executor.predict_parallel_guidance(
-            FakePipeline(),
-            plan,
-            index=0,
-            timestep=torch.tensor(1.0),
-            state=state,
-            forward_ctx=forward_ctx,
-            denoise_ctx=denoise_ctx,
-        )
-
-        expected_video = executor._guide_modality(
-            state.video,
-            video_predictions,
-            video_sigma,
-            plan.spec.video,
-        )
-        expected_audio = executor._guide_modality(
-            state.audio,
-            audio_predictions,
-            audio_sigma,
-            plan.spec.audio,
-        )
-        torch.testing.assert_close(actual_video, expected_video)
-        torch.testing.assert_close(actual_audio, expected_audio)
-        assert group.call_index == 4
-        assert captured["hidden_states"].shape[0] == 2
-        perturbations = captured["attention_kwargs"]["ltx_perturbation_kwargs"]
-        if guidance_rank == 0:
-            assert "video_self_attention_mask" in perturbations
-            assert "a2v_cross_attention_mask" not in perturbations
-        else:
-            assert "video_self_attention_mask" not in perturbations
-            assert "a2v_cross_attention_mask" in perturbations
-
 
 class TestCFGParallelForwardPath:
     """Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
@@ -252,7 +96,7 @@ class TestCFGParallelForwardPath:
         pipe.tokenizer_max_length = 4
         pipe.vae_spatial_compression_ratio = 32
         pipe.vae_temporal_compression_ratio = 8
-        monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 1)
+        monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 1)
 
         class StopAtEncodePromptError(Exception):
             pass
@@ -394,10 +238,10 @@ class TestCFGParallelForwardPath:
                     return [audio_pos, audio_neg]
                 raise AssertionError(f"Unexpected gathered tensor: {tensor}")
 
-        monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2)
-        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_world_size", lambda: 2)
-        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_rank", lambda: cfg_rank)
-        monkeypatch.setattr(ltx2_guidance, "get_guidance_parallel_group", lambda: FakeCfgGroup())
+        monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda: cfg_rank)
+        monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
 
         def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None):
             scheduler.sigmas = torch.tensor([0.25, 0.25], device=device)

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -9,8 +9,6 @@ from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransforme
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
 
 def test_ltx_rms_norm_no_affine_matches_official_torch_module():
     norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -1,375 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from multiprocessing.reduction import ForkingPickler
-from types import SimpleNamespace
-from unittest.mock import Mock
-
-import numpy as np
-import pytest
-import torch
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
-
-
-def test_pipeline_import_registry_and_component_discovery():
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-    from vllm_omni.diffusion.registry import (
-        _DIFFUSION_MODELS,
-        _DIFFUSION_POST_PROCESS_FUNCS,
-    )
-
-    assert _DIFFUSION_MODELS["MiniMaxH3Pipeline"] == (
-        "minimax_h3",
-        "pipeline_minimax_h3",
-        "MiniMaxH3Pipeline",
-    )
-    assert _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3Pipeline"] == "get_minimax_h3_post_process_func"
-    assert MiniMaxH3Pipeline._dit_modules == ["transformer"]
-    assert MiniMaxH3Pipeline._encoder_modules == ["text_encoder"]
-    assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"]
-
-
-def test_joint_postprocess_is_multiprocessing_picklable():
-    from vllm_omni.diffusion.models.minimax_h3 import (
-        get_minimax_h3_post_process_func,
-    )
-
-    postprocess = get_minimax_h3_post_process_func(SimpleNamespace())
-    postprocess = ForkingPickler.loads(ForkingPickler.dumps(postprocess))
-    video = torch.linspace(0, 1, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5)
-    audio = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6)
-
-    result = postprocess((video, audio), output_type="np")
-
-    assert isinstance(result["video"], list)
-    assert result["video"][0].shape == (2, 4, 5, 3)
-    np.testing.assert_array_equal(result["audio"], audio.numpy())
-    assert result["audio_sample_rate"] == 32000
-    assert result["fps"] == 24
-
-
-def test_cfg_parallel_is_rejected_for_distilled_checkpoint():
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-
-    od_config = SimpleNamespace(
-        parallel_config=SimpleNamespace(cfg_parallel_size=2),
-    )
-    with pytest.raises(ValueError, match="CFG-distilled"):
-        MiniMaxH3Pipeline(od_config=od_config)
-
-
-def test_shape_contract_matches_three_reference_tasks():
-    from vllm_omni.diffusion.models.minimax_h3.time_request import (
-        MINIMAX_H3_SHAPE_PLANNER,
-        minimax_h3_align_frame_count,
-    )
-
-    assert minimax_h3_align_frame_count(round(8.7 * 24)) == 209
-    assert MINIMAX_H3_SHAPE_PLANNER.video_latent_t(209) == 62
-    assert MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(8.7) == 348
-    assert minimax_h3_align_frame_count(round(5.0 * 24)) == 124
-    assert MINIMAX_H3_SHAPE_PLANNER.video_latent_t(124) == 37
-    assert MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(5.0) == 200
-    assert MINIMAX_H3_SHAPE_PLANNER.frame_count_from_video_latent_t(62) == 209
-
-
-def test_shifted_sigma_schedule_matches_reference_values():
-    from vllm_omni.diffusion.models.minimax_h3.time_request import (
-        minimax_h3_time_shift_sigmas,
-    )
-
-    sigmas = minimax_h3_time_shift_sigmas(num_steps=5, shift_scale=12.0)
-
-    assert sigmas == pytest.approx(
-        [1.0, 0.9729729891, 0.9230769277, 0.8000000119, 0.0],
-        abs=1e-7,
-    )
-
-
-def test_reference_image_resize_contract():
-    from PIL import Image
-
-    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
-        _reference_image_shape,
-    )
-
-    assert _reference_image_shape(Image.new("RGB", (1080, 1440))) == (
-        2048,
-        2720,
-    )
-    with pytest.raises(ValueError, match="aspect ratio"):
-        _reference_image_shape(Image.new("RGB", (100, 501)))
-
-
-def test_encoder_forward_uses_hook_compatible_encode_entrypoint():
-    from vllm_omni.diffusion.models.minimax_h3.encoder import (
-        MiniMaxH3Qwen3VLEncoder,
-    )
-
-    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
-    torch.nn.Module.__init__(encoder)
-    expected = torch.ones(2, 3)
-    encoder.encode_ids = Mock(return_value=expected)
-    input_ids = torch.tensor([1, 2])
-    pixel_values = torch.ones(1, 4)
-    image_grid_thw = torch.tensor([[1, 1, 1]])
-
-    actual = encoder(
-        input_ids,
-        pixel_values=pixel_values,
-        image_grid_thw=image_grid_thw,
-    )
-
-    assert actual is expected
-    encoder.encode_ids.assert_called_once_with(
-        input_ids,
-        pixel_values=pixel_values,
-        image_grid_thw=image_grid_thw,
-    )
-
-
-def test_encoder_forward_forwards_video_inputs():
-    from vllm_omni.diffusion.models.minimax_h3.encoder import (
-        MiniMaxH3Qwen3VLEncoder,
-    )
-
-    encoder = object.__new__(MiniMaxH3Qwen3VLEncoder)
-    torch.nn.Module.__init__(encoder)
-    expected = torch.ones(2, 3)
-    encoder.encode_ids = Mock(return_value=expected)
-    input_ids = torch.tensor([1, 2])
-    pixel_values_videos = torch.ones(1, 4)
-    video_grid_thw = torch.tensor([[2, 1, 1]])
-
-    actual = encoder(
-        input_ids,
-        pixel_values_videos=pixel_values_videos,
-        video_grid_thw=video_grid_thw,
-    )
-
-    assert actual is expected
-    encoder.encode_ids.assert_called_once_with(
-        input_ids,
-        pixel_values=None,
-        image_grid_thw=None,
-        pixel_values_videos=pixel_values_videos,
-        video_grid_thw=video_grid_thw,
-    )
-
-
-def test_reference_video_shape_uses_h3_adapt_shape_policy():
-    from vllm_omni.diffusion.models.minimax_h3.reference_video import (
-        _reference_video_shape,
-    )
-
-    assert _reference_video_shape(1280, 720) == (1344, 768)
-    assert _reference_video_shape(3844, 2160) == (1344, 768)
-
-
-def test_text_encoder_stub_constructs_without_group_or_weights():
-    from vllm_omni.diffusion.models.minimax_h3.encoder import (
-        MiniMaxH3Qwen3VLEncoder,
-    )
-
-    encoder = MiniMaxH3Qwen3VLEncoder(
-        "/nonexistent/text_encoder",
-        device=torch.device("cpu"),
-        load_model=False,
-        encoder_group=None,
-    )
-    assert not encoder.is_loaded
-    assert encoder.tp_size == 1
-    # The stub has no parameters, so it never contributes to the runner's
-    # strict missing-parameter check on non-encoder ranks.
-    assert list(encoder.named_parameters()) == []
-
-
-def test_no_offload_keeps_text_encoder_resident():
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    pipeline.od_config = SimpleNamespace(
-        enable_cpu_offload=False,
-        enable_layerwise_offload=False,
-    )
-    pipeline.text_encoder = Mock()
-    expected = torch.ones(2, 3)
-    pipeline.text_encoder.encode_ids.return_value = expected
-    input_ids = torch.tensor([1, 2])
-
-    actual = pipeline._encode_text_hidden(input_ids, {})
-
-    assert actual is expected
-    pipeline.text_encoder.load_to_device.assert_called_once_with()
-    pipeline.text_encoder.offload_to_cpu.assert_not_called()
-
-
-def test_model_offload_uses_hooked_text_encoder_call():
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    pipeline.od_config = SimpleNamespace(
-        enable_cpu_offload=True,
-        enable_layerwise_offload=False,
-    )
-    expected = torch.ones(2, 3)
-    pipeline.text_encoder = Mock(return_value=expected)
-    input_ids = torch.tensor([1, 2])
-    vision_kwargs = {"pixel_values": torch.ones(1, 4)}
-
-    actual = pipeline._encode_text_hidden(input_ids, vision_kwargs)
-
-    assert actual is expected
-    pipeline.text_encoder.assert_called_once_with(input_ids, **vision_kwargs)
-    pipeline.text_encoder.load_to_device.assert_not_called()
-
-
-def test_layerwise_offload_releases_text_encoder():
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    pipeline.od_config = SimpleNamespace(
-        enable_cpu_offload=False,
-        enable_layerwise_offload=True,
-    )
-    pipeline.text_encoder = Mock()
-    expected = torch.ones(2, 3)
-    pipeline.text_encoder.encode_ids.return_value = expected
-
-    actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
-
-    assert actual is expected
-    pipeline.text_encoder.load_to_device.assert_called_once_with()
-    pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
-
-
-def test_video_vae_keeps_reference_fp32_weights(monkeypatch):
-    from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module
-
-    class FakeRemote(torch.nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.model = torch.nn.Linear(1, 1).half()
-
-    monkeypatch.setattr(
-        vae_module,
-        "_load_component_config",
-        lambda _path: {
-            "latent_channels": 1,
-            "latents_mean": [0.0],
-            "latents_std": [1.0],
-        },
-    )
-    monkeypatch.setattr(
-        vae_module,
-        "_load_remote_component",
-        lambda _path, _config: FakeRemote(),
-    )
-
-    video_vae = vae_module.MiniMaxH3VideoVAE(
-        "unused",
-        device=torch.device("cpu"),
-    )
-
-    assert next(video_vae.parameters()).dtype == torch.float32
-
-
-def test_video_vae_encode_uses_configured_parallel_tiling():
-    from vllm_omni.diffusion.models.minimax_h3.vae import (
-        MiniMaxH3VideoVAE,
-    )
-
-    class FakeModel(torch.nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.weight = torch.nn.Parameter(torch.ones(1))
-            self.parallel_tiling = True
-            self.encode_calls = []
-
-        def encode_videos(self, frames, *, use_fp16_latent):
-            assert self.parallel_tiling
-            self.encode_calls.append((frames, use_fp16_latent))
-            return [torch.ones(1, 1, 2, 2, 2)]
-
-    video_vae = object.__new__(MiniMaxH3VideoVAE)
-    torch.nn.Module.__init__(video_vae)
-    video_vae.model = FakeModel()
-    video_vae.config_dict = {
-        "latent_channels": 1,
-        "latents_mean": [0.0],
-        "latents_std": [1.0],
-    }
-
-    rows, shape = video_vae.encode_video("frames")
-
-    assert video_vae.model.parallel_tiling
-    assert video_vae.model.encode_calls == [("frames", True)]
-    assert rows.shape == (2, 4)
-    assert shape == (2, 2, 2)
-
-
-def test_distributed_video_vae_encodes_references_sequentially(monkeypatch):
-    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
-    from vllm_omni.diffusion.models.minimax_h3 import (
-        pipeline_minimax_h3 as pipeline_module,
-    )
-
-    prepared = [
-        {"prepared_path": "video-1.mp4"},
-        {"prepared_path": "video-2.mp4"},
-    ]
-
-    class FakeVideoVAE:
-        def __init__(self):
-            self.calls = []
-
-        def is_distributed_enabled(self):
-            return True
-
-        def encode_video(self, frames):
-            self.calls.append(frames)
-            index = len(self.calls)
-            return torch.full((1, 2), index, dtype=torch.float32), (index, 2, 3)
-
-    pipeline = object.__new__(MiniMaxH3Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    pipeline.device = torch.device("cpu")
-    pipeline.video_vae = FakeVideoVAE()
-
-    monkeypatch.setattr(
-        pipeline_module,
-        "_dit_rank_world",
-        lambda: ("dit-group", 1, 4),
-    )
-
-    def fake_broadcast_object_list(values, *, src, group, device):
-        assert values == [None]
-        assert (src, group, device) == (0, "dit-group", torch.device("cpu"))
-        values[0] = prepared
-
-    monkeypatch.setattr(
-        pipeline_module.dist,
-        "broadcast_object_list",
-        fake_broadcast_object_list,
-    )
-    monkeypatch.setattr(
-        pipeline_module,
-        "load_video_frames",
-        lambda path: f"frames:{path}",
-    )
-
-    rows, shapes = pipeline._encode_video_conditions(None, count=2)
-
-    assert pipeline.video_vae.calls == [
-        "frames:video-1.mp4",
-        "frames:video-2.mp4",
-    ]
-    torch.testing.assert_close(
-        rows,
-        torch.tensor([[1.0, 1.0], [2.0, 2.0]]),
-    )
-    assert shapes == [(1, 2, 3), (2, 2, 3)]