[Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951)

Signed-off-by: wangyu <410167048@qq.com>

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

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

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

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

.buildkite/cuda/test-merge.yml

@@ -8,14 +8,24 @@ steps:
   - group: ":card_index_dividers: Simple Test"
     depends_on: upload-merge-pipeline
     steps:
-      - label: "Simple · Diffusion & Model Executor Test"
+      - label: "Simple · Model Executor Test"
         commands:
-          - "timeout 20m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        mirror_hardwares: l4_1
+
+      - label: "Simple · Diffusion Test"
+        commands:
+          - "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+        mirror_hardwares: l4_1
+
+      - label: "Simple · Engine&Entrypoints Test"
+        commands:
+          - "timeout 20m pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
         mirror_hardwares: l4_1
 
       - label: "Simple · Other Test"
         commands:
-          - "timeout 20m pytest -sv -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
+          - "timeout 20m pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
         mirror_hardwares: l4_1
 
   - group: ":card_index_dividers: Diffusion Test"
@@ -27,12 +37,6 @@ steps:
           - pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model"
         mirror_hardwares: l4_1
 
-      - label: "Diffusion · Offloader Test"
-        timeout_in_minutes: 40
-        commands:
-          - timeout 40m pytest -s -v tests/diffusion/offloader -m "advanced_model and cuda"
-        mirror_hardwares: l4_1
-
   - label: "Distributed Test"
     depends_on: upload-merge-pipeline
     timeout_in_minutes: 60
@@ -135,39 +139,7 @@ steps:
         timeout_in_minutes: 50
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'advanced_model and cuda' --run-level 'advanced_model'
-        agents:
-          queue: "mithril-h100-pool"
-        plugins:
-          - kubernetes:
-              podSpec:
-                containers:
-                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                    resources:
-                      limits:
-                        nvidia.com/gpu: 2
-                    volumeMounts:
-                      - name: devshm
-                        mountPath: /dev/shm
-                      - name: hf-cache
-                        mountPath: /root/.cache/huggingface
-                    env:
-                      - name: HF_HOME
-                        value: /root/.cache/huggingface
-                      - name: HF_TOKEN
-                        valueFrom:
-                          secretKeyRef:
-                            name: hf-token-secret
-                            key: token
-                nodeSelector:
-                  node.kubernetes.io/instance-type: gpu-h100-sxm
-                volumes:
-                  - name: devshm
-                    emptyDir:
-                      medium: Memory
-                  - name: hf-cache
-                    hostPath:
-                      path: /mnt/hf-cache
-                      type: DirectoryOrCreate
+        mirror_hardwares: h100_2
 
       - label: "Omni · MiniCPM-o 4.5 Duplex Test"
         source_file_dependencies:
@@ -210,22 +182,6 @@ steps:
             "
         mirror_hardwares: l4_1
 
-      - label: "TTS · SoulX-Singer Test"
-        source_file_dependencies:
-          - tests/e2e/offline_inference/test_soulxsinger.py
-          - vllm_omni/diffusion/models/soulx_singer/
-          - vllm_omni/deploy/soulxsinger_svc.yaml
-          - vllm_omni/deploy/soulxsinger_svs.yaml
-        timeout_in_minutes: 30
-        commands:
-          - |
-            timeout 30m bash -c "
-              export VLLM_LOGGING_LEVEL=DEBUG
-              export VLLM_WORKER_MULTIPROC_METHOD=spawn
-              pytest -s -v tests/e2e/offline_inference/test_soulxsinger.py -k 'svc or precomputed' -m 'advanced_model and cuda' --run-level 'advanced_model'
-            "
-        mirror_hardwares: l4_1
-
       - label: "Diffusion · Qwen Image Edit Test"
         source_file_dependencies:
           - tests/e2e/online_serving/test_qwen_image_edit.py

.buildkite/cuda/test-nightly.yml

@@ -458,6 +458,11 @@ steps:
     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:
+      - group: ":full_moon: Diffusion · Distributed Test"
+        timeout_in_minutes: 30
+        commands:
+          - timeout 30m pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda' --run-level "full_model"
+        mirror_hardwares: l4_4
 
       - label: ":full_moon: Diffusion · Offloader Test"
         timeout_in_minutes: 40
@@ -473,6 +478,7 @@ steps:
       - 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
 

.buildkite/cuda/test-ready.yml

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

docs/contributing/ci/test_system_overview.md

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

docs/contributing/ci/test_writing_guide.md

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

docs/design/feature/tensor_parallel.md

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

examples/model_display_names.yml

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

recipes/JD/JoyAI-VL-Interaction.md

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

recipes/NVIDIA/GR00T-N1.7.md

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

tests/config/test_config_factory.py

@@ -1145,7 +1145,7 @@ class TestResolveScheduler:
 
 class TestDeployConfigLoading:
     def test_load_minicpmo_duplex_deploy_config(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "minicpmo_4_5_duplex.yaml"
+        deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
 
         deploy = load_deploy_config(deploy_path)
         pipeline = resolve_pipeline_config("minicpmo_4_5")
@@ -1182,7 +1182,7 @@ class TestDeployConfigLoading:
         stage2_devices: str,
         stage1_replicas: int,
     ):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / filename
+        deploy_path = Path(get_deploy_config_path(filename))
 
         deploy = load_deploy_config(deploy_path)
         pipeline = resolve_pipeline_config("minicpmo_4_5")
@@ -1226,7 +1226,7 @@ stages:
         assert {s.yaml_engine_args.get("custom_voice_dir") for s in stages} == {str(custom_voice_dir)}
 
     def test_load_qwen3_omni_moe_deploy_config(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         deploy = load_deploy_config(deploy_path)
         assert len(deploy.stages) == 3
         assert deploy.async_chunk is True
@@ -1234,7 +1234,7 @@ stages:
         assert deploy.platforms is not None
 
     def test_load_voxtral_tts_deploy_config_schema_fields(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "voxtral_tts.yaml"
+        deploy_path = Path(get_deploy_config_path("voxtral_tts.yaml"))
         deploy = load_deploy_config(deploy_path)
         assert deploy.stages[0].config_format == "mistral"
         assert deploy.stages[0].load_format == "mistral"
@@ -1244,7 +1244,7 @@ stages:
         )
 
     def test_load_ming_flash_omni_deploy_config_schema_fields(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml"
+        deploy_path = Path(get_deploy_config_path("ming_flash_omni.yaml"))
         deploy = load_deploy_config(deploy_path)
         assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}}
         assert "compilation_config" not in deploy.stages[0].engine_extras
@@ -1384,7 +1384,7 @@ stages:
         assert not list(stage_config_dir.glob("*.yaml"))
 
     def test_merge_pipeline_deploy(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -1411,7 +1411,7 @@ stages:
         )
         assert isinstance(pipeline, PipelineConfig)
 
-        base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not base.exists():
             pytest.skip("Deploy config not found")
 
@@ -1831,7 +1831,7 @@ class TestQwen3TTSPipeline:
 
     def test_per_stage_model_arch_flows_through_merge(self, tmp_path):
         """Verify the new ps.model_arch override survives merge_pipeline_deploy."""
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_tts.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_tts.yaml"))
         if not deploy_path.exists():
             pytest.skip("qwen3_tts deploy yaml not found")
 
@@ -1955,7 +1955,7 @@ class TestMingFlashOmniPipeline:
 
     def test_full_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni.yaml parses and merges with the registered pipeline."""
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml"
+        deploy_path = Path(get_deploy_config_path("ming_flash_omni.yaml"))
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni deploy yaml not found")
 
@@ -1975,7 +1975,7 @@ class TestMingFlashOmniPipeline:
 
     def test_tts_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni_tts.yaml parses and routes to the TTS-only pipeline."""
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_tts.yaml"
+        deploy_path = Path(get_deploy_config_path("ming_flash_omni_tts.yaml"))
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni_tts deploy yaml not found")
 
@@ -2013,7 +2013,7 @@ class TestMingFlashOmniPipeline:
 
     def test_thinker_only_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni_thinker_only.yaml parses and routes to the thinker-only pipeline."""
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_thinker_only.yaml"
+        deploy_path = Path(get_deploy_config_path("ming_flash_omni_thinker_only.yaml"))
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni_thinker_only deploy yaml not found")
 
@@ -2073,7 +2073,7 @@ class TestMingFlashOmniPipeline:
 
     def test_image_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni_image.yaml parses and routes to the image pipeline."""
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_image.yaml"
+        deploy_path = Path(get_deploy_config_path("ming_flash_omni_image.yaml"))
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni_image deploy yaml not found")
 
@@ -2135,7 +2135,7 @@ class TestBaseConfigInheritance:
 
     def test_pure_inheritance_overlay(self, tmp_path):
         """An overlay with only ``base_config`` inherits everything."""
-        base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not base.exists():
             pytest.skip("Base deploy config not found")
 
@@ -2148,7 +2148,7 @@ class TestBaseConfigInheritance:
 
     def test_single_field_overlay(self, tmp_path):
         """An overlay overriding one stage field merges with the base."""
-        base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        base = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not base.exists():
             pytest.skip("Base deploy config not found")
 
@@ -2165,7 +2165,7 @@ class TestPlatformOverrides:
     """Test platform-specific deploy config overrides."""
 
     def test_npu_overrides(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -2179,7 +2179,7 @@ class TestPlatformOverrides:
         assert deploy.stages[2].enforce_eager is False
 
     def test_xpu_overrides(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -2191,7 +2191,7 @@ class TestPlatformOverrides:
         assert deploy.stages[0].engine_extras.get("max_cudagraph_capture_size") == 0
 
     def test_unknown_platform_noop(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -2236,7 +2236,7 @@ class TestCLIOverrideFlow:
     """Test --stage-overrides JSON merge into StageConfig."""
 
     def test_stage_overrides_merge(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -2254,7 +2254,7 @@ class TestCLIOverrideFlow:
         assert stages[0].runtime_overrides["gpu_memory_utilization"] == 0.5
 
     def test_global_override_applies_to_all(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 
@@ -2274,7 +2274,7 @@ class TestCLIOverrideFlow:
 
 class TestAuraOmniDeploy:
     def test_aura_omni_deploy_forces_pipeline_override(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"
+        deploy_path = Path(get_deploy_config_path("aura_omni.yaml"))
         deploy = load_deploy_config(deploy_path)
 
         assert deploy.pipeline == "aura_omni"
@@ -2285,7 +2285,7 @@ class TestAuraOmniDeploy:
         stages, _ = StageConfigFactory._create_legacy_from_registry(
             pipeline_cfg,
             cli_overrides={},
-            deploy_config_path=str(Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"),
+            deploy_config_path=get_deploy_config_path("aura_omni.yaml"),
         )
 
         assert [stage.model_stage for stage in stages] == [
@@ -2599,7 +2599,12 @@ class TestSentinelDefaultPrecedence:
         from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy
 
         pipeline_path = (
-            Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "models" / "qwen3_omni" / "pipeline.py"
+            Path(__file__).parent.parent.parent
+            / "vllm_omni"
+            / "model_executor"
+            / "models"
+            / "qwen3_omni"
+            / "pipeline.py"
         )
         pipeline = runpy.run_path(str(pipeline_path))["QWEN3_OMNI_PIPELINE"]
 
@@ -2669,7 +2674,7 @@ class TestSamplingConstraintsPrecedence:
     """Test that pipeline sampling_constraints override deploy defaults."""
 
     def test_constraints_win(self):
-        deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
+        deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
 

tests/conftest.py

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

tests/core/sched/test_omni_scheduler_mixin_timeouts.py

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

tests/core/test_prefix_cache.py

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

tests/diffusion/ar_diffusion/test_engine_routing.py

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

tests/diffusion/ar_diffusion/test_kv_cache.py

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

tests/diffusion/ar_diffusion/test_kv_orchestration.py

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

tests/diffusion/ar_diffusion/test_paged_attention.py

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

tests/diffusion/ar_diffusion/test_pipeline_bridge.py

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

tests/diffusion/ar_diffusion/test_slot_mapping.py

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

tests/diffusion/attention/test_attention_config.py

@@ -288,6 +288,13 @@ class TestBuildAttentionConfig:
 class TestOmniDiffusionConfigAttentionParsing:
     """Test OmniDiffusionConfig attention shorthand and structured config."""
 
+    @pytest.fixture(autouse=True)
+    def _clear_diffusion_attention_backend_env(self, monkeypatch):
+        # OmniDiffusionConfig.__post_init__ applies DIFFUSION_ATTENTION_BACKEND via
+        # build_attention_config(); clear it so these tests assert config defaults,
+        # not whatever the process inherited from CI / sibling tests.
+        monkeypatch.delenv("DIFFUSION_ATTENTION_BACKEND", raising=False)
+
     def test_diffusion_attention_backend_sets_default(self):
         config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="SAGE_ATTN")
         assert isinstance(config.diffusion_attention_config, AttentionConfig)

tests/diffusion/attention/test_attention_sp.py

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

tests/diffusion/attention/test_flash_attn.py

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

tests/diffusion/attention/test_piecewise_attn.py

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

tests/diffusion/attention/test_sage_attn3.py

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

tests/diffusion/attention/test_ulysses_uaa.py

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

tests/diffusion/distributed/test_comm.py

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

tests/diffusion/distributed/test_pipeline_parallel.py

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

tests/diffusion/distributed/test_tensor_parallel.py

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

tests/diffusion/distributed/test_vae_decode_parallelism.py

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

tests/diffusion/hooks/test_hook_registry.py

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

tests/diffusion/layers/test_rotary_emb_equivalence.py

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

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -9,6 +9,8 @@ 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/sensenova_u1/test_sensenova_u1_cfg.py

@@ -6,6 +6,8 @@ import torch
 
 from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def _pipeline_without_init() -> SenseNovaU1Pipeline:
     return object.__new__(SenseNovaU1Pipeline)

tests/diffusion/offloader/test_diffusion_cpu_offload.py

@@ -5,7 +5,6 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.helpers import skip_if_gated_repo_inaccessible
 from tests.helpers.env import DeviceMemoryMonitor
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner
@@ -16,32 +15,26 @@ AUDIO_MODEL = {
     "stabilityai/stable-audio-open-1.0": {"cuda": 100, "rocm": None},
 }
 
-IMAGE_VIDEO_MODELS = {
+IMAGE_MODELS = {
     "riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100},
-    "Tongyi-MAI/Z-Image-Turbo": {"cuda": 2500, "rocm": 2100},
-    "OmniGen2/OmniGen2": {"cuda": 2500, "rocm": 2100},
 }
 
-MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
+MODELS = {**AUDIO_MODEL, **IMAGE_MODELS}
 
 MODEL_MARKS = {
     "riverclouds/qwen_image_random": pytest.mark.core_model,
-    "Tongyi-MAI/Z-Image-Turbo": pytest.mark.advanced_model,
     "stabilityai/stable-audio-open-1.0": pytest.mark.full_model,
-    "OmniGen2/OmniGen2": pytest.mark.full_model,
 }
 
 _GATED_MODELS = {"stabilityai/stable-audio-open-1.0"}
 
-# Aliased for backward compatibility (imported by test_diffusion_layerwise_offload.py).
-_skip_if_gated_repo_inaccessible = skip_if_gated_repo_inaccessible
 
 AUDIO_MODEL_PARAMS = {
     "runner_params": {},
     "sampler_params": {},
 }
 
-IMAGE_VIDEO_MODELS_PARAMS = {
+IMAGE_MODEL_PARAMS = {
     "runner_params": {},
     "sampler_params": {
         "height": 256,
@@ -60,7 +53,7 @@ def inference(model_name: str, offload: bool = True):
     if model_name in AUDIO_MODEL:
         params = AUDIO_MODEL_PARAMS
     else:
-        params = IMAGE_VIDEO_MODELS_PARAMS
+        params = IMAGE_MODEL_PARAMS
 
     with OmniRunner(
         model_name,
@@ -111,19 +104,12 @@ def check_audio_determinism(audio1, audio2, atol=1e-2):
     [pytest.param(name, marks=MODEL_MARKS[name]) for name in MODELS],
 )
 def test_cpu_offload_diffusion_model(model_name: str):
-    if model_name == "OmniGen2/OmniGen2":
-        pytest.skip("issue #4537")
-    if model_name in _GATED_MODELS:
-        _skip_if_gated_repo_inaccessible(model_name)
     try:
         offload_peak_memory, output_offload = inference(model_name, offload=True)
         cleanup_dist_env_and_memory()
         no_offload_peak_memory, output_no_offload = inference(model_name, offload=False)
     except ValueError as exc:
-        # omni_snapshot_download wraps GatedRepoError in a ValueError.
-        # If the pre-flight guard above did not catch it (e.g. partial
-        # HF_TOKEN where config.json is accessible but weight shards are
-        # blocked), skip instead of failing.
+        # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing.
         if "Access to model" in str(exc) and "is restricted" in str(exc):
             pytest.skip(
                 f"Skipping: gated HF repo {model_name!r} inaccessible "

tests/diffusion/offloader/test_diffusion_layerwise_offload.py

@@ -3,11 +3,8 @@ import pytest
 import torch
 from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
 
-from tests.diffusion.offloader.test_diffusion_cpu_offload import (
-    _GATED_MODELS,
-    _skip_if_gated_repo_inaccessible,
-)
 from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
@@ -23,6 +20,19 @@ IMAGE_VIDEO_MODELS = {
 
 MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}
 
+LAYERWISE_MODEL_CASES = [
+    pytest.param(
+        "stabilityai/stable-audio-open-1.0",
+        marks=[pytest.mark.full_model, pytest.mark.diffusion],
+        id="stable_audio",
+    ),
+    pytest.param(
+        "riverclouds/qwen_image_random",
+        marks=[pytest.mark.core_model, pytest.mark.diffusion],
+        id="qwen_image_random",
+    ),
+]
+
 AUDIO_MODEL_PARAMS = {
     "runner_params": {},
     "sampler_params": {},
@@ -91,7 +101,8 @@ def run_inference(
     return peak, output
 
 
-@pytest.mark.parametrize("model_name", list(MODELS.keys()))
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("model_name", LAYERWISE_MODEL_CASES)
 def test_layerwise_offload_diffusion_model(model_name: str):
     """Test that layerwise offloading reduces GPU memory usage.
 
@@ -100,8 +111,6 @@ def test_layerwise_offload_diffusion_model(model_name: str):
     offloader keeps only a single transformer block on GPU at a time, with
     prefetching for compute-memory overlap.
     """
-    if model_name in _GATED_MODELS:
-        _skip_if_gated_repo_inaccessible(model_name)
     try:
         # Run without layerwise offloading (baseline)
         no_offload_peak_memory, output_no_offload = run_inference(model_name, layerwise_offload=False)
@@ -111,10 +120,7 @@ def test_layerwise_offload_diffusion_model(model_name: str):
         layerwise_offload_peak_memory, output_offload = run_inference(model_name, layerwise_offload=True)
         cleanup_dist_env_and_memory()
     except ValueError as exc:
-        # omni_snapshot_download wraps GatedRepoError in a ValueError.
-        # If the pre-flight guard above did not catch it (e.g. partial
-        # HF_TOKEN where config.json is accessible but weight shards are
-        # blocked), skip instead of failing.
+        # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing.
         if "Access to model" in str(exc) and "is restricted" in str(exc):
             pytest.skip(
                 f"Skipping: gated HF repo {model_name!r} inaccessible "

tests/diffusion/quantization/test_gguf_diffusion.py

@@ -5,10 +5,10 @@
 Validates that GGUF-quantized diffusion models generate valid images and
 use less peak GPU memory than BF16 baseline.
 
-Requires vllm-omni to be installed alongside the plugin.
+Requires vllm-omni to be installed alongside the vllm-gguf-plugin package.
 
 Usage:
-    pytest tests/diffusion/test_gguf_memory.py -v
+    pytest tests/diffusion/quantization/test_gguf_diffusion.py -v
 """
 
 from __future__ import annotations
@@ -112,7 +112,6 @@ def _generate_single_stage_image(
 @hardware_test(res={"cuda": "L4"})
 @pytest.mark.full_model
 @pytest.mark.diffusion
-@pytest.mark.slow
 @pytest.mark.parametrize("model", [Z_IMAGE_CONFIG, FLUX_CONFIG], ids=["Z-Image-Turbo", "FLUX.2-klein"])
 def test_single_stage_diffusion_gguf(model: DiffusionGGUFTestConfig, monkeypatch: pytest.MonkeyPatch) -> None:
     """Z-Image-Turbo GGUF generates valid images and uses less memory than BF16."""

tests/diffusion/quantization/test_quantization_quality.py

@@ -431,6 +431,8 @@ def _free_gpu_memory():
 # ---------------------------------------------------------------------------
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch):
     from benchmarks.diffusion.quantization_quality import _generate_image as benchmark_generate_image
     from vllm_omni.outputs import OmniRequestOutput

tests/diffusion/quantization/test_trajectory_similarity_tool.py

@@ -22,6 +22,8 @@ from vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity import
     summarize_output_image_metrics,
 )
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
 
 def test_compute_tensor_metrics_identical_tensors():
     metrics = compute_tensor_metrics(torch.ones(2, 3), torch.ones(2, 3))

tests/diffusion/test_diffusion_config_fields.py

@@ -29,7 +29,7 @@ def test_diffusion_stage_configs_only_contain_valid_fields():
     Regression test for https://github.com/vllm-project/vllm-omni/issues/2563
     """
     # Scan both main configs and test configs
-    repo_root = Path(__file__).parent.parent
+    repo_root = Path(__file__).parent.parent.parent
     config_dirs = [
         repo_root / "vllm_omni" / "model_executor" / "stage_configs",
     ]

tests/diffusion/test_diffusion_engine.py

@@ -14,6 +14,7 @@ import torch
 from pytest_mock import MockerFixture
 
 import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.diffusion_engine import (
     DiffusionEngine,
@@ -30,6 +31,8 @@ from vllm_omni.diffusion.sched.interface import (
 )
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
 
 @dataclass
 class DiffusionSchedulerOutput:
@@ -56,6 +59,8 @@ class MockScheduler:
     def __init__(self):
         self._waiting_queue = []
         self._step_id = 0
+        # Match RequestScheduler API used by request-batch admission wait.
+        self.max_num_running_reqs = 5
 
     def add_request(self, request):
         self._waiting_queue.append(request)
@@ -64,6 +69,12 @@ class MockScheduler:
     def has_requests(self):
         return len(self._waiting_queue) > 0
 
+    def num_waiting_requests(self) -> int:
+        return len(self._waiting_queue)
+
+    def num_running_requests(self) -> int:
+        return 0
+
     def schedule(self) -> DiffusionSchedulerOutput:
         if not self._waiting_queue:
             return DiffusionSchedulerOutput(step_id=self._step_id)
@@ -582,9 +593,7 @@ def test_move_tensor_tree_returns_non_tensor_values_unchanged() -> None:
     assert moved is value
 
 
-@pytest.mark.diffusion
-@pytest.mark.cuda
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None:
     tensor = torch.arange(8, dtype=torch.float32, device="cuda")
     other = torch.arange(4, dtype=torch.int64, device="cuda")
@@ -610,6 +619,7 @@ async def _consume_final_output(generator):
     return final_output
 
 
+@pytest.mark.cpu
 @pytest.mark.asyncio
 async def test_async_add_req_and_stream_response():
     engine = object.__new__(DiffusionEngine)
@@ -621,13 +631,24 @@ async def test_async_add_req_and_stream_response():
     engine._cv = threading.Condition(engine._rpc_lock)
     engine._init_lock = asyncio.Lock()
     engine._closed = False
-    engine.od_config = SimpleNamespace(streaming_output=False)
+    # Enable admission wait so concurrent adds land in one schedule wave;
+    # otherwise the first request can execute alone (~1s) and the rest in a
+    # second wave (~2s), failing the latency-spread assertion.
+    engine.od_config = SimpleNamespace(
+        streaming_output=False,
+        request_batch_max_wait_ms=500.0,
+    )
     engine._loop_started = False
     engine.main_loop = None
-    engine.supports_request_batch = False
+    engine.supports_request_batch = True
+    engine.step_execution = False
     engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
 
-    engine._finalize_finished_request = lambda rid, out, err: out.result
+    def _finalize(rid, out, err=None):
+        # Stream consumers stop on ``finished``; keep result_data for assertions.
+        return SimpleNamespace(result_data=out.result.result_data, finished=True)
+
+    engine._finalize_finished_request = _finalize
 
     def mock_execute_batch(sched_output):
         request_ids = sched_output.scheduled_request_ids

tests/diffusion/test_diffusion_engine_metrics.py

@@ -12,6 +12,11 @@ from __future__ import annotations
 import ast
 import os
 
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
 _ENGINE_PATH = os.path.normpath(
     os.path.join(
         os.path.dirname(__file__),

tests/docs/test_generate_examples.py

@@ -6,6 +6,9 @@ from pathlib import Path
 
 import pytest
 
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
 ROOT_DIR = Path(__file__).parents[2]
 HOOK_PATH = ROOT_DIR / "docs/mkdocs/hooks/generate_examples.py"
 SPEC = importlib.util.spec_from_file_location("generate_examples", HOOK_PATH)

tests/dreamzero/openpi_client_helper.py

@@ -1,307 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import uuid
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-
-try:
-    import cv2
-except ImportError:  # pragma: no cover - optional e2e dependency
-    cv2 = None
-
-try:
-    import websockets.sync.client as websockets_client
-except ImportError:  # pragma: no cover - optional e2e dependency
-    websockets_client = None
-
-try:
-    from openpi_client import msgpack_numpy
-except ImportError:  # pragma: no cover - optional e2e dependency
-    msgpack_numpy = None
-
-PING_INTERVAL_SECS = 300
-PING_TIMEOUT_SECS = 3600
-DEFAULT_HOST = "127.0.0.1"
-DEFAULT_PORT = 8000
-DEFAULT_PATH = "/v1/realtime/robot/openpi"
-DEFAULT_PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
-ACTION_HORIZON = 24
-DEFAULT_ACTION_DIM = 8
-RELATIVE_OFFSETS = [-23, -16, -8, 0]
-CAMERA_FILES = {
-    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
-    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
-    "observation/wrist_image_left": "wrist_image_left.mp4",
-}
-
-
-def _decode_action_response(response: bytes | str) -> np.ndarray:
-    if isinstance(response, str):
-        raise RuntimeError(f"Inference failed: {response}")
-    decoded = msgpack_numpy.unpackb(response)
-    if isinstance(decoded, dict) and decoded.get("type") == "error":
-        message = decoded.get("message", decoded)
-        raise RuntimeError(f"Inference failed: {message}")
-    return np.asarray(decoded, dtype=np.float32)
-
-
-def require_dependencies() -> None:
-    missing = []
-    if cv2 is None:
-        missing.append("opencv-python")
-    if websockets_client is None:
-        missing.append("websockets")
-    if msgpack_numpy is None:
-        missing.append("openpi-client")
-    if missing:
-        raise ModuleNotFoundError(f"DreamZero OpenPI test dependencies are missing: {', '.join(missing)}")
-
-
-@dataclass(frozen=True)
-class DreamZeroServerMetadata:
-    image_resolution: tuple[int, int]
-    n_external_cameras: int
-    needs_wrist_camera: bool
-    needs_stereo_camera: bool
-    needs_session_id: bool
-    action_space: str
-
-    @classmethod
-    def from_dict(cls, payload: dict[str, Any]) -> DreamZeroServerMetadata:
-        required_keys = (
-            "image_resolution",
-            "n_external_cameras",
-            "needs_wrist_camera",
-            "needs_stereo_camera",
-            "needs_session_id",
-            "action_space",
-        )
-        missing_keys = [key for key in required_keys if key not in payload]
-        if missing_keys:
-            raise ValueError(f"Missing DreamZero metadata keys: {missing_keys}")
-
-        image_resolution = payload["image_resolution"]
-        if not isinstance(image_resolution, (list, tuple)) or len(image_resolution) != 2:
-            raise ValueError(f"Invalid image_resolution: {image_resolution!r}")
-
-        return cls(
-            image_resolution=(int(image_resolution[0]), int(image_resolution[1])),
-            n_external_cameras=int(payload["n_external_cameras"]),
-            needs_wrist_camera=bool(payload["needs_wrist_camera"]),
-            needs_stereo_camera=bool(payload["needs_stereo_camera"]),
-            needs_session_id=bool(payload["needs_session_id"]),
-            action_space=str(payload["action_space"]),
-        )
-
-
-class OpenPIWebsocketClient:
-    def __init__(
-        self,
-        *,
-        host: str = DEFAULT_HOST,
-        port: int = DEFAULT_PORT,
-        path: str = DEFAULT_PATH,
-    ) -> None:
-        require_dependencies()
-        self._uri = f"ws://{host}:{port}{path}"
-        self._packer = msgpack_numpy.Packer()
-        self._ws, self._server_metadata = self._connect()
-
-    def _connect(self):
-        conn = websockets_client.connect(
-            self._uri,
-            compression=None,
-            max_size=None,
-            ping_interval=PING_INTERVAL_SECS,
-            ping_timeout=PING_TIMEOUT_SECS,
-        )
-        metadata = msgpack_numpy.unpackb(conn.recv())
-        if not isinstance(metadata, dict):
-            raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
-        return conn, metadata
-
-    def get_server_metadata(self) -> dict[str, Any]:
-        return dict(self._server_metadata)
-
-    def infer(self, obs: dict[str, Any]) -> np.ndarray:
-        payload = dict(obs)
-        payload["endpoint"] = "infer"
-        self._ws.send(self._packer.pack(payload))
-        response = self._ws.recv()
-        return _decode_action_response(response)
-
-    def reset(self, reset_info: dict[str, Any] | None = None) -> str:
-        payload = dict(reset_info or {})
-        payload["endpoint"] = "reset"
-        self._ws.send(self._packer.pack(payload))
-        response = self._ws.recv()
-        if isinstance(response, str):
-            return response
-        decoded = msgpack_numpy.unpackb(response)
-        if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
-            raise RuntimeError(f"Unexpected reset response: {decoded!r}")
-        return str(decoded["status"])
-
-    def close(self) -> None:
-        self._ws.close()
-
-
-def load_all_frames(video_path: Path) -> np.ndarray:
-    require_dependencies()
-    cap = cv2.VideoCapture(str(video_path))
-    frames = []
-    while True:
-        ok, frame = cap.read()
-        if not ok:
-            break
-        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-    cap.release()
-    if not frames:
-        raise RuntimeError(f"No frames loaded from {video_path}")
-    return np.stack(frames, axis=0)
-
-
-def load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
-    camera_frames: dict[str, np.ndarray] = {}
-    for camera_key, file_name in CAMERA_FILES.items():
-        video_path = video_dir / file_name
-        if not video_path.exists():
-            raise FileNotFoundError(f"Missing DreamZero test asset: {video_path}")
-        camera_frames[camera_key] = load_all_frames(video_path)
-    return camera_frames
-
-
-def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
-    chunks: list[list[int]] = []
-    current_frame = 23
-    for _ in range(num_chunks):
-        indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
-        if indices[-1] >= total_frames:
-            break
-        chunks.append(indices)
-        current_frame += ACTION_HORIZON
-    return chunks
-
-
-def make_obs_from_video(
-    camera_frames: dict[str, np.ndarray],
-    frame_indices: list[int],
-    *,
-    prompt: str,
-    session_id: str,
-) -> dict[str, Any]:
-    obs: dict[str, Any] = {}
-    for camera_key, all_frames in camera_frames.items():
-        selected = all_frames[frame_indices]
-        obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
-
-    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
-    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
-    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
-    obs["prompt"] = prompt
-    obs["session_id"] = session_id
-    return obs
-
-
-def build_demo_observations(
-    camera_frames: dict[str, np.ndarray],
-    *,
-    prompt: str,
-    session_id: str,
-    num_chunks: int = 2,
-) -> list[dict[str, Any]]:
-    if num_chunks < 1:
-        raise ValueError("num_chunks must be at least 1")
-
-    total_frames = min(frames.shape[0] for frames in camera_frames.values())
-    observations = [
-        make_obs_from_video(
-            camera_frames,
-            [0],
-            prompt=prompt,
-            session_id=session_id,
-        )
-    ]
-    for indices in build_frame_schedule(total_frames, num_chunks - 1):
-        observations.append(
-            make_obs_from_video(
-                camera_frames,
-                indices,
-                prompt=prompt,
-                session_id=session_id,
-            )
-        )
-    return observations
-
-
-def validate_session_result(
-    result: dict[str, Any],
-    *,
-    expected_action_horizon: int = ACTION_HORIZON,
-    expected_action_dim: int = DEFAULT_ACTION_DIM,
-) -> None:
-    metadata = DreamZeroServerMetadata.from_dict(result["metadata"])
-    if metadata.image_resolution != (180, 320):
-        raise AssertionError(f"Unexpected image_resolution: {metadata.image_resolution}")
-    if metadata.n_external_cameras != 2:
-        raise AssertionError(f"Unexpected n_external_cameras: {metadata.n_external_cameras}")
-    if not metadata.needs_wrist_camera:
-        raise AssertionError("DreamZero test expects wrist camera metadata")
-    if metadata.action_space != "joint_position":
-        raise AssertionError(f"Unexpected action_space: {metadata.action_space}")
-
-    actions = result["actions"]
-    if len(actions) != 3:
-        raise AssertionError(f"Expected 3 action tensors, got {len(actions)}")
-    for index, action in enumerate(actions):
-        if action.shape != (expected_action_horizon, expected_action_dim):
-            raise AssertionError(
-                f"Action {index} shape mismatch: expected "
-                f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
-            )
-        if not np.isfinite(action).all():
-            raise AssertionError(f"Action {index} contains non-finite values")
-
-    if result["reset_status"] != "reset successful":
-        raise AssertionError(f"Unexpected reset status: {result['reset_status']!r}")
-
-
-def run_policy_session(
-    *,
-    host: str = DEFAULT_HOST,
-    port: int = DEFAULT_PORT,
-    path: str = DEFAULT_PATH,
-    video_dir: Path,
-    prompt: str = DEFAULT_PROMPT,
-    session_id: str | None = None,
-    num_chunks: int = 2,
-) -> dict[str, Any]:
-    session_id = session_id or str(uuid.uuid4())
-    camera_frames = load_camera_frames(video_dir)
-    observations = build_demo_observations(
-        camera_frames,
-        prompt=prompt,
-        session_id=session_id,
-        num_chunks=num_chunks,
-    )
-
-    client = OpenPIWebsocketClient(host=host, port=port, path=path)
-    try:
-        metadata = client.get_server_metadata()
-        actions = [client.infer(obs) for obs in observations]
-        reset_status = client.reset({})
-        actions.append(client.infer(observations[0]))
-        return {
-            "metadata": metadata,
-            "actions": actions,
-            "reset_status": reset_status,
-            "session_id": session_id,
-        }
-    finally:
-        client.close()

tests/dreamzero/test_openpi_client_helper.py

@@ -1,46 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import numpy as np
-import pytest
-
-from tests.dreamzero import openpi_client_helper
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class FakeMsgpackNumpy:
-    payloads = {}
-
-    @staticmethod
-    def packb(obj):
-        key = str(len(FakeMsgpackNumpy.payloads)).encode()
-        FakeMsgpackNumpy.payloads[key] = obj
-        return key
-
-    @staticmethod
-    def unpackb(data):
-        return FakeMsgpackNumpy.payloads[data]
-
-
-def test_decode_action_response_surfaces_structured_error(monkeypatch):
-    monkeypatch.setattr(openpi_client_helper, "msgpack_numpy", FakeMsgpackNumpy)
-    payload = FakeMsgpackNumpy.packb(
-        {
-            "type": "error",
-            "message": "Internal inference error",
-        }
-    )
-
-    with pytest.raises(RuntimeError, match="Internal inference error"):
-        openpi_client_helper._decode_action_response(payload)
-
-
-def test_decode_action_response_converts_action_payload_to_float32(monkeypatch):
-    monkeypatch.setattr(openpi_client_helper, "msgpack_numpy", FakeMsgpackNumpy)
-    payload = FakeMsgpackNumpy.packb(np.asarray([[1.0, 2.0]], dtype=np.float64))
-
-    actions = openpi_client_helper._decode_action_response(payload)
-
-    assert actions.dtype == np.float32
-    np.testing.assert_array_equal(actions, np.asarray([[1.0, 2.0]], dtype=np.float32))

tests/dreamzero/upstream/openpi_test_client_ar.py

@@ -1,381 +0,0 @@
-#!/usr/bin/env python3
-"""
-Copied from the DreamZero repository's `test_client_AR.py`:
-https://github.com/dreamzero0/dreamzero/blob/main/test_client_AR.py
-
-Kept here for end-to-end compatibility / parity testing against the vLLM
-OpenPI server.
-
-Test client for AR_droid policy server using roboarena interface.
-
-Sends real video frames from debug_image/ directory instead of zero dummy images.
-
-Frame schedule (matching debug_inference.py):
-  - Step 0 (initial): send frame [0]             (1 frame, H W 3)
-  - Step 1: send frames [0, 7, 15, 23]           (4 frames, 4 H W 3)
-  - Step 2: send frames [24, 31, 39, 47]         (4 frames)
-  - Step 3: send frames [48, 55, 63, 71]         (4 frames)
-  - ...
-
-Expected server configuration:
-    - image_resolution: (180, 320)
-    - n_external_cameras: 2
-    - needs_wrist_camera: True
-    - action_space: "joint_position"
-
-Usage:
-    # Start server with roboarena interface:
-    torchrun --nproc_per_node=8 socket_test_optimized_AR.py --port 8000
-
-    # Run against the original DreamZero websocket server (default path):
-    python test_client_AR.py --host <server_host> --port 8000
-
-    # Run against the vLLM OpenPI server:
-    python test_client_AR.py --host <server_host> --port 8000 --path /v1/realtime/robot/openpi
-
-    # Use zero images instead of real video (old behavior):
-    python test_client_AR.py --host <server_host> --port 8000 --use-zero-images
-"""
-
-import argparse
-import logging
-import os
-import sys
-import time
-import uuid
-from pathlib import Path
-
-import cv2
-import numpy as np
-from openpi_client import msgpack_numpy
-
-_DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO")
-DREAMZERO_REPO = Path(_DREAMZERO_REPO_ENV).expanduser() if _DREAMZERO_REPO_ENV else None
-
-
-def _import_upstream_policy_modules():
-    if DREAMZERO_REPO is None:
-        raise ImportError("Set DREAMZERO_REPO to an upstream DreamZero checkout before using this helper.")
-    if DREAMZERO_REPO.exists() and str(DREAMZERO_REPO) not in sys.path:
-        sys.path.insert(0, str(DREAMZERO_REPO))
-
-    import eval_utils.policy_server as policy_server
-    from eval_utils.policy_client import WebsocketClientPolicy
-
-    return policy_server, WebsocketClientPolicy
-
-
-policy_server, WebsocketClientPolicy = _import_upstream_policy_modules()
-
-VIDEO_DIR = os.environ.get(
-    "DREAMZERO_VIDEO_DIR",
-    str(DREAMZERO_REPO / "debug_image") if DREAMZERO_REPO is not None else "debug_image",
-)
-
-# roboarena key -> video filename
-CAMERA_FILES = {
-    "observation/exterior_image_0_left": "exterior_image_1_left.mp4",
-    "observation/exterior_image_1_left": "exterior_image_2_left.mp4",
-    "observation/wrist_image_left": "wrist_image_left.mp4",
-}
-
-# Frame schedule constants (matching debug_inference.py)
-RELATIVE_OFFSETS = [-23, -16, -8, 0]
-ACTION_HORIZON = 24
-DEFAULT_WEBSOCKET_PATH = ""
-
-
-class OpenPIWebsocketClientPolicy(WebsocketClientPolicy):
-    """DreamZero websocket client with a configurable path suffix.
-
-    The original DreamZero client connects to ``ws://host:port``.
-    vLLM serves the compatible robot policy endpoint at
-    ``/v1/realtime/robot/openpi`` when ``path`` is set accordingly.
-    """
-
-    def __init__(
-        self,
-        host: str = "0.0.0.0",
-        port: int = 8000,
-        path: str = DEFAULT_WEBSOCKET_PATH,
-    ) -> None:
-        self._uri = f"ws://{host}:{port}{path}"
-        self._packer = msgpack_numpy.Packer()
-        self._ws, self._server_metadata = self._wait_for_server()
-
-
-def load_all_frames(video_path: str) -> np.ndarray:
-    """Load all frames from a video file. Returns (N, H, W, 3) uint8 array (RGB)."""
-    cap = cv2.VideoCapture(video_path)
-    frames = []
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-    cap.release()
-    if not frames:
-        raise RuntimeError(f"No frames loaded from {video_path}")
-    return np.stack(frames, axis=0)
-
-
-def load_camera_frames() -> dict[str, np.ndarray]:
-    """Load all video frames for each camera from the debug_image/ directory.
-
-    Returns:
-        Dict mapping roboarena camera keys to (N, H, W, 3) uint8 arrays.
-    """
-    camera_frames: dict[str, np.ndarray] = {}
-    for cam_key, fname in CAMERA_FILES.items():
-        path = os.path.join(VIDEO_DIR, fname)
-        camera_frames[cam_key] = load_all_frames(path)
-        logging.info(f"Loaded {cam_key}: {camera_frames[cam_key].shape}")
-    return camera_frames
-
-
-def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
-    """Build the frame index schedule for multi-frame chunks.
-
-    Returns a list of frame-index lists. Each inner list has 4 indices.
-    """
-    chunks: list[list[int]] = []
-    current_frame = 23  # first anchor frame
-    for _ in range(num_chunks):
-        indices = [max(current_frame + off, 0) for off in RELATIVE_OFFSETS]
-        if indices[-1] >= total_frames:
-            logging.info(f"Frame {indices[-1]} >= {total_frames}, stopping at {len(chunks)} chunks")
-            break
-        chunks.append(indices)
-        current_frame += ACTION_HORIZON
-    return chunks
-
-
-def _make_obs_from_video(
-    camera_frames: dict[str, np.ndarray],
-    frame_indices: list[int],
-    prompt: str,
-    session_id: str,
-) -> dict:
-    """Build an observation dict from real video frames.
-
-    For 1 frame: each image key is (H, W, 3).
-    For 4 frames: each image key is (4, H, W, 3).
-    """
-    obs: dict = {}
-    for cam_key, all_frames in camera_frames.items():
-        selected = all_frames[frame_indices]  # (T, H, W, 3)
-        if len(frame_indices) == 1:
-            selected = selected[0]  # (H, W, 3)
-        obs[cam_key] = selected
-
-    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
-    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
-    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
-    obs["prompt"] = prompt
-    obs["session_id"] = session_id
-    return obs
-
-
-def _make_zero_observation(
-    server_config: policy_server.PolicyServerConfig,
-    prompt: str = "pick up the object",
-    session_id: str | None = None,
-) -> dict:
-    """Create a dummy observation matching AR_droid expectations.
-
-    AR_droid expects:
-        - 2 external cameras (exterior_image_0_left, exterior_image_1_left)
-        - 1 wrist camera (wrist_image_left)
-        - Image resolution: 180x320 (H x W)
-        - joint_position: 7 DoF
-        - gripper_position: 1 DoF
-    """
-    obs = {}
-
-    # Determine image resolution
-    if server_config.image_resolution is not None:
-        h, w = server_config.image_resolution
-    else:
-        # Default for AR_droid
-        h, w = 180, 320
-
-    # External cameras (0-indexed in roboarena)
-    for i in range(server_config.n_external_cameras):
-        obs[f"observation/exterior_image_{i}_left"] = np.zeros((h, w, 3), dtype=np.uint8)
-        if server_config.needs_stereo_camera:
-            obs[f"observation/exterior_image_{i}_right"] = np.zeros((h, w, 3), dtype=np.uint8)
-
-    # Wrist camera
-    if server_config.needs_wrist_camera:
-        obs["observation/wrist_image_left"] = np.zeros((h, w, 3), dtype=np.uint8)
-        if server_config.needs_stereo_camera:
-            obs["observation/wrist_image_right"] = np.zeros((h, w, 3), dtype=np.uint8)
-
-    # Session ID - should be passed in to ensure consistency within a session
-    if server_config.needs_session_id:
-        import uuid
-
-        # Generate unique session ID if not provided
-        obs["session_id"] = session_id if session_id else str(uuid.uuid4())
-
-    # State observations (AR_droid: 7 DoF arm + 1 gripper)
-    obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
-    obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
-    obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
-
-    # Language prompt
-    obs["prompt"] = prompt
-
-    return obs
-
-
-def test_ar_droid_policy_server(
-    host: str = "localhost",
-    port: int = 8000,
-    path: str = DEFAULT_WEBSOCKET_PATH,
-    num_chunks: int = 15,
-    prompt: str = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
-    use_zero_images: bool = False,
-):
-    """Test the AR_droid policy server with roboarena interface.
-
-    When use_zero_images is False (default), loads real video frames from
-    debug_image/ and follows the frame schedule from debug_inference.py.
-    """
-    logging.info(f"Connecting to AR_droid server at ws://{host}:{port}{path} ...")
-
-    client = OpenPIWebsocketClientPolicy(host=host, port=port, path=path)
-
-    # Validate server metadata
-    metadata = client.get_server_metadata()
-    logging.info(f"Server metadata: {metadata}")
-    assert isinstance(metadata, dict), "Metadata should be a dict"
-
-    try:
-        server_config = policy_server.PolicyServerConfig(**metadata)
-    except Exception as e:
-        logging.error(f"Error parsing metadata: {e}")
-        raise e
-
-    # Validate expected AR_droid configuration
-    logging.info(f"Server config: {server_config}")
-    assert server_config.n_external_cameras == 2, f"Expected 2 external cameras, got {server_config.n_external_cameras}"
-    assert server_config.needs_wrist_camera, "Expected wrist camera to be enabled"
-    assert server_config.action_space == "joint_position", (
-        f"Expected joint_position action space, got {server_config.action_space}"
-    )
-
-    logging.info("Server configuration validated for AR_droid")
-
-    # Generate unique session ID for this test run
-    session_id = str(uuid.uuid4())
-    logging.info(f"Session ID: {session_id}")
-
-    # ── Zero-image fallback mode ──────────────────────────────────────
-    if use_zero_images:
-        logging.info("Using ZERO dummy images (legacy mode)")
-        for i in range(num_chunks):
-            obs = _make_zero_observation(server_config, prompt=prompt, session_id=session_id)
-            logging.info(f"Inference {i + 1}/{num_chunks}: prompt='{prompt}'")
-            t0 = time.time()
-            actions = client.infer(obs)
-            dt = time.time() - t0
-            _log_action(actions, dt)
-
-        logging.info("Sending reset...")
-        client.reset({})
-        logging.info("Done (zero-image mode).")
-        return
-
-    # ── Real video frame mode ─────────────────────────────────────────
-    logging.info("Loading real video frames from debug_image/ directory")
-    camera_frames = load_camera_frames()
-
-    total_frames = min(v.shape[0] for v in camera_frames.values())
-    logging.info(f"Total frames available: {total_frames}")
-
-    # Build frame schedule
-    chunks = build_frame_schedule(total_frames, num_chunks)
-
-    logging.info("Frame schedule:")
-    logging.info("  Initial: [0]")
-    for i, indices in enumerate(chunks):
-        logging.info(f"  Chunk {i}: {indices}")
-
-    # Step 0: initial single frame
-    logging.info("=== Initial: frame [0] ===")
-    obs = _make_obs_from_video(camera_frames, [0], prompt, session_id)
-    t0 = time.time()
-    actions = client.infer(obs)
-    dt = time.time() - t0
-    _log_action(actions, dt)
-
-    # Subsequent chunks: send 4 frames at a time
-    for chunk_idx, frame_indices in enumerate(chunks):
-        logging.info(f"=== Chunk {chunk_idx}: frames {frame_indices} ===")
-        obs = _make_obs_from_video(camera_frames, frame_indices, prompt, session_id)
-        t0 = time.time()
-        actions = client.infer(obs)
-        dt = time.time() - t0
-        _log_action(actions, dt)
-
-    # Reset triggers video save on the server
-    logging.info("Sending reset to save video...")
-    client.reset({})
-
-    logging.info("Done.")
-
-
-def _log_action(actions: np.ndarray, dt: float) -> None:
-    """Pretty-print action shape, range, and timing."""
-    assert isinstance(actions, np.ndarray), f"Expected numpy array, got {type(actions)}"
-    assert actions.ndim == 2, f"Expected 2D array, got shape {actions.shape}"
-    assert actions.shape[-1] == 8, f"Expected 8 action dims (7 joints + 1 gripper), got {actions.shape[-1]}"
-    logging.info(f"  Action shape: {actions.shape}, range: [{actions.min():.4f}, {actions.max():.4f}], time: {dt:.2f}s")
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Test AR_droid policy server with real video frames from debug_image/")
-    parser.add_argument("--host", default="localhost", help="Server hostname")
-    parser.add_argument("--port", type=int, default=8000, help="Server port")
-    parser.add_argument(
-        "--path",
-        default=DEFAULT_WEBSOCKET_PATH,
-        help="WebSocket path suffix (default: empty string for the original DreamZero server)",
-    )
-    parser.add_argument(
-        "--num-chunks",
-        type=int,
-        default=15,
-        help="Number of 4-frame chunks to send after the initial frame (default: 15)",
-    )
-    parser.add_argument(
-        "--prompt",
-        default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
-        help="Language prompt for the policy",
-    )
-    parser.add_argument(
-        "--use-zero-images",
-        action="store_true",
-        help="Use zero dummy images instead of real video frames (legacy mode)",
-    )
-
-    args = parser.parse_args()
-
-    logging.basicConfig(
-        level=logging.INFO,
-        format="%(asctime)s [%(levelname)s] %(message)s",
-    )
-
-    test_ar_droid_policy_server(
-        host=args.host,
-        port=args.port,
-        path=args.path,
-        num_chunks=args.num_chunks,
-        prompt=args.prompt,
-        use_zero_images=args.use_zero_images,
-    )
-
-
-if __name__ == "__main__":
-    main()

tests/dreamzero/upstream/test_openpi_e2e_source_parity.py

@@ -1,358 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Formal OpenPI end-to-end parity: upstream DreamZero server vs `vllm serve`.
-
-This test uses DreamZero's own client-side observation builders from
-`${DREAMZERO_REPO}/test_client_AR.py`, and client-side websocket protocol from
-`${DREAMZERO_REPO}/eval_utils/policy_client.py`.
-
-The only client-side adaptation for vLLM is the websocket path:
-DreamZero's upstream server serves at `/`, while vLLM serves OpenPI at
-`/v1/realtime/robot/openpi`.
-
-Current scope for this test:
-- default two-GPU run (`nproc_per_node=2` on upstream, `--cfg-parallel-size 2` on `vllm serve`)
-- non-`torch.compile` (upstream launched through
-  `upstream_socket_server_no_compile.py`, vLLM with `--enforce-eager`)
-- non-DiT-cache / non-skip-schedule (`NUM_DIT_STEPS=16`)
-
-Serving contract locked by this test:
-- upstream DreamZero still boots from the local checkpoint directory
-- vLLM boots from the official DreamZero HF repo name (`GEAR-Dreams/DreamZero-DROID`)
-  rather than a prepared local bundle path
-"""
-
-from __future__ import annotations
-
-import os
-import shutil
-import subprocess
-import sys
-import time
-from pathlib import Path
-
-import numpy as np
-import pytest
-import torch
-
-from tests.helpers.runtime import get_open_port
-
-msgpack_numpy = pytest.importorskip("openpi_client.msgpack_numpy")
-
-_DREAMZERO_REPO_ENV = os.environ.get("DREAMZERO_REPO")
-DREAMZERO_REPO = Path(_DREAMZERO_REPO_ENV).expanduser() if _DREAMZERO_REPO_ENV else None
-if DREAMZERO_REPO is not None and str(DREAMZERO_REPO) not in sys.path:
-    sys.path.insert(0, str(DREAMZERO_REPO))
-
-try:
-    import test_client_AR as dreamzero_client
-    from eval_utils.policy_client import WebsocketClientPolicy
-except Exception:  # pragma: no cover - guarded by pytest skip below
-    dreamzero_client = None
-    WebsocketClientPolicy = None
-_BaseWebsocketClientPolicy = WebsocketClientPolicy if WebsocketClientPolicy is not None else object
-
-CHECKPOINT_DIR = DREAMZERO_REPO / "checkpoints" / "dreamzero" if DREAMZERO_REPO is not None else None
-VLLM_MODEL = os.environ.get("VLLM_DREAMZERO_MODEL", "GEAR-Dreams/DreamZero-DROID")
-SERVICE_READY_TIMEOUT_S = int(os.environ.get("OPENPI_SERVICE_READY_TIMEOUT_S", "900"))
-PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
-SESSION_ID = "openpi-e2e-parity-session"
-
-pytestmark = [
-    pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU required"),
-    pytest.mark.skipif(
-        dreamzero_client is None or WebsocketClientPolicy is None,
-        reason="DreamZero client modules are required on PYTHONPATH",
-    ),
-    pytest.mark.skipif(
-        DREAMZERO_REPO is None or not DREAMZERO_REPO.exists(),
-        reason="DreamZero source repo is required at DREAMZERO_REPO",
-    ),
-    pytest.mark.skipif(
-        CHECKPOINT_DIR is None or not CHECKPOINT_DIR.exists(), reason="DreamZero local checkpoint is required"
-    ),
-]
-
-
-class OpenPIWebsocketClientPolicy(_BaseWebsocketClientPolicy):
-    """DreamZero client protocol with an OpenPI websocket path suffix."""
-
-    def __init__(
-        self,
-        host: str = "127.0.0.1",
-        port: int = 8000,
-        path: str = "/v1/realtime/robot/openpi",
-    ) -> None:
-        self._uri = f"ws://{host}:{port}{path}"
-        self._packer = msgpack_numpy.Packer()
-        self._ws, self._server_metadata = self._wait_for_server()
-
-
-def _vllm_executable() -> str:
-    fallback = Path(sys.executable).with_name("vllm")
-    if fallback.exists():
-        return str(fallback)
-    exe = shutil.which("vllm")
-    if exe:
-        return exe
-    raise FileNotFoundError("Unable to locate `vllm` executable in current environment.")
-
-
-def _cfg_parallel_size() -> int:
-    return int(os.environ.get("OPENPI_E2E_CFG_PARALLEL_SIZE", "2"))
-
-
-def _pick_test_gpus() -> list[str]:
-    cfg_parallel_size = _cfg_parallel_size()
-    override = os.environ.get("OPENPI_E2E_GPUS") or os.environ.get("OPENPI_E2E_GPU")
-    if override is not None:
-        gpus = [part.strip() for part in override.split(",") if part.strip()]
-        if not gpus:
-            raise ValueError("OPENPI_E2E_GPUS is set but empty.")
-        if len(gpus) < cfg_parallel_size:
-            raise RuntimeError(f"Need {cfg_parallel_size} GPUs, but OPENPI_E2E_GPUS only provided {gpus}.")
-        return gpus
-
-    query = subprocess.check_output(
-        [
-            "nvidia-smi",
-            "--query-gpu=index,memory.used",
-            "--format=csv,noheader,nounits",
-        ],
-        text=True,
-    )
-    gpu_rows = []
-    for line in query.strip().splitlines():
-        gpu_index, used_mb = [part.strip() for part in line.split(",", maxsplit=1)]
-        gpu_rows.append((int(used_mb), gpu_index))
-    gpu_rows.sort()
-    gpus = [gpu_index for _, gpu_index in gpu_rows[: max(cfg_parallel_size, 1)]]
-    if len(gpus) < cfg_parallel_size:
-        raise RuntimeError(
-            f"Need {cfg_parallel_size} GPUs for cfg_parallel_size={cfg_parallel_size}, "
-            f"but found only {len(gpus)} candidates."
-        )
-    return gpus
-
-
-def _torchrun_argv(script: str, port: int) -> list[str]:
-    return [
-        sys.executable,
-        "-m",
-        "torch.distributed.run",
-        "--standalone",
-        "--nproc_per_node",
-        str(_cfg_parallel_size()),
-        script,
-        "--port",
-        str(port),
-        "--model_path",
-        str(CHECKPOINT_DIR),
-    ]
-
-
-def _run_upstream_service(port: int, log_path: Path) -> subprocess.Popen[str]:
-    env = os.environ.copy()
-    env.setdefault("PYTHONPATH", "")
-    env["PYTHONPATH"] = f"{Path.cwd()}:{DREAMZERO_REPO}:{env['PYTHONPATH']}".rstrip(":")
-    env["CUDA_VISIBLE_DEVICES"] = ",".join(_pick_test_gpus())
-    env.setdefault("NO_ALBUMENTATIONS_UPDATE", "1")
-    env["ATTENTION_BACKEND"] = "torch"
-    env.setdefault("ENABLE_TENSORRT", "false")
-    env["ENABLE_DIT_CACHE"] = "false"
-    env["NUM_DIT_STEPS"] = "16"
-    env["DYNAMIC_CACHE_SCHEDULE"] = "false"
-    argv = _torchrun_argv(
-        str(Path("tests/dreamzero/upstream/upstream_socket_server_no_compile.py")),
-        port,
-    )
-    log_file = log_path.open("w")
-    proc = subprocess.Popen(
-        argv,
-        stdout=log_file,
-        stderr=subprocess.STDOUT,
-        text=True,
-        cwd=str(Path.cwd()),
-        env=env,
-    )
-    proc._codex_log_file = log_file  # type: ignore[attr-defined]
-    return proc
-
-
-def _run_vllm_service(port: int, log_path: Path) -> subprocess.Popen[str]:
-    env = os.environ.copy()
-    gpus = _pick_test_gpus()
-    cfg_parallel_size = _cfg_parallel_size()
-    if cfg_parallel_size > len(gpus):
-        raise RuntimeError(
-            f"cfg_parallel_size={cfg_parallel_size} requires at least {cfg_parallel_size} GPUs, but only got {gpus}."
-        )
-    env["CUDA_VISIBLE_DEVICES"] = ",".join(gpus[:cfg_parallel_size])
-    env.setdefault("ATTENTION_BACKEND", "torch")
-    env.setdefault("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA")
-    env.setdefault("MASTER_PORT", str(get_open_port()))
-    argv = [
-        _vllm_executable(),
-        "serve",
-        VLLM_MODEL,
-        "--omni",
-        "--host",
-        "127.0.0.1",
-        "--port",
-        str(port),
-        "--served-model-name",
-        "dreamzero-droid",
-        "--enforce-eager",
-    ]
-    if cfg_parallel_size > 1:
-        argv.extend(["--cfg-parallel-size", str(cfg_parallel_size)])
-    log_file = log_path.open("w")
-    proc = subprocess.Popen(
-        argv,
-        stdout=log_file,
-        stderr=subprocess.STDOUT,
-        text=True,
-        env=env,
-        cwd=str(Path.cwd()),
-    )
-    proc._codex_log_file = log_file  # type: ignore[attr-defined]
-    return proc
-
-
-def _stop_process(proc: subprocess.Popen[str]) -> None:
-    log_file = getattr(proc, "_codex_log_file", None)
-    if proc.poll() is None:
-        proc.terminate()
-        try:
-            proc.wait(timeout=30)
-        except subprocess.TimeoutExpired:  # pragma: no cover - cleanup path
-            proc.kill()
-            proc.wait(timeout=10)
-    if log_file is not None:
-        log_file.close()
-
-
-def _build_obs_sequence() -> tuple[dict, dict]:
-    camera_frames = dreamzero_client.load_camera_frames()
-    chunks = dreamzero_client.build_frame_schedule(
-        min(v.shape[0] for v in camera_frames.values()),
-        1,
-    )
-    obs0 = dreamzero_client._make_obs_from_video(camera_frames, [0], PROMPT, SESSION_ID)
-    obs1 = dreamzero_client._make_obs_from_video(camera_frames, chunks[0], PROMPT, SESSION_ID)
-    return obs0, obs1
-
-
-def _wait_for_client_ready(client_factory, timeout_s: float, proc=None, log_path: Path | None = None):
-    deadline = time.time() + timeout_s
-    last_err: Exception | None = None
-    while time.time() < deadline:
-        if proc is not None and proc.poll() is not None:
-            details = ""
-            if log_path is not None and log_path.exists():
-                details = log_path.read_text(errors="replace")[-8000:]
-            raise RuntimeError(f"Service exited before becoming ready with code {proc.returncode}.\n{details}")
-        try:
-            return client_factory()
-        except Exception as exc:  # pragma: no cover - retry path
-            last_err = exc
-            time.sleep(1)
-    raise TimeoutError(f"Timed out waiting for websocket service: {last_err}")
-
-
-def _collect_outputs_with_client(client) -> tuple[dict, list[np.ndarray]]:
-    metadata = client.get_server_metadata()
-    obs0, obs1 = _build_obs_sequence()
-    outputs = [
-        client.infer(dict(obs0)),
-        client.infer(dict(obs1)),
-    ]
-    assert _normalize_reset_response(client.reset({})) == "reset successful"
-    outputs.append(client.infer(dict(obs0)))
-    client._ws.close()
-    return metadata, outputs
-
-
-def _normalize_reset_response(response) -> str:
-    if isinstance(response, str):
-        return response
-    decoded = msgpack_numpy.unpackb(response)
-    if isinstance(decoded, dict):
-        return str(decoded.get("status"))
-    return str(decoded)
-
-
-def _normalize_metadata(metadata: dict) -> dict:
-    normalized = dict(metadata)
-    if isinstance(normalized.get("image_resolution"), tuple):
-        normalized["image_resolution"] = list(normalized["image_resolution"])
-    return normalized
-
-
-def _assert_logs_clean(log_path: Path) -> None:
-    text = log_path.read_text(errors="replace")
-    if "SignalException: Process" in text and "got signal: 15" in text:
-        text = text.split("Traceback (most recent call last):", 1)[0]
-    assert "Traceback" not in text, text
-    assert "RuntimeError:" not in text, text
-
-
-def _assert_upstream_log_matches_vllm_baseline(log_path: Path) -> None:
-    text = log_path.read_text(errors="replace")
-    assert "DIT Compute Steps 8 steps" not in text, text
-    assert "DIT Compute Steps 16 steps" in text, text
-
-
-def test_openpi_service_matches_upstream_server_noncompile(tmp_path: Path) -> None:
-    expected_metadata = {
-        "image_resolution": [180, 320],
-        "n_external_cameras": 2,
-        "needs_wrist_camera": True,
-        "needs_stereo_camera": False,
-        "needs_session_id": True,
-        "action_space": "joint_position",
-    }
-
-    upstream_port = get_open_port()
-    upstream_log = tmp_path / "dreamzero_upstream.log"
-    upstream_proc = _run_upstream_service(upstream_port, upstream_log)
-    try:
-        upstream_client = _wait_for_client_ready(
-            lambda: WebsocketClientPolicy(host="127.0.0.1", port=upstream_port),
-            timeout_s=SERVICE_READY_TIMEOUT_S,
-            proc=upstream_proc,
-            log_path=upstream_log,
-        )
-        upstream_metadata, upstream_outputs = _collect_outputs_with_client(upstream_client)
-    finally:
-        _stop_process(upstream_proc)
-    _assert_logs_clean(upstream_log)
-    _assert_upstream_log_matches_vllm_baseline(upstream_log)
-
-    vllm_port = get_open_port()
-    vllm_log = tmp_path / "vllm_openpi.log"
-    vllm_proc = _run_vllm_service(vllm_port, vllm_log)
-    try:
-        vllm_client = _wait_for_client_ready(
-            lambda: OpenPIWebsocketClientPolicy(host="127.0.0.1", port=vllm_port),
-            timeout_s=SERVICE_READY_TIMEOUT_S,
-            proc=vllm_proc,
-            log_path=vllm_log,
-        )
-        vllm_metadata, vllm_outputs = _collect_outputs_with_client(vllm_client)
-    finally:
-        _stop_process(vllm_proc)
-    _assert_logs_clean(vllm_log)
-
-    assert _normalize_metadata(upstream_metadata) == expected_metadata
-    assert _normalize_metadata(vllm_metadata) == expected_metadata
-
-    for idx, (actual, expected) in enumerate(zip(vllm_outputs, upstream_outputs, strict=True)):
-        np.testing.assert_allclose(
-            actual,
-            expected,
-            rtol=1e-2,
-            atol=1e-3,
-            err_msg=f"OpenPI step {idx} output mismatch",
-        )

tests/e2e/accuracy/conftest.py

@@ -1,54 +1,15 @@
 from __future__ import annotations
 
-import os
-import subprocess
-from contextlib import contextmanager
-from dataclasses import dataclass
 from io import BytesIO
 from pathlib import Path
 
 import pytest
 import requests
-import torch
 from PIL import Image
 
-from tests.helpers.runtime import OmniServer, OmniServerParams
-
 
 def pytest_addoption(parser):
     group = parser.getgroup("accuracy-e2e")
-    group.addoption("--gebench-root", action="store", default=None, help="Local GEBench dataset root")
-    group.addoption("--gedit-root", action="store", default=None, help="Local GEdit-Bench dataset root")
-    group.addoption(
-        "--gebench-model", action="store", default="Qwen/Qwen-Image-2512", help="Generate model for GEBench smoke"
-    )
-    group.addoption(
-        "--gedit-model", action="store", default="Qwen/Qwen-Image-Edit", help="Generate model for GEdit-Bench smoke"
-    )
-    group.addoption(
-        "--accuracy-judge-model",
-        action="store",
-        default="QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ",
-        help="Judge model path",
-    )
-    group.addoption("--accuracy-gpu", action="store", default="0", help="Single GPU id used sequentially")
-    group.addoption("--gebench-port", action="store", type=int, default=8093, help="Generate port for GEBench")
-    group.addoption("--gedit-port", action="store", type=int, default=8093, help="Generate port for GEdit-Bench")
-    group.addoption(
-        "--gebench-samples-per-type",
-        action="store",
-        type=int,
-        default=10,
-        help="Balanced sample count per GEBench type",
-    )
-    group.addoption(
-        "--gedit-samples-per-group",
-        action="store",
-        type=int,
-        default=20,
-        help="Balanced sample count per GEdit task group",
-    )
-    group.addoption("--accuracy-workers", action="store", type=int, default=1, help="Worker count for accuracy benches")
     group.addoption(
         "--wan22-i2v-image-source",
         action="store",
@@ -77,99 +38,6 @@ def pytest_addoption(parser):
     )
 
 
-def _hf_cache_root() -> Path:
-    return Path(os.environ.get("HF_HOME", "/root/.cache/huggingface"))
-
-
-def _dataset_cache_dirs(dataset_id: str) -> list[Path]:
-    cache_root = _hf_cache_root() / "hub" / f"datasets--{dataset_id.replace('/', '--')}" / "snapshots"
-    if not cache_root.exists():
-        return []
-    return sorted(
-        (path for path in cache_root.iterdir() if path.is_dir()), key=lambda path: path.stat().st_mtime, reverse=True
-    )
-
-
-def _ensure_dataset_snapshot(dataset_id: str) -> Path:
-    candidates = _dataset_cache_dirs(dataset_id)
-    if candidates:
-        return candidates[0]
-
-    subprocess.run(
-        ["hf", "download", "--repo-type", "dataset", dataset_id],
-        check=True,
-    )
-    candidates = _dataset_cache_dirs(dataset_id)
-    if not candidates:
-        raise FileNotFoundError(
-            f"Dataset {dataset_id} was downloaded but no snapshot was found under {_hf_cache_root()}"
-        )
-    return candidates[0]
-
-
-def _resolve_dataset_root(request: pytest.FixtureRequest, option_name: str, dataset_id: str) -> Path:
-    value = request.config.getoption(option_name)
-    if value:
-        path = Path(value)
-        if not path.exists():
-            pytest.skip(f"Dataset path does not exist: {path}")
-        return path
-    return _ensure_dataset_snapshot(dataset_id)
-
-
-@dataclass
-class AccuracyServerConfig:
-    generate_params: OmniServerParams
-    judge_params: OmniServerParams
-    run_level: str
-    model_prefix: str
-
-    @contextmanager
-    def generate_server(self):
-        params = self.generate_params
-        model = self.model_prefix + params.model
-        server_args = params.server_args or []
-        if params.use_omni and params.stage_init_timeout is not None:
-            server_args = ["--stage-init-timeout", str(params.stage_init_timeout), *server_args]
-        with OmniServer(
-            model,
-            server_args,
-            port=params.port,
-            env_dict=params.env_dict,
-            use_omni=params.use_omni,
-        ) as server:
-            yield server
-
-    @contextmanager
-    def judge_server(self):
-        params = self.judge_params
-        model = self.model_prefix + params.model
-        server_args = params.server_args or []
-        with OmniServer(
-            model,
-            server_args,
-            port=params.port,
-            env_dict=params.env_dict,
-            use_omni=params.use_omni,
-        ) as server:
-            yield server
-
-
-@pytest.fixture(scope="session")
-def gebench_dataset_root(request: pytest.FixtureRequest) -> Path:
-    return _resolve_dataset_root(request, "gebench_root", "stepfun-ai/GEBench")
-
-
-@pytest.fixture(scope="session")
-def gedit_dataset_root(request: pytest.FixtureRequest) -> Path:
-    return _resolve_dataset_root(request, "gedit_root", "stepfun-ai/GEdit-Bench")
-
-
-@pytest.fixture(scope="session")
-def accuracy_workers(request: pytest.FixtureRequest) -> int:
-    return int(request.config.getoption("accuracy_workers"))
-
-
 @pytest.fixture(scope="session")
 def wan22_i2v_image_source(request: pytest.FixtureRequest) -> str | None:
     value = request.config.getoption("wan22_i2v_image_source")
@@ -192,16 +60,6 @@ def hunyuanvideo15_online_timeout_seconds(request: pytest.FixtureRequest) -> int
     return int(request.config.getoption("hunyuanvideo15_online_timeout_seconds"))
 
 
-@pytest.fixture(scope="session")
-def gebench_samples_per_type(request: pytest.FixtureRequest) -> int:
-    return int(request.config.getoption("gebench_samples_per_type"))
-
-
-@pytest.fixture(scope="session")
-def gedit_samples_per_group(request: pytest.FixtureRequest) -> int:
-    return int(request.config.getoption("gedit_samples_per_group"))
-
-
 @pytest.fixture(scope="session")
 def accuracy_artifact_root() -> Path:
     root = Path(__file__).resolve().parent / "artifacts"
@@ -209,12 +67,6 @@ def accuracy_artifact_root() -> Path:
     return root
 
 
-@pytest.fixture(scope="session")
-def accuracy_assets_root() -> Path:
-    root = Path(__file__).resolve().parent / "assets"
-    return root
-
-
 @pytest.fixture(scope="session")
 def qwen_bear_image(accuracy_artifact_root: Path):
     """Download the Qwen bear image from the URL and save it to the accuracy artifact root."""
@@ -249,80 +101,3 @@ def rabbit_image(accuracy_artifact_root: Path):
     image.save(image_path)
     yield image
     image.close()
-
-
-def _build_accuracy_server_config(
-    *,
-    generate_model: str,
-    judge_model: str,
-    shared_gpu: str,
-    port: int,
-    run_level: str,
-    model_prefix: str,
-) -> AccuracyServerConfig:
-    if torch.accelerator.device_count() < 1:
-        pytest.skip("Need at least 1 CUDA GPU for accuracy benchmark smoke tests.")
-
-    if not generate_model:
-        pytest.skip("No generate model configured for accuracy benchmark test.")
-    generate_server_args = ["--num-gpus", "1"]
-    judge_server_args = [
-        "--max-model-len",
-        "32768",
-        "--gpu-memory-utilization",
-        "0.8",
-    ]
-
-    judge_env = {"CUDA_VISIBLE_DEVICES": shared_gpu}
-
-    return AccuracyServerConfig(
-        generate_params=OmniServerParams(
-            model=generate_model,
-            port=port,
-            server_args=generate_server_args,
-            env_dict={"CUDA_VISIBLE_DEVICES": shared_gpu},
-            use_omni=True,
-            stage_init_timeout=300,
-        ),
-        judge_params=OmniServerParams(
-            model=judge_model,
-            port=port,
-            server_args=judge_server_args,
-            env_dict=judge_env,
-            use_omni=False,
-        ),
-        run_level=run_level,
-        model_prefix=model_prefix,
-    )
-
-
-@pytest.fixture
-def gebench_accuracy_servers(
-    request: pytest.FixtureRequest,
-    run_level: str,
-    model_prefix: str,
-) -> AccuracyServerConfig:
-    return _build_accuracy_server_config(
-        generate_model=request.config.getoption("gebench_model"),
-        judge_model=request.config.getoption("accuracy_judge_model"),
-        shared_gpu=str(request.config.getoption("accuracy_gpu")),
-        port=int(request.config.getoption("gebench_port")),
-        run_level=run_level,
-        model_prefix=model_prefix,
-    )
-
-
-@pytest.fixture
-def gedit_accuracy_servers(
-    request: pytest.FixtureRequest,
-    run_level: str,
-    model_prefix: str,
-) -> AccuracyServerConfig:
-    return _build_accuracy_server_config(
-        generate_model=request.config.getoption("gedit_model"),
-        judge_model=request.config.getoption("accuracy_judge_model"),
-        shared_gpu=str(request.config.getoption("accuracy_gpu")),
-        port=int(request.config.getoption("gedit_port")),
-        run_level=run_level,
-        model_prefix=model_prefix,
-    )

tests/e2e/accuracy/helpers.py

@@ -3,8 +3,12 @@ import os
 import re
 import shutil
 import subprocess
+import sys
 import time
+from abc import ABC, abstractmethod
 from base64 import b64decode, b64encode
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
 from fractions import Fraction
 from hashlib import sha1
 from io import BytesIO
@@ -53,8 +57,6 @@ def env_to_apply_ftfy_mock_in_subproc(env: dict[str, str] | None = None) -> dict
 
 
 def reset_artifact_dir(path: Path) -> Path:
-    import shutil
-
     if path.exists():
         shutil.rmtree(path)
     path.mkdir(parents=True, exist_ok=True)
@@ -609,3 +611,344 @@ def download_images(urls: list[str]) -> list[Image.Image]:
         images.append(image)
 
     return images
+
+
+_DREAMZERO_UPSTREAM_SERVER_SCRIPT = (
+    Path(__file__).resolve().parent / "upstream_mock" / "upstream_socket_server_no_compile.py"
+)
+_DREAMZERO_DEFAULT_MODEL = "GEAR-Dreams/DreamZero-DROID"
+_DREAMZERO_DEFAULT_READY_TIMEOUT_S = int(os.environ.get("OPENPI_SERVICE_READY_TIMEOUT_S", "900"))
+_DREAMZERO_DEFAULT_PROMPT = (
+    "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
+)
+_DREAMZERO_DEFAULT_SESSION_ID = "openpi-e2e-parity-session"
+_DREAMZERO_OPENPI_VLLM_PATH = "/v1/realtime/robot/openpi"
+
+
+def require_dreamzero_parity_env() -> tuple[Path, Path]:
+    repo_env = os.environ.get("DREAMZERO_REPO")
+    if repo_env is None:
+        raise ValueError("DREAMZERO_REPO environment variable is required.")
+    repo = Path(repo_env).expanduser()
+    if not repo.is_dir():
+        raise FileNotFoundError(f"DREAMZERO_REPO does not exist: {repo}")
+    checkpoint_dir = repo / "checkpoints" / "dreamzero"
+    if not checkpoint_dir.is_dir():
+        raise FileNotFoundError(f"DreamZero checkpoint not found: {checkpoint_dir}")
+    if str(repo) not in sys.path:
+        sys.path.insert(0, str(repo))
+    return repo, checkpoint_dir
+
+
+def load_dreamzero_upstream_modules():
+    require_dreamzero_parity_env()
+    import test_client_AR as dreamzero_upstream_client
+    from eval_utils.policy_client import WebsocketClientPolicy
+
+    return dreamzero_upstream_client, WebsocketClientPolicy
+
+
+def dreamzero_cfg_parallel_size() -> int:
+    return int(os.environ.get("OPENPI_E2E_CFG_PARALLEL_SIZE", "2"))
+
+
+def _dreamzero_parity_vllm_env() -> dict[str, str]:
+    from tests.helpers.env import pick_least_used_device_indices
+    from tests.helpers.runtime import get_open_port
+
+    parallel_size = dreamzero_cfg_parallel_size()
+    gpus = [str(index) for index in pick_least_used_device_indices(parallel_size)]
+    if parallel_size > len(gpus):
+        raise RuntimeError(
+            f"cfg_parallel_size={parallel_size} requires at least {parallel_size} GPUs, but only got {gpus}."
+        )
+    return {
+        "CUDA_VISIBLE_DEVICES": ",".join(gpus[:parallel_size]),
+        "ATTENTION_BACKEND": "torch",
+        "DIFFUSION_ATTENTION_BACKEND": "TORCH_SDPA",
+        "VLLM_DISABLE_COMPILE_CACHE": "1",
+        "MASTER_PORT": str(get_open_port()),
+    }
+
+
+def dreamzero_parity_vllm_params(model: str | None = None):
+    from tests.helpers.runtime import OmniServerParams
+    from tests.helpers.stage_config import get_deploy_config_path
+
+    parallel_size = dreamzero_cfg_parallel_size()
+    server_args = [
+        "--served-model-name",
+        "dreamzero-droid",
+        "--enforce-eager",
+    ]
+    if parallel_size > 1:
+        server_args.extend(["--cfg-parallel-size", str(parallel_size)])
+    return OmniServerParams(
+        model=model or os.environ.get("VLLM_DREAMZERO_MODEL", _DREAMZERO_DEFAULT_MODEL),
+        stage_config_path=get_deploy_config_path("dreamzero_tp1_cfg2.yaml"),
+        server_args=server_args,
+        env_dict=_dreamzero_parity_vllm_env(),
+    )
+
+
+def _dreamzero_omni_server_from_params(params):
+    from tests.helpers.runtime import OmniServer
+
+    server_args = list(params.server_args or [])
+    if params.stage_config_path is not None:
+        server_args += ["--stage-configs-path", params.stage_config_path]
+    return OmniServer(
+        params.model,
+        server_args,
+        port=params.port,
+        env_dict=params.env_dict,
+        use_omni=params.use_omni,
+    )
+
+
+@contextmanager
+def dreamzero_parity_vllm_client(
+    params=None,
+    *,
+    ready_timeout_s: float = _DREAMZERO_DEFAULT_READY_TIMEOUT_S,
+) -> Iterator["DreamZeroParityClient"]:
+    """Start vLLM via ``OmniServer`` and return a ready OpenPI parity client."""
+    with _dreamzero_omni_server_from_params(params or dreamzero_parity_vllm_params()) as server:
+        client = _dreamzero_wait_for_client_ready(
+            lambda: _dreamzero_create_openpi_vllm_client(server.host, server.port),
+            timeout_s=ready_timeout_s,
+            proc=server.proc,
+        )
+        yield DreamZeroParityClient(client)
+
+
+def _dreamzero_create_openpi_vllm_client(host: str, port: int) -> Any:
+    from openpi_client import msgpack_numpy
+
+    _, WebsocketClientPolicy = load_dreamzero_upstream_modules()
+
+    class _OpenPIWebsocketClientPolicy(WebsocketClientPolicy):
+        def __init__(
+            self,
+            host: str = "127.0.0.1",
+            port: int = 8000,
+            path: str = _DREAMZERO_OPENPI_VLLM_PATH,
+        ) -> None:
+            self._uri = f"ws://{host}:{port}{path}"
+            self._packer = msgpack_numpy.Packer()
+            self._ws, self._server_metadata = self._wait_for_server()
+
+    return _OpenPIWebsocketClientPolicy(host=host, port=port)
+
+
+def _dreamzero_stop_process(proc: subprocess.Popen[str]) -> None:
+    log_file = getattr(proc, "_codex_log_file", None)
+    if proc.poll() is None:
+        proc.terminate()
+        try:
+            proc.wait(timeout=30)
+        except subprocess.TimeoutExpired:  # pragma: no cover - cleanup path
+            proc.kill()
+            proc.wait(timeout=10)
+    if log_file is not None:
+        log_file.close()
+
+
+def _dreamzero_wait_for_client_ready(
+    client_factory: Callable[[], Any],
+    *,
+    timeout_s: float,
+    proc: subprocess.Popen[str] | None = None,
+    log_path: Path | None = None,
+) -> Any:
+    deadline = time.time() + timeout_s
+    last_err: Exception | None = None
+    while time.time() < deadline:
+        if proc is not None and proc.poll() is not None:
+            details = ""
+            if log_path is not None and log_path.exists():
+                details = log_path.read_text(errors="replace")[-8000:]
+            raise RuntimeError(f"Service exited before becoming ready with code {proc.returncode}.\n{details}")
+        try:
+            return client_factory()
+        except Exception as exc:  # pragma: no cover - retry path
+            last_err = exc
+            time.sleep(1)
+    raise TimeoutError(f"Timed out waiting for websocket service: {last_err}")
+
+
+def _dreamzero_normalize_reset_response(response: Any) -> str:
+    from openpi_client import msgpack_numpy
+
+    if isinstance(response, str):
+        return response
+    decoded = msgpack_numpy.unpackb(response)
+    if isinstance(decoded, dict):
+        return str(decoded.get("status"))
+    return str(decoded)
+
+
+def normalize_dreamzero_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
+    normalized = dict(metadata)
+    if isinstance(normalized.get("image_resolution"), tuple):
+        normalized["image_resolution"] = list(normalized["image_resolution"])
+    return normalized
+
+
+def assert_dreamzero_logs_clean(log_path: Path) -> None:
+    text = log_path.read_text(errors="replace")
+    if "SignalException: Process" in text and "got signal: 15" in text:
+        text = text.split("Traceback (most recent call last):", 1)[0]
+    assert "Traceback" not in text, text
+    assert "RuntimeError:" not in text, text
+
+
+def assert_dreamzero_upstream_log_matches_vllm_baseline(log_path: Path) -> None:
+    text = log_path.read_text(errors="replace")
+    assert "DIT Compute Steps 8 steps" not in text, text
+    assert "DIT Compute Steps 16 steps" in text, text
+
+
+class DreamZeroParityClient:
+    """Connected OpenPI websocket client for DreamZero parity inference."""
+
+    def __init__(self, client: Any) -> None:
+        self._client = client
+
+    @property
+    def raw_client(self) -> Any:
+        return self._client
+
+    def build_obs_sequence(self) -> tuple[dict[str, Any], dict[str, Any]]:
+        dreamzero_upstream_client, _ = load_dreamzero_upstream_modules()
+        camera_frames = dreamzero_upstream_client.load_camera_frames()
+        chunks = dreamzero_upstream_client.build_frame_schedule(
+            min(v.shape[0] for v in camera_frames.values()),
+            1,
+        )
+        obs0 = dreamzero_upstream_client._make_obs_from_video(
+            camera_frames,
+            [0],
+            _DREAMZERO_DEFAULT_PROMPT,
+            _DREAMZERO_DEFAULT_SESSION_ID,
+        )
+        obs1 = dreamzero_upstream_client._make_obs_from_video(
+            camera_frames,
+            chunks[0],
+            _DREAMZERO_DEFAULT_PROMPT,
+            _DREAMZERO_DEFAULT_SESSION_ID,
+        )
+        return obs0, obs1
+
+    def collect_parity_outputs(self) -> tuple[dict[str, Any], list[np.ndarray]]:
+        metadata = self._client.get_server_metadata()
+        obs0, obs1 = self.build_obs_sequence()
+        outputs = [
+            self._client.infer(dict(obs0)),
+            self._client.infer(dict(obs1)),
+        ]
+        assert _dreamzero_normalize_reset_response(self._client.reset({})) == "reset successful"
+        outputs.append(self._client.infer(dict(obs0)))
+        self._client._ws.close()
+        return metadata, outputs
+
+
+class _DreamZeroParityServer(ABC):
+    host = "127.0.0.1"
+
+    def __init__(
+        self,
+        *,
+        log_path: Path,
+        port: int | None = None,
+        ready_timeout_s: float = _DREAMZERO_DEFAULT_READY_TIMEOUT_S,
+    ) -> None:
+        from tests.helpers.runtime import get_open_port
+
+        self.log_path = log_path
+        self.port = get_open_port() if port is None else port
+        self.ready_timeout_s = ready_timeout_s
+        self.proc: subprocess.Popen[str] | None = None
+        self._log_file = None
+
+    @abstractmethod
+    def _build_argv(self) -> list[str]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _build_env(self) -> dict[str, str]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _create_client(self) -> Any:
+        raise NotImplementedError
+
+    def _start(self) -> None:
+        self.log_path.parent.mkdir(parents=True, exist_ok=True)
+        self._log_file = self.log_path.open("w")
+        self.proc = subprocess.Popen(
+            self._build_argv(),
+            stdout=self._log_file,
+            stderr=subprocess.STDOUT,
+            text=True,
+            cwd=str(Path.cwd()),
+            env=self._build_env(),
+        )
+        self.proc._codex_log_file = self._log_file  # type: ignore[attr-defined]
+
+    def __enter__(self) -> DreamZeroParityClient:
+        self._start()
+        client = _dreamzero_wait_for_client_ready(
+            self._create_client,
+            timeout_s=self.ready_timeout_s,
+            proc=self.proc,
+            log_path=self.log_path,
+        )
+        return DreamZeroParityClient(client)
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self.proc is not None:
+            _dreamzero_stop_process(self.proc)
+            self.proc = None
+
+
+class DreamZeroUpstreamServer(_DreamZeroParityServer):
+    """Official DreamZero upstream reference server launched via torchrun."""
+
+    def _build_argv(self) -> list[str]:
+        _, checkpoint_dir = require_dreamzero_parity_env()
+        return [
+            sys.executable,
+            "-m",
+            "torch.distributed.run",
+            "--standalone",
+            "--nproc_per_node",
+            str(dreamzero_cfg_parallel_size()),
+            str(_DREAMZERO_UPSTREAM_SERVER_SCRIPT),
+            "--port",
+            str(self.port),
+            "--model_path",
+            str(checkpoint_dir),
+        ]
+
+    def _build_env(self) -> dict[str, str]:
+        from tests.helpers.env import pick_least_used_device_indices
+
+        dreamzero_repo, _ = require_dreamzero_parity_env()
+        env = os.environ.copy()
+        env.setdefault("PYTHONPATH", "")
+        env["PYTHONPATH"] = f"{Path.cwd()}:{dreamzero_repo}:{env['PYTHONPATH']}".rstrip(":")
+        env["CUDA_VISIBLE_DEVICES"] = ",".join(
+            str(index) for index in pick_least_used_device_indices(max(dreamzero_cfg_parallel_size(), 1))
+        )
+        env.setdefault("NO_ALBUMENTATIONS_UPDATE", "1")
+        env["ATTENTION_BACKEND"] = "torch"
+        env.setdefault("ENABLE_TENSORRT", "false")
+        env["ENABLE_DIT_CACHE"] = "false"
+        env["NUM_DIT_STEPS"] = "16"
+        env["DYNAMIC_CACHE_SCHEDULE"] = "false"
+        return env
+
+    def _create_client(self) -> Any:
+        _, WebsocketClientPolicy = load_dreamzero_upstream_modules()
+        return WebsocketClientPolicy(host=self.host, port=self.port)

tests/e2e/accuracy/test_dreamzero.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DreamZero OpenPI upstream parity accuracy test."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from tests.e2e.accuracy.helpers import (
+    DreamZeroUpstreamServer,
+    assert_dreamzero_logs_clean,
+    assert_dreamzero_upstream_log_matches_vllm_baseline,
+    dreamzero_parity_vllm_client,
+    load_dreamzero_upstream_modules,
+    normalize_dreamzero_metadata,
+)
+from tests.helpers.mark import hardware_test
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.slow]
+
+EXPECTED_METADATA = {
+    "image_resolution": [180, 320],
+    "n_external_cameras": 2,
+    "needs_wrist_camera": True,
+    "needs_stereo_camera": False,
+    "needs_session_id": True,
+    "action_space": "joint_position",
+}
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+def test_dreamzero_openpi_upstream_parity(tmp_path: Path) -> None:
+    """Assert vLLM OpenPI actions match the upstream DreamZero reference server."""
+    load_dreamzero_upstream_modules()
+
+    upstream_log = tmp_path / "dreamzero_upstream.log"
+    with DreamZeroUpstreamServer(log_path=upstream_log) as upstream_client:
+        upstream_metadata, upstream_outputs = upstream_client.collect_parity_outputs()
+    assert_dreamzero_logs_clean(upstream_log)
+    assert_dreamzero_upstream_log_matches_vllm_baseline(upstream_log)
+
+    with dreamzero_parity_vllm_client() as vllm_client:
+        vllm_metadata, vllm_outputs = vllm_client.collect_parity_outputs()
+
+    assert normalize_dreamzero_metadata(upstream_metadata) == EXPECTED_METADATA
+    assert normalize_dreamzero_metadata(vllm_metadata) == EXPECTED_METADATA
+
+    for idx, (vllm_action, upstream_action) in enumerate(zip(vllm_outputs, upstream_outputs, strict=True)):
+        np.testing.assert_allclose(
+            vllm_action,
+            upstream_action,
+            rtol=1e-2,
+            atol=1e-3,
+            err_msg=f"OpenPI step {idx} output mismatch",
+        )

tests/e2e/accuracy/test_gebench_h100_smoke.py

@@ -1,97 +0,0 @@
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-import pytest
-
-from benchmarks.accuracy.text_to_image.gbench import main as gbench_main
-from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir
-from tests.helpers.mark import hardware_test
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
-
-
-@pytest.mark.benchmark
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skip(reason="#3256")
-def test_gebench_h100_smoke(
-    gebench_accuracy_servers,
-    accuracy_artifact_root: Path,
-    gebench_dataset_root: Path,
-    accuracy_workers: int,
-) -> None:
-    model_label = infer_model_label(gebench_accuracy_servers.generate_params.model).lower()
-    output_root = reset_artifact_dir(accuracy_artifact_root / f"gebench_{model_label}")
-
-    with gebench_accuracy_servers.generate_server() as generate_server:
-        for data_type in ("type3", "type4"):
-            assert (
-                gbench_main(
-                    [
-                        "generate",
-                        "--dataset-root",
-                        str(gebench_dataset_root),
-                        "--output-root",
-                        str(output_root),
-                        "--base-url",
-                        f"http://{generate_server.host}:{generate_server.port}",
-                        "--model",
-                        generate_server.model,
-                        "--data-type",
-                        data_type,
-                        "--width",
-                        "768",
-                        "--height",
-                        "576",
-                        "--output-compression",
-                        "98",
-                        "--num-inference-steps",
-                        "8",
-                        "--workers",
-                        str(accuracy_workers),
-                    ]
-                )
-                == 0
-            )
-
-    with gebench_accuracy_servers.judge_server() as judge_server:
-        for data_type in ("type3", "type4"):
-            assert (
-                gbench_main(
-                    [
-                        "evaluate",
-                        "--dataset-root",
-                        str(gebench_dataset_root),
-                        "--output-root",
-                        str(output_root),
-                        "--data-type",
-                        data_type,
-                        "--judge-base-url",
-                        f"http://{judge_server.host}:{judge_server.port}",
-                        "--judge-model",
-                        judge_server.model,
-                        "--judge-api-key",
-                        "EMPTY",
-                        "--workers",
-                        str(accuracy_workers),
-                    ]
-                )
-                == 0
-            )
-
-    assert gbench_main(["summarize", "--output-root", str(output_root)]) == 0
-
-    summary = json.loads((output_root / "summary.json").read_text(encoding="utf-8"))
-    assert "generation" in summary
-    assert "evaluation" in summary
-
-    for data_type in ("type3", "type4"):
-        assert data_type in summary["generation"]["by_type"]
-        assert summary["generation"]["by_type"][data_type]["count"] > 0
-        assert data_type in summary["evaluation"]["by_type"]
-        assert summary["evaluation"]["by_type"][data_type]["count"] > 0
-
-    assert summary["evaluation"]["overall_mean"] >= 0.45
-    assert summary["evaluation"]["by_type"]["type3"]["overall_mean"] >= 0.45
-    assert summary["evaluation"]["by_type"]["type4"]["overall_mean"] >= 0.45

tests/e2e/accuracy/test_gedit_bench_h100_smoke.py

@@ -1,115 +0,0 @@
-from __future__ import annotations
-
-import json
-from pathlib import Path
-
-import pytest
-
-from benchmarks.accuracy.image_to_image.gedit_bench import GROUPS
-from benchmarks.accuracy.image_to_image.gedit_bench import main as gedit_main
-from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir
-from tests.helpers.mark import hardware_test
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
-
-
-@pytest.mark.benchmark
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.skip(reason="#3257")
-def test_gedit_bench_h100_smoke(
-    gedit_accuracy_servers,
-    accuracy_artifact_root: Path,
-    gedit_dataset_root: Path,
-    gedit_samples_per_group: int,
-    accuracy_workers: int,
-) -> None:
-    model_label = infer_model_label(gedit_accuracy_servers.generate_params.model).lower()
-    output_root = reset_artifact_dir(accuracy_artifact_root / f"gedit_results_{model_label}")
-    score_root = reset_artifact_dir(accuracy_artifact_root / f"gedit_scores_{model_label}")
-    model_name = model_label
-
-    with gedit_accuracy_servers.generate_server() as generate_server:
-        assert (
-            gedit_main(
-                [
-                    "generate",
-                    "--dataset-ref",
-                    str(gedit_dataset_root),
-                    "--output-root",
-                    str(output_root),
-                    "--base-url",
-                    f"http://{generate_server.host}:{generate_server.port}",
-                    "--model",
-                    generate_server.model,
-                    "--model-name",
-                    model_name,
-                    "--task-type",
-                    "all",
-                    "--instruction-language",
-                    "all",
-                    "--samples-per-group",
-                    str(gedit_samples_per_group),
-                    "--workers",
-                    str(accuracy_workers),
-                ]
-            )
-            == 0
-        )
-
-    with gedit_accuracy_servers.judge_server() as judge_server:
-        assert (
-            gedit_main(
-                [
-                    "evaluate",
-                    "--dataset-ref",
-                    str(gedit_dataset_root),
-                    "--output-root",
-                    str(output_root),
-                    "--model-name",
-                    model_name,
-                    "--save-dir",
-                    str(score_root),
-                    "--task-type",
-                    "all",
-                    "--instruction-language",
-                    "all",
-                    "--judge-base-url",
-                    f"http://{judge_server.host}:{judge_server.port}",
-                    "--judge-model",
-                    judge_server.model,
-                    "--judge-api-key",
-                    "EMPTY",
-                    "--samples-per-group",
-                    str(gedit_samples_per_group),
-                    "--workers",
-                    str(accuracy_workers),
-                ]
-            )
-            == 0
-        )
-
-    csv_path = score_root / f"{model_name}_all_all_vie_score.csv"
-    assert gedit_main(["summarize", "--csv-path", str(csv_path), "--language", "all"]) == 0
-
-    summary_path = score_root / f"{model_name}_all_all_summary.json"
-    summary = json.loads(summary_path.read_text(encoding="utf-8"))
-    assert set(summary["languages"]) == {"en", "cn"}
-
-    for language in ["en", "cn"]:
-        language_summary = summary["languages"][language]
-        assert language_summary["overall"]["count"] is not None
-        assert language_summary["intersection"]["count"] is not None
-        assert language_summary["overall"]["Q_SC"] is not None
-        assert language_summary["overall"]["Q_PQ"] is not None
-        assert language_summary["overall"]["Q_O"] is not None
-
-        for group in GROUPS:
-            group_summary = language_summary["by_group"][group]
-            assert set(group_summary) == {"count", "Q_SC", "Q_PQ", "Q_O"}
-
-    assert summary["languages"]["en"]["overall"]["Q_SC"] >= 6.95
-    assert summary["languages"]["en"]["overall"]["Q_PQ"] >= 5.8
-    assert summary["languages"]["en"]["overall"]["Q_O"] >= 6.15
-    assert summary["languages"]["cn"]["overall"]["Q_SC"] >= 6.9
-    assert summary["languages"]["cn"]["overall"]["Q_PQ"] >= 5.7
-    assert summary["languages"]["cn"]["overall"]["Q_O"] >= 6.1

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -28,6 +28,7 @@ from tests.e2e.accuracy.helpers import (
     model_output_dir,
 )
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
 from tests.helpers.runtime import OmniRunner, OmniServer
 from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
     build_prompt_tokens,
@@ -94,6 +95,7 @@ def _default_ar_dit_devices() -> tuple[str, str]:
 
 AR_DEVICES, DIT_DEVICES = _default_ar_dit_devices()
 MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
+HUNYUAN_IMAGE_REF_PATH = get_asset_path("hunyuan/hunyuan_image_ref.png")
 NUM_INFERENCE_STEPS = 50
 GUIDANCE_SCALE = 2.5
 
@@ -414,7 +416,7 @@ def _run_online(stage_configs_path: str, output_path: Path) -> tuple[Image.Image
     torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE,
     reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)",
 )
-def test_image_to_image_alignment_online(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None:
+def test_image_to_image_alignment_online(accuracy_artifact_root: Path) -> None:
     """Online API test: same pipeline, same seed as offline → PSNR >= 10 dB."""
     if importlib.util.find_spec("FlagEmbedding") is None:
         raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding")
@@ -431,7 +433,7 @@ def test_image_to_image_alignment_online(accuracy_artifact_root: Path, accuracy_
     scorer = SemanticSimilarityScorer()
     clip_scorer = CLIPScorer()
     cot_results = scorer.text_similarity(online_cot, COT_REF)
-    image_ref = Image.open(str(accuracy_assets_root / "hunyuan_image_ref.png")).convert("RGB")
+    image_ref = Image.open(HUNYUAN_IMAGE_REF_PATH).convert("RGB")
     image_clip_score = clip_scorer.image_image_score(online_image, image_ref)
     ssim_value, psnr_value = compute_image_ssim_psnr(prediction=online_image, reference=image_ref, compare_mode="RGB")
 
@@ -529,7 +531,7 @@ def _run_dit_model(
     torch.accelerator.device_count() < AR_TP_SIZE + DIT_TP_SIZE,
     reason=f"Needs {AR_TP_SIZE + DIT_TP_SIZE}+ GPUs ({AR_TP_SIZE} AR + {DIT_TP_SIZE} DiT)",
 )
-def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None:
+def test_image_to_image_alignment(accuracy_artifact_root: Path) -> None:
     if importlib.util.find_spec("FlagEmbedding") is None:
         raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding")
     from tabulate import tabulate  # lazy import
@@ -545,7 +547,7 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_
     scorer = SemanticSimilarityScorer()
     clip_scorer = CLIPScorer()
     cot_results = scorer.text_similarity(omni_cot, COT_REF)
-    image_ref = Image.open(str(accuracy_assets_root / "hunyuan_image_ref.png")).convert("RGB")
+    image_ref = Image.open(HUNYUAN_IMAGE_REF_PATH).convert("RGB")
     image_clip_score = clip_scorer.image_image_score(omni_image, image_ref)
     ssim_value, psnr_value = compute_image_ssim_psnr(prediction=omni_image, reference=image_ref, compare_mode="RGB")
 

tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

@@ -15,6 +15,7 @@ from PIL import Image
 
 from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
 from tests.helpers.runtime import OmniServer
 
 pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
@@ -33,8 +34,8 @@ P99_THRESHOLD = 3e-1
 SSIM_THRESHOLD = 0.97
 PSNR_THRESHOLD = 30.0
 
+BASELINE_PATH = get_asset_path("hunyuan/hunyuan_baseline.png")
 _REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
-BASELINE_PATH = _REPO_ROOT / "tests" / "assets" / "hunyuan" / "hunyuan_baseline.png"
 _OFFLINE_SCRIPT = _REPO_ROOT / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"
 
 # DiT-only deploy config with trust_remote_code (based on hunyuan_image3_dit.yaml).

tests/e2e/accuracy/test_qwen_image_edit.py

@@ -14,6 +14,7 @@ from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
 from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
 from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
 from tests.helpers.runtime import OmniServer
 
 pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
@@ -46,10 +47,9 @@ EDIT_2511_PROMPT = (
     "（帽子紧贴头顶，未贴合部分可调整帽子摆放角度达到完全贴合）， 头部与身体过渡必须自然，无拼接痕迹，"
     "无额外元素出现在头部四周，纯白背景，不要出现绿色圆圈"
 )
-EDIT_2511_ASSET_DIR = Path(__file__).resolve().parent / "assets"
 EDIT_2511_IMAGE_PATHS = [
-    EDIT_2511_ASSET_DIR / "qwen_image_edit_2511_test1.png",
-    EDIT_2511_ASSET_DIR / "qwen_image_edit_2511_test2.png",
+    get_asset_path("qwen_image_edit/qwen_image_edit_2511_test1.png"),
+    get_asset_path("qwen_image_edit/qwen_image_edit_2511_test2.png"),
 ]
 
 PROMPT_SINGLE_IMAGE = "The input is a 2D cartoon bear mascot. Restyle it into a painterly oil artwork with warm colors while preserving the main structure."

tests/e2e/accuracy/upstream_mock/upstream_socket_server_no_compile.py

@@ -24,7 +24,7 @@ PyTorch SDPA (`ATTENTION_BACKEND=torch`) for this subprocess only.
 Usage:
     PYTHONPATH="${DREAMZERO_REPO}" \\
       .venv/bin/python -m torch.distributed.run --standalone --nproc_per_node=2 \\
-      tests/dreamzero/upstream/upstream_socket_server_no_compile.py --port 18081 \\
+      tests/e2e/accuracy/upstream_mock/upstream_socket_server_no_compile.py --port 18081 \\
       --model_path "${DREAMZERO_REPO}/checkpoints/dreamzero"
 """
 

tests/e2e/features/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/e2e/features/comfyui/test_comfyui_integration.py

@@ -6,6 +6,8 @@ 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
@@ -52,7 +54,7 @@ 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
 

tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py

@@ -16,7 +16,7 @@ exercises when it calls ``await self.engine.list_loras()`` before
 dispatching a generation request.
 
 Usage:
-    pytest tests/e2e/offline_inference/custom_pipeline/test_async_omni_collective_rpc.py -v -s
+    pytest tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py -v -s
 """
 
 from __future__ import annotations
@@ -34,12 +34,8 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
 MODEL = "tiny-random/Qwen-Image"
-CUSTOM_PIPELINE_CLASS = (
-    "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
-)
-WORKER_EXTENSION_CLASS = (
-    "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
-)
+CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest"
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
 
 
 def _create_inline_engine() -> AsyncOmni:

tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -24,12 +24,8 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 
-CUSTOM_PIPELINE_CLASS = (
-    "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
-)
-WORKER_EXTENSION_CLASS = (
-    "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
-)
+CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest"
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
 
 # Use your specified HF repo name/path directly here
 MODEL = "tiny-random/Qwen-Image"

tests/e2e/features/custom_pipeline/test_worker_extension.py

@@ -7,37 +7,19 @@ from contextlib import ExitStack
 
 import pytest
 
-from tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob import (
+from tests.e2e.features.helpers.custom_pipeline import (
     QwenImagePipelineWithLogProbForTest,
 )
-from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
-    vLLMOmniColocateWorkerExtensionForTest,
-)
 from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 
 CUSTOM_PIPELINE_CLASS = QwenImagePipelineWithLogProbForTest
-WORKER_EXTENSION_CLASS = (
-    "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
-)
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
 MODEL = "tiny-random/Qwen-Image"
 
 pytestmark = [pytest.mark.core_model]
 
 
-@pytest.mark.cpu
-def test_worker_extension_inheritance():
-    assert issubclass(vLLMOmniColocateWorkerExtensionForTest, CustomPipelineWorkerExtension)
-
-
-@pytest.mark.cpu
-def test_worker_extension_test_function():
-    assert (
-        vLLMOmniColocateWorkerExtensionForTest.test_extension_name() == "vllm-omni-colocate-worker-extension-for-test"
-    )
-
-
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 @pytest.mark.asyncio

tests/e2e/features/helpers/custom_pipeline.py

@@ -1,27 +1,299 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Detailed custom Qwen-Image pipeline for E2E tests.
+"""Test doubles for custom-pipeline E2E and RLHF integration.
 
-This pipeline follows the structure of the user's reference implementation:
-- supports pre-tokenized prompt IDs via OmniCustomPrompt-style dict input
-- uses an SDE scheduler that can return step logprobs
-- returns structured trajectory_* fields (latents, timesteps, log_probs)
-  consistent with the BAGEL trajectory recording design
+Three layers:
+
+* ``FlowMatchSDEDiscreteSchedulerForTest`` — SDE scheduler with per-step logprobs.
+* ``vLLMOmniColocateWorkerExtensionForTest`` — verl-omni colocate worker extension mirror.
+* ``QwenImagePipelineWithLogProbForTest`` — Qwen-Image pipeline returning trajectory_* fields.
 """
 
 from __future__ import annotations
 
+import ctypes
+import logging
+import math
 import os
+import platform
+import signal
+from dataclasses import dataclass
 from typing import Any, Literal
 
 import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+from msgspec import field
 
-from tests.e2e.offline_inference.custom_pipeline.flow_match_sde_scheduler import FlowMatchSDEDiscreteSchedulerForTest
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.models.qwen_image import QwenImagePipeline
+from vllm_omni.diffusion.worker.diffusion_worker import CustomPipelineWorkerExtension
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.lora.request import LoRARequest as OmniLoRARequest
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Scheduler layer
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class FlowMatchSDEDiscreteSchedulerOutputForTest(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+        log_prob (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
+            The log probability of the previous sample.
+        prev_sample_mean (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
+            The mean of the computed sample of previous timestep.
+        std_dev_t (`torch.FloatTensor` of shape `(batch_size, 1, 1)`):
+            The standard deviation used to compute `prev_sample`.
+    """
+
+    prev_sample: torch.FloatTensor
+    log_prob: torch.FloatTensor | None
+    prev_sample_mean: torch.FloatTensor
+    std_dev_t: torch.FloatTensor
+
+
+class FlowMatchSDEDiscreteSchedulerForTest(FlowMatchEulerDiscreteScheduler):
+    """SDE version of the FlowMatchEulerDiscreteScheduler.
+    The implementation is based on FlowGRPO paper (https://arxiv.org/abs/2505.05470)
+    and diffusers v0.37 branch.
+    """
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        sample: torch.FloatTensor,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: torch.Generator | None = None,
+        per_token_timesteps: torch.Tensor | None = None,
+        return_dict: bool = True,
+        noise_level: float = 0.7,
+        prev_sample: torch.FloatTensor | None = None,
+        sde_type: Literal["sde", "cps"] = "sde",
+        return_logprobs: bool = True,
+    ) -> FlowMatchSDEDiscreteSchedulerOutputForTest | tuple:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Modified from https://github.com/yifan123/flow_grpo/blob/main/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            s_churn (`float`):
+            s_tmin  (`float`):
+            s_tmax  (`float`):
+            s_noise (`float`, defaults to 1.0):
+                Scaling factor for noise added to the sample.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            per_token_timesteps (`torch.Tensor`, *optional*):
+                The timesteps for each token in the sample.
+            return_dict (`bool`):
+                Whether or not to return a
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchSDEDiscreteSchedulerOutput`] or tuple.
+            noise_level (`float`, *optional*, defaults to 0.7):
+                The noise level used in the SDE.
+            prev_sample (`torch.FloatTensor`, *optional*):
+                The sample from the previous timestep. If not provided, it will be sampled inside the function.
+            sde_type (`str`, *optional*, defaults to "sde"):
+                The type of SDE to use. Choose between "sde" and "cps".
+            return_logprobs (`bool`, *optional*, defaults to True):
+                Whether to return log probabilities of the previous sample.
+        """
+
+        if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
+            raise ValueError(
+                (
+                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                    " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                    " one of the `scheduler.timesteps` as a timestep."
+                ),
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+        if prev_sample is not None:
+            prev_sample = prev_sample.to(torch.float32)
+
+        prev_sample, log_prob, prev_sample_mean, std_dev_t = self.sample_previous_step(
+            sample=sample,
+            model_output=model_output,
+            generator=generator,
+            per_token_timesteps=per_token_timesteps,
+            noise_level=noise_level,
+            prev_sample=prev_sample,
+            sde_type=sde_type,
+            return_logprobs=return_logprobs,
+        )
+
+        # upon completion increase step index by one
+        self._step_index += 1
+        if per_token_timesteps is None:
+            # Cast sample back to model compatible dtype
+            prev_sample = prev_sample.to(model_output.dtype)
+
+        if not return_dict:
+            return (prev_sample, log_prob, prev_sample_mean, std_dev_t)
+
+        return FlowMatchSDEDiscreteSchedulerOutputForTest(
+            prev_sample=prev_sample, log_prob=log_prob, prev_sample_mean=prev_sample_mean, std_dev_t=std_dev_t
+        )
+
+    def sample_previous_step(
+        self,
+        sample: torch.Tensor,
+        model_output: torch.Tensor,
+        timestep: torch.FloatTensor | None = None,
+        generator: torch.Generator | None = None,
+        per_token_timesteps: torch.Tensor | None = None,
+        noise_level: float = 0.7,
+        prev_sample: torch.Tensor | None = None,
+        sde_type: Literal["cps", "sde"] = "sde",
+        return_logprobs: bool = True,
+    ):
+        assert sde_type in ["sde", "cps"]
+        assert sample.dtype == torch.float32
+        if prev_sample is not None:
+            assert prev_sample.dtype == torch.float32
+
+        if per_token_timesteps is not None:
+            raise NotImplementedError("per_token_timesteps is not supported yet for FlowMatchSDEDiscreteScheduler.")
+        else:
+            if timestep is None:
+                sigma_idx = self.step_index
+                sigma = self.sigmas[sigma_idx]
+                sigma_prev = self.sigmas[sigma_idx + 1]
+            else:
+                sigma_idx = torch.tensor([self.index_for_timestep(t) for t in timestep])
+                sigma = self.sigmas[sigma_idx].view(-1, *([1] * (len(sample.shape) - 1)))
+                sigma_prev = self.sigmas[sigma_idx + 1].view(-1, *([1] * (len(sample.shape) - 1)))
+
+            sigma_max = self.sigmas[1]
+            dt = sigma_prev - sigma
+
+        if sde_type == "sde":
+            std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * noise_level
+
+            prev_sample_mean = (
+                sample * (1 + std_dev_t**2 / (2 * sigma) * dt)
+                + model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt
+            )
+
+            if prev_sample is None:
+                variance_noise = randn_tensor(
+                    model_output.shape,
+                    generator=generator,
+                    device=model_output.device,
+                    dtype=model_output.dtype,
+                )
+                prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1 * dt) * variance_noise
+
+            if return_logprobs:
+                log_prob = (
+                    -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2))
+                    - torch.log(std_dev_t * torch.sqrt(-1 * dt))
+                    - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
+                )
+            else:
+                log_prob = None
+
+        elif sde_type == "cps":
+            std_dev_t = sigma_prev * math.sin(noise_level * math.pi / 2)
+            pred_original_sample = sample - sigma * model_output
+            noise_estimate = sample + model_output * (1 - sigma)
+            prev_sample_mean = pred_original_sample * (1 - sigma_prev) + noise_estimate * torch.sqrt(
+                sigma_prev**2 - std_dev_t**2
+            )
+
+            if prev_sample is None:
+                variance_noise = randn_tensor(
+                    model_output.shape,
+                    generator=generator,
+                    device=model_output.device,
+                    dtype=model_output.dtype,
+                )
+                prev_sample = prev_sample_mean + std_dev_t * variance_noise
+
+            if return_logprobs:
+                log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2)
+            else:
+                log_prob = None
+
+        # mean along all but batch dimension
+        log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) if log_prob is not None else None
+        return prev_sample, log_prob, prev_sample_mean, std_dev_t
+
+
+# ---------------------------------------------------------------------------
+# Worker extension layer
+# ---------------------------------------------------------------------------
+
+
+class OmniTensorLoRARequestForTest(OmniLoRARequest):
+    peft_config: dict = field(default=None)
+    lora_tensors: dict = field(default=None)
+
+
+def set_death_signal() -> None:
+    """Verbatim port of ``verl.workers.rollout.vllm_rollout.utils.set_death_signal``."""
+    if platform.system() != "Linux":
+        return
+    try:
+        libc = ctypes.CDLL("libc.so.6")
+        libc.prctl(1, signal.SIGKILL)
+        if os.getppid() == 1:
+            os.kill(os.getpid(), signal.SIGKILL)
+    except Exception:  # noqa: BLE001
+        # libc.so.6 may not be present (musl, alpine) — best-effort only.
+        pass
+
+
+class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension):
+    """Mirror of ``vLLMOmniColocateWorkerExtension`` (verl-omni).
+
+    The production ``__new__`` runs ``set_death_signal`` + ``VLLMOmniHijack.hijack``
+    on every vLLM worker process. Replicating both ensures this test
+    reproduces the same monkey-patched environment.
+    """
+
+    def __new__(cls, **kwargs):
+        set_death_signal()
+        return super().__new__(cls)
+
+    @staticmethod
+    def test_extension_name() -> str:
+        """Return a stable identifier for assertions in unit tests."""
+        return "vllm-omni-colocate-worker-extension-for-test"
+
+
+# ---------------------------------------------------------------------------
+# Pipeline layer
+# ---------------------------------------------------------------------------
 
 
 def _maybe_to_cpu(v):
@@ -30,8 +302,6 @@ def _maybe_to_cpu(v):
     return v
 
 
-# Custom pipeline class for QwenImage that returns log probabilities during the diffusion process.
-# This is for test
 class QwenImagePipelineWithLogProbForTest(QwenImagePipeline):
     supports_request_batch = False
 

tests/e2e/features/helpers/verl_omni_server.py

@@ -1,30 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""E2E test that follows the EXACT same flow as
-``verl-omni/tests/workers/rollout/rollout_vllm/test_vllm_omni_generate.py``
-but inlines all the ``verl`` / ``verl_omni`` symbols the flow depends on
-(stripped down to only what the diffusion-mode test actually exercises).
+"""Ray-serializable verl-omni HTTP server test double for RLHF E2E tests.
 
-Flow (identical to the verl-omni test):
-
-1. Build two ``OmegaConf`` configs (``DiffusionRolloutConfig`` /
-   ``DiffusionModelConfig`` shape).
-2. ``ServerCls = ray.remote(vLLMOmniHttpServerLocal)`` then
-   ``ServerCls.options(...).remote(config=..., model_config=...,
-   rollout_mode=RolloutMode.STANDALONE, workers=[], replica_rank=0,
-   node_rank=0, gpus_per_node=1, nnodes=1, cuda_visible_devices="0")``.
-3. ``ray.get(server.launch_server.remote())``.
-4. ``server.generate.remote(prompt_ids=..., sampling_params=..., request_id=...)``
-   and assert the returned ``DiffusionOutput``.
-
-The inlined classes / helpers mirror their ``verl`` / ``verl_omni``
-counterparts; bits not reached by this test (AR mode, LoRA,
-quantization, run_uvicorn, headless, profiler, port wiring) are removed.
-
-Usage:
-    pytest tests/e2e/offline_inference/rlhf_test/\
-test_ray_async_omni_qwen_image_generate.py -v -s
+The actor class must live in an importable module (not a pytest test file) so
+Ray workers can unpickle ``vLLMOmniHttpServerLocal`` via
+``tests.e2e.features.helpers.verl_omni_server``.
 """
 
 from __future__ import annotations
@@ -35,23 +16,15 @@ import logging
 import os
 from dataclasses import asdict
 from enum import Enum
-from functools import lru_cache
 from typing import Any
-from uuid import uuid4
 
-import pytest
-import ray
-import torch
 import torchvision.transforms as T
-from huggingface_hub import snapshot_download
 from omegaconf import DictConfig, OmegaConf
 from pydantic import BaseModel, ConfigDict
-from transformers import AutoTokenizer
 from vllm.entrypoints.openai.api_server import build_app
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 import vllm_omni.entrypoints.cli.serve
-from tests.helpers.mark import hardware_test
 from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.api_server import omni_init_app_state
@@ -60,33 +33,8 @@ from vllm_omni.outputs import OmniRequestOutput
 
 logger = logging.getLogger(__name__)
 
-CUSTOM_PIPELINE_CLASS = (
-    "tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob.QwenImagePipelineWithLogProbForTest"
-)
-WORKER_EXTENSION_CLASS = (
-    "tests.e2e.offline_inference.custom_pipeline.worker_extension.vLLMOmniColocateWorkerExtensionForTest"
-)
-
-# Use your specified HF repo name/path directly here
-MODEL = "tiny-random/Qwen-Image"
-
-TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct"
-
-
-def _resolve_model_path(repo_id: str) -> str:
-    """Resolve an HF repo ID to a local snapshot path, downloading if needed."""
-    # Allow overriding with a pre-existing local path (skips download).
-    if os.path.isdir(repo_id):
-        return repo_id
-    return snapshot_download(repo_id=repo_id)
-
-
-_MIN_PROMPT_TOKENS = 35
-
-
-# ---------------------------------------------------------------------
-#       Inlined: verl.workers.rollout.replica.RolloutMode
-# ---------------------------------------------------------------------
+CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest"
+WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest"
 
 
 class RolloutMode(Enum):
@@ -95,11 +43,6 @@ class RolloutMode(Enum):
     STANDALONE = "standalone"
 
 
-# ---------------------------------------------------------------------
-#       Inlined: verl_omni.workers.rollout.replica.DiffusionOutput
-# ---------------------------------------------------------------------
-
-
 class DiffusionOutput(BaseModel):
     """Pydantic mirror of ``verl_omni.workers.rollout.replica.DiffusionOutput``."""
 
@@ -112,11 +55,6 @@ class DiffusionOutput(BaseModel):
     extra_fields: dict[str, Any] = {}
 
 
-# ---------------------------------------------------------------------
-#       Inlined: verl.utils.tokenizer.normalize_token_ids
-# ---------------------------------------------------------------------
-
-
 def normalize_token_ids(tokenized_output) -> list[int]:
     token_ids = tokenized_output
     if isinstance(tokenized_output, dict):
@@ -142,11 +80,6 @@ def normalize_token_ids(tokenized_output) -> list[int]:
     return out
 
 
-# ---------------------------------------------------------------------
-#  Inlined: verl.workers.rollout.vllm_rollout.utils.build_cli_args_from_config
-# ---------------------------------------------------------------------
-
-
 def build_cli_args_from_config(config: dict[str, Any]) -> list[str]:
     cli_args: list[str] = []
     for k, v in config.items():
@@ -166,34 +99,6 @@ def build_cli_args_from_config(config: dict[str, Any]) -> list[str]:
     return cli_args
 
 
-# ---------------------------------------------------------------------
-#                  Tokenizer helper (test-side)
-# ---------------------------------------------------------------------
-
-
-@lru_cache(maxsize=1)
-def _get_tokenizer():
-    return AutoTokenizer.from_pretrained(TOKENIZER_MODEL, trust_remote_code=True)
-
-
-def _tokenize_prompt(text: str) -> list[int]:
-    tokenizer = _get_tokenizer()
-    messages = [{"role": "user", "content": text}]
-    token_ids = normalize_token_ids(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False))
-    assert len(token_ids) > _MIN_PROMPT_TOKENS, (
-        f"Prompt too short ({len(token_ids)} tokens, need >{_MIN_PROMPT_TOKENS}). "
-        "The pipeline drops the first 34 chat-template prefix tokens; "
-        "use a longer prompt so content tokens remain after the drop."
-    )
-    return token_ids
-
-
-# ---------------------------------------------------------------------
-#  Inlined: vLLMOmniHttpServer (verl-omni) + relevant vLLMHttpServer (verl)
-#  bits, stripped to the diffusion-mode path the test exercises.
-# ---------------------------------------------------------------------
-
-
 def _cfg_get(cfg: Any, key: str, default: Any = None) -> Any:
     """Uniform getter that works for both DictConfig and dataclass-like objects."""
     if isinstance(cfg, DictConfig):
@@ -202,13 +107,8 @@ def _cfg_get(cfg: Any, key: str, default: Any = None) -> Any:
 
 
 class vLLMOmniHttpServerLocal:
-    """Self-contained simulation of ``vLLMOmniHttpServer`` for diffusion mode.
+    """Self-contained simulation of ``vLLMOmniHttpServer`` for diffusion mode."""
 
-    Constructor signature matches ``vLLMHttpServer.__init__`` exactly so the
-    Ray-actor call site in the test is identical to the verl-omni one.
-    """
-
-    # ------------------------------------------------------------------ init
     def __init__(
         self,
         config,
@@ -224,8 +124,6 @@ class vLLMOmniHttpServerLocal:
         os.environ.setdefault("CUDA_VISIBLE_DEVICES", cuda_visible_devices)
         os.environ["VERL_REPLICA_RANK"] = str(replica_rank)
 
-        # Resolve configs. verl uses omega_conf_to_dataclass; here we keep
-        # them as DictConfig — every read site below uses ``_cfg_get``.
         self.config = OmegaConf.create(config) if not isinstance(config, DictConfig) else config
         self.model_config = OmegaConf.create(model_config) if not isinstance(model_config, DictConfig) else model_config
 
@@ -237,8 +135,6 @@ class vLLMOmniHttpServerLocal:
         self.nnodes = nnodes
         self.global_steps = 0
         self.engine: AsyncOmni | None = None
-
-        # Diffusion-mode parity with ``vLLMOmniHttpServer._post_init``.
         self._to_tensor = T.PILToTensor()
 
         logger.info(
@@ -249,11 +145,7 @@ class vLLMOmniHttpServerLocal:
             cuda_visible_devices,
         )
 
-    # ---------------------------------------------------------- launch_server
     async def launch_server(self) -> None:
-        """Diffusion-only mirror of ``vLLMHttpServer.launch_server``."""
-
-        # vLLMOmniHttpServer._get_engine_kwargs_key() returns "vllm_omni".
         engine_kwargs_root = _cfg_get(self.config, "engine_kwargs", {}) or {}
         if isinstance(engine_kwargs_root, (dict, DictConfig)):
             engine_kwargs = engine_kwargs_root.get("vllm_omni", {}) or {}
@@ -261,10 +153,8 @@ class vLLMOmniHttpServerLocal:
             engine_kwargs = {}
         engine_kwargs = {k: v for k, v in dict(engine_kwargs).items() if v is not None}
 
-        # vLLMOmniHttpServer._preprocess_engine_kwargs (diffusion path).
         engine_kwargs.pop("output_mode", None)
 
-        # ``compilation_config`` handling — verbatim from parent launch_server.
         compilation_config = engine_kwargs.pop("compilation_config", None) or {}
         if isinstance(compilation_config, str):
             compilation_config = json.loads(compilation_config)
@@ -273,8 +163,6 @@ class vLLMOmniHttpServerLocal:
         compilation_config.setdefault("cudagraph_mode", "FULL_AND_PIECEWISE")
         compilation_config = json.dumps(compilation_config)
 
-        # Args dict (subset of parent's; the diffusion override drops AR-only
-        # entries via ``_get_override_generation_config`` returning ``{}``).
         args = {
             "dtype": _cfg_get(self.config, "dtype", "bfloat16"),
             "load_format": _cfg_get(self.config, "load_format", "auto"),
@@ -303,8 +191,6 @@ class vLLMOmniHttpServerLocal:
         model_path = str(_cfg_get(self.model_config, "local_path", None) or _cfg_get(self.model_config, "path"))
         server_args_list = ["serve", model_path] + build_cli_args_from_config(args)
 
-        # Same parser pipeline as parent launch_server (with vllm-omni's
-        # cmd_init injected, matching ``_get_cli_modules`` + ``_get_cli_description``).
         parser = FlexibleArgumentParser(description="vLLM-Omni CLI")
         subparsers = parser.add_subparsers(required=False, dest="subparser")
         cmds: dict[str, Any] = {}
@@ -316,19 +202,12 @@ class vLLMOmniHttpServerLocal:
         if parsed.subparser in cmds:
             cmds[parsed.subparser].validate(parsed)
 
-        # node_rank == 0 → run_server (the test never sets node_rank != 0).
         await self.run_server(parsed)
 
-    # ------------------------------------------------------------- run_server
     async def run_server(self, args: argparse.Namespace) -> None:
-        """Diffusion branch of ``vLLMOmniHttpServer.run_server``."""
-
         engine_args = OmniEngineArgs.from_cli_args(args)
         engine_args = asdict(engine_args)
 
-        # Diffusion path: pin the in-repo test pipeline. The production code
-        # calls ``VllmOmniPipelineBase.get_pipeline_path(architecture)``; we
-        # bypass the registry and bind the test pipeline directly.
         engine_args["enable_dummy_pipeline"] = True
         engine_args["custom_pipeline_args"] = {"pipeline_class": CUSTOM_PIPELINE_CLASS}
 
@@ -337,10 +216,7 @@ class vLLMOmniHttpServerLocal:
         await omni_init_app_state(engine_client, app.state, args)
 
         self.engine = engine_client
-        # ``run_uvicorn`` intentionally skipped — the test drives the engine
-        # via Ray RPC, not HTTP.
 
-    # ---------------------------------------------------------------- generate
     async def generate(
         self,
         prompt_ids: list[int],
@@ -351,8 +227,6 @@ class vLLMOmniHttpServerLocal:
         negative_prompt_ids: list[int] | None = None,
         priority: int = 0,  # noqa: ARG002 (signature parity)
     ) -> DiffusionOutput:
-        # Verbatim copy of ``vLLMOmniHttpServer._generate_diffusion`` minus
-        # the LoRA branch (the test does not enable ``lora_as_adapter``).
         prompt_ids = normalize_token_ids(prompt_ids)
 
         multi_modal_data: dict[str, Any] = {}
@@ -444,149 +318,3 @@ class vLLMOmniHttpServerLocal:
             num_preempted=num_preempted,
             extra_fields=extra_fields,
         )
-
-
-# ---------------------------------------------------------------------
-#                          Fixtures & tests
-# ---------------------------------------------------------------------
-
-
-@pytest.fixture
-def init_server():
-    """Create and launch a vLLMOmniHttpServerLocal Ray actor with Qwen-Image."""
-    model_path = _resolve_model_path(MODEL)
-
-    ray.init(
-        runtime_env={
-            "env_vars": {
-                "TOKENIZERS_PARALLELISM": "true",
-                "NCCL_DEBUG": "WARN",
-                "VLLM_LOGGING_LEVEL": "INFO",
-            }
-        },
-        ignore_reinit_error=True,
-    )
-
-    rollout_cfg = OmegaConf.create(
-        {
-            "name": "vllm_omni",
-            "mode": "async",
-            "tensor_model_parallel_size": 1,
-            "data_parallel_size": 1,
-            "pipeline_model_parallel_size": 1,
-            "gpu_memory_utilization": 0.8,
-            "max_num_batched_tokens": 8192,
-            # The custom log-prob pipeline only supports serial requests.
-            "max_num_seqs": 1,
-            "max_model_len": 1058,
-            "dtype": "bfloat16",
-            "load_format": "auto",
-            "enforce_eager": True,
-            "enable_chunked_prefill": False,
-            "enable_prefix_caching": False,
-            "enable_sleep_mode": False,
-            "free_cache_engine": True,
-            "disable_log_stats": True,
-            "logprobs_mode": "processed_logprobs",
-            "scheduling_policy": "fcfs",
-            "seed": 0,
-            "n": 4,
-            "pipeline": {
-                "height": 512,
-                "width": 512,
-                "num_inference_steps": 10,
-            },
-        }
-    )
-
-    model_cfg = OmegaConf.create(
-        {
-            "path": str(model_path),
-            "local_path": str(model_path),
-            "tokenizer_path": TOKENIZER_MODEL,
-            "trust_remote_code": True,
-            "load_tokenizer": True,
-        }
-    )
-
-    ServerCls = ray.remote(vLLMOmniHttpServerLocal)
-    server = ServerCls.options(
-        runtime_env={
-            "env_vars": {
-                "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
-                "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
-                "NCCL_CUMEM_ENABLE": "0",
-            }
-        },
-        max_concurrency=10,
-        num_gpus=1,
-    ).remote(
-        config=rollout_cfg,
-        model_config=model_cfg,
-        rollout_mode=RolloutMode.STANDALONE,
-        workers=[],
-        replica_rank=0,
-        node_rank=0,
-        gpus_per_node=1,
-        nnodes=1,
-        cuda_visible_devices="0",
-    )
-
-    ray.get(server.launch_server.remote())
-
-    yield server
-
-    ray.shutdown()
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_generate(init_server):
-    """Concurrent generate() calls covering basic output, logprobs, and multi-request correctness."""
-    server = init_server
-
-    prompts = [
-        "a beautiful sunset over the ocean with vibrant orange and purple clouds "
-        "reflecting on the calm water surface near a rocky coastline",
-        "a fluffy orange cat sitting on a wooden windowsill looking outside at "
-        "a garden full of colorful flowers on a bright sunny afternoon",
-        "a majestic mountain landscape covered with fresh white snow under a "
-        "clear blue sky with pine trees in the foreground and a frozen lake",
-        "a futuristic city at night with neon lights glowing on tall glass "
-        "skyscrapers and flying vehicles soaring between the buildings",
-    ]
-
-    refs = []
-    for i, prompt in enumerate(prompts):
-        rid = f"test_{i}_{uuid4().hex[:8]}"
-        ref = server.generate.remote(
-            prompt_ids=_tokenize_prompt(prompt),
-            sampling_params={
-                "num_inference_steps": 10,
-                "true_cfg_scale": 4.0,
-                "height": 512,
-                "width": 512,
-                "logprobs": i == 0,  # first request includes logprobs
-            },
-            request_id=rid,
-        )
-        refs.append(ref)
-
-    results = ray.get(refs, timeout=600)
-
-    for i, output in enumerate(results):
-        assert isinstance(output, DiffusionOutput), f"Request {i}: expected DiffusionOutput"
-        img = output.diffusion_output
-        assert isinstance(img, torch.Tensor) and img.ndim == 3, (
-            f"Request {i}: expected CHW torch.Tensor, got {type(img).__name__} ndim={getattr(img, 'ndim', None)}"
-        )
-        c, h, w = img.shape
-        assert c == 3, f"Request {i}: expected 3 channels (CHW), got {c}"
-        assert h > 0 and w > 0, f"Request {i}: image dimensions must be positive"
-        assert output.stop_reason in ("completed", "aborted", None), (
-            f"Request {i}: unexpected stop_reason {output.stop_reason!r}"
-        )
-        assert 0.0 <= float(img[0, 0, 0]) <= 1.0, f"Request {i}: pixel values must be in [0, 1]"
-
-    print(f"All {len(prompts)} concurrent requests returned valid DiffusionOutput")

tests/e2e/features/rlhf_test/test_verl_omni_e2e.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E test that follows the EXACT same flow as
+``verl-omni/tests/workers/rollout/rollout_vllm/test_vllm_omni_generate.py``
+but inlines all the ``verl`` / ``verl_omni`` symbols the flow depends on
+(stripped down to only what the diffusion-mode test actually exercises).
+
+Flow (identical to the verl-omni test):
+
+1. Build two ``OmegaConf`` configs (``DiffusionRolloutConfig`` /
+   ``DiffusionModelConfig`` shape).
+2. ``ServerCls = ray.remote(vLLMOmniHttpServerLocal)`` then
+   ``ServerCls.options(...).remote(config=..., model_config=...,
+   rollout_mode=RolloutMode.STANDALONE, workers=[], replica_rank=0,
+   node_rank=0, gpus_per_node=1, nnodes=1, cuda_visible_devices="0")``.
+3. ``ray.get(server.launch_server.remote())``.
+4. ``server.generate.remote(prompt_ids=..., sampling_params=..., request_id=...)``
+   and assert the returned ``DiffusionOutput``.
+
+The Ray actor lives in ``tests.e2e.features.helpers.verl_omni_server`` so
+workers can import it after directory migration (pytest test modules are not
+importable as ``test_verl_omni_e2e`` in Ray child processes).
+
+Usage:
+    pytest tests/e2e/features/rlhf_test/test_verl_omni_e2e.py -v -s
+"""
+
+from __future__ import annotations
+
+import os
+from functools import lru_cache
+from uuid import uuid4
+
+import pytest
+import ray
+import torch
+from huggingface_hub import snapshot_download
+from omegaconf import OmegaConf
+from transformers import AutoTokenizer
+
+from tests.e2e.features.helpers.verl_omni_server import (
+    DiffusionOutput,
+    RolloutMode,
+    normalize_token_ids,
+    vLLMOmniHttpServerLocal,
+)
+from tests.helpers.mark import hardware_test
+
+MODEL = "tiny-random/Qwen-Image"
+TOKENIZER_MODEL = "Qwen/Qwen2-1.5B-Instruct"
+_MIN_PROMPT_TOKENS = 35
+
+
+def _resolve_model_path(repo_id: str) -> str:
+    if os.path.isdir(repo_id):
+        return repo_id
+    return snapshot_download(repo_id=repo_id)
+
+
+@lru_cache(maxsize=1)
+def _get_tokenizer():
+    return AutoTokenizer.from_pretrained(TOKENIZER_MODEL, trust_remote_code=True)
+
+
+def _tokenize_prompt(text: str) -> list[int]:
+    tokenizer = _get_tokenizer()
+    messages = [{"role": "user", "content": text}]
+    token_ids = normalize_token_ids(tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False))
+    assert len(token_ids) > _MIN_PROMPT_TOKENS, (
+        f"Prompt too short ({len(token_ids)} tokens, need >{_MIN_PROMPT_TOKENS}). "
+        "The pipeline drops the first 34 chat-template prefix tokens; "
+        "use a longer prompt so content tokens remain after the drop."
+    )
+    return token_ids
+
+
+@pytest.fixture
+def init_server():
+    """Create and launch a vLLMOmniHttpServerLocal Ray actor with Qwen-Image."""
+    model_path = _resolve_model_path(MODEL)
+
+    ray.init(
+        runtime_env={
+            "env_vars": {
+                "TOKENIZERS_PARALLELISM": "true",
+                "NCCL_DEBUG": "WARN",
+                "VLLM_LOGGING_LEVEL": "INFO",
+            }
+        },
+        ignore_reinit_error=True,
+    )
+
+    rollout_cfg = OmegaConf.create(
+        {
+            "name": "vllm_omni",
+            "mode": "async",
+            "tensor_model_parallel_size": 1,
+            "data_parallel_size": 1,
+            "pipeline_model_parallel_size": 1,
+            "gpu_memory_utilization": 0.8,
+            "max_num_batched_tokens": 8192,
+            # QwenImagePipeline (and the logprob custom subclass) only supports
+            # serial request execution; max_num_seqs>1 enables request-level batching.
+            "max_num_seqs": 1,
+            "max_model_len": 1058,
+            "dtype": "bfloat16",
+            "load_format": "auto",
+            "enforce_eager": True,
+            "enable_chunked_prefill": False,
+            "enable_prefix_caching": False,
+            "enable_sleep_mode": False,
+            "free_cache_engine": True,
+            "disable_log_stats": True,
+            "logprobs_mode": "processed_logprobs",
+            "scheduling_policy": "fcfs",
+            "seed": 0,
+            "n": 4,
+            "pipeline": {
+                "height": 512,
+                "width": 512,
+                "num_inference_steps": 10,
+            },
+        }
+    )
+
+    model_cfg = OmegaConf.create(
+        {
+            "path": str(model_path),
+            "local_path": str(model_path),
+            "tokenizer_path": TOKENIZER_MODEL,
+            "trust_remote_code": True,
+            "load_tokenizer": True,
+        }
+    )
+
+    ServerCls = ray.remote(vLLMOmniHttpServerLocal)
+    server = ServerCls.options(
+        runtime_env={
+            "env_vars": {
+                "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
+                "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
+                "NCCL_CUMEM_ENABLE": "0",
+            }
+        },
+        max_concurrency=10,
+        num_gpus=1,
+    ).remote(
+        config=rollout_cfg,
+        model_config=model_cfg,
+        rollout_mode=RolloutMode.STANDALONE,
+        workers=[],
+        replica_rank=0,
+        node_rank=0,
+        gpus_per_node=1,
+        nnodes=1,
+        cuda_visible_devices="0",
+    )
+
+    ray.get(server.launch_server.remote())
+
+    yield server
+
+    ray.shutdown()
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_generate(init_server):
+    """Concurrent generate() calls covering basic output, logprobs, and multi-request correctness."""
+    server = init_server
+
+    prompts = [
+        "a beautiful sunset over the ocean with vibrant orange and purple clouds "
+        "reflecting on the calm water surface near a rocky coastline",
+        "a fluffy orange cat sitting on a wooden windowsill looking outside at "
+        "a garden full of colorful flowers on a bright sunny afternoon",
+        "a majestic mountain landscape covered with fresh white snow under a "
+        "clear blue sky with pine trees in the foreground and a frozen lake",
+        "a futuristic city at night with neon lights glowing on tall glass "
+        "skyscrapers and flying vehicles soaring between the buildings",
+    ]
+
+    refs = []
+    for i, prompt in enumerate(prompts):
+        rid = f"test_{i}_{uuid4().hex[:8]}"
+        ref = server.generate.remote(
+            prompt_ids=_tokenize_prompt(prompt),
+            sampling_params={
+                "num_inference_steps": 10,
+                "true_cfg_scale": 4.0,
+                "height": 512,
+                "width": 512,
+                "logprobs": i == 0,
+            },
+            request_id=rid,
+        )
+        refs.append(ref)
+
+    results = ray.get(refs, timeout=600)
+
+    for i, output in enumerate(results):
+        assert isinstance(output, DiffusionOutput), f"Request {i}: expected DiffusionOutput"
+        img = output.diffusion_output
+        assert isinstance(img, torch.Tensor) and img.ndim == 3, (
+            f"Request {i}: expected CHW torch.Tensor, got {type(img).__name__} ndim={getattr(img, 'ndim', None)}"
+        )
+        c, h, w = img.shape
+        assert c == 3, f"Request {i}: expected 3 channels (CHW), got {c}"
+        assert h > 0 and w > 0, f"Request {i}: image dimensions must be positive"
+        assert output.stop_reason in ("completed", "aborted", None), (
+            f"Request {i}: unexpected stop_reason {output.stop_reason!r}"
+        )
+        assert 0.0 <= float(img[0, 0, 0]) <= 1.0, f"Request {i}: pixel values must be in [0, 1]"
+
+    print(f"All {len(prompts)} concurrent requests returned valid DiffusionOutput")

tests/e2e/offline_inference/custom_pipeline/__init__.py

@@ -1,14 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from tests.e2e.offline_inference.custom_pipeline.qwen_image_pipeline_with_logprob import (
-    QwenImagePipelineWithLogProbForTest,
-)
-from tests.e2e.offline_inference.custom_pipeline.worker_extension import (
-    vLLMOmniColocateWorkerExtensionForTest,
-)
-
-__all__ = [
-    "QwenImagePipelineWithLogProbForTest",
-    "vLLMOmniColocateWorkerExtensionForTest",
-]

tests/e2e/offline_inference/custom_pipeline/flow_match_sde_scheduler.py

@@ -1,220 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from __future__ import annotations
-
-import math
-from dataclasses import dataclass
-from typing import Literal
-
-import torch
-from diffusers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import BaseOutput
-from diffusers.utils.torch_utils import randn_tensor
-
-
-@dataclass
-class FlowMatchSDEDiscreteSchedulerOutputForTest(BaseOutput):
-    """
-    Output class for the scheduler's `step` function output.
-
-    Args:
-        prev_sample (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
-            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
-            denoising loop.
-        log_prob (`torch.FloatTensor` of shape `(batch_size,)`, *optional*):
-            The log probability of the previous sample.
-        prev_sample_mean (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_channels)` for images):
-            The mean of the computed sample of previous timestep.
-        std_dev_t (`torch.FloatTensor` of shape `(batch_size, 1, 1)`):
-            The standard deviation used to compute `prev_sample`.
-    """
-
-    prev_sample: torch.FloatTensor
-    log_prob: torch.FloatTensor | None
-    prev_sample_mean: torch.FloatTensor
-    std_dev_t: torch.FloatTensor
-
-
-class FlowMatchSDEDiscreteSchedulerForTest(FlowMatchEulerDiscreteScheduler):
-    """SDE version of the FlowMatchEulerDiscreteScheduler.
-    The implementation is based on FlowGRPO paper (https://arxiv.org/abs/2505.05470)
-    and diffusers v0.37 branch.
-    """
-
-    def step(
-        self,
-        model_output: torch.FloatTensor,
-        timestep: float | torch.FloatTensor,
-        sample: torch.FloatTensor,
-        s_churn: float = 0.0,
-        s_tmin: float = 0.0,
-        s_tmax: float = float("inf"),
-        s_noise: float = 1.0,
-        generator: torch.Generator | None = None,
-        per_token_timesteps: torch.Tensor | None = None,
-        return_dict: bool = True,
-        noise_level: float = 0.7,
-        prev_sample: torch.FloatTensor | None = None,
-        sde_type: Literal["sde", "cps"] = "sde",
-        return_logprobs: bool = True,
-    ) -> FlowMatchSDEDiscreteSchedulerOutputForTest | tuple:
-        """
-        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
-        process from the learned model outputs (most often the predicted noise).
-
-        Modified from https://github.com/yifan123/flow_grpo/blob/main/flow_grpo/diffusers_patch/sd3_sde_with_logprob.py
-
-        Args:
-            model_output (`torch.FloatTensor`):
-                The direct output from learned diffusion model.
-            timestep (`float`):
-                The current discrete timestep in the diffusion chain.
-            sample (`torch.FloatTensor`):
-                A current instance of a sample created by the diffusion process.
-            s_churn (`float`):
-            s_tmin  (`float`):
-            s_tmax  (`float`):
-            s_noise (`float`, defaults to 1.0):
-                Scaling factor for noise added to the sample.
-            generator (`torch.Generator`, *optional*):
-                A random number generator.
-            per_token_timesteps (`torch.Tensor`, *optional*):
-                The timesteps for each token in the sample.
-            return_dict (`bool`):
-                Whether or not to return a
-                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchSDEDiscreteSchedulerOutput`] or tuple.
-            noise_level (`float`, *optional*, defaults to 0.7):
-                The noise level used in the SDE.
-            prev_sample (`torch.FloatTensor`, *optional*):
-                The sample from the previous timestep. If not provided, it will be sampled inside the function.
-            sde_type (`str`, *optional*, defaults to "sde"):
-                The type of SDE to use. Choose between "sde" and "cps".
-            return_logprobs (`bool`, *optional*, defaults to True):
-                Whether to return log probabilities of the previous sample.
-        """
-
-        if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
-            raise ValueError(
-                (
-                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
-                    " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
-                    " one of the `scheduler.timesteps` as a timestep."
-                ),
-            )
-
-        if self.step_index is None:
-            self._init_step_index(timestep)
-
-        # Upcast to avoid precision issues when computing prev_sample
-        sample = sample.to(torch.float32)
-        if prev_sample is not None:
-            prev_sample = prev_sample.to(torch.float32)
-
-        prev_sample, log_prob, prev_sample_mean, std_dev_t = self.sample_previous_step(
-            sample=sample,
-            model_output=model_output,
-            generator=generator,
-            per_token_timesteps=per_token_timesteps,
-            noise_level=noise_level,
-            prev_sample=prev_sample,
-            sde_type=sde_type,
-            return_logprobs=return_logprobs,
-        )
-
-        # upon completion increase step index by one
-        self._step_index += 1
-        if per_token_timesteps is None:
-            # Cast sample back to model compatible dtype
-            prev_sample = prev_sample.to(model_output.dtype)
-
-        if not return_dict:
-            return (prev_sample, log_prob, prev_sample_mean, std_dev_t)
-
-        return FlowMatchSDEDiscreteSchedulerOutputForTest(
-            prev_sample=prev_sample, log_prob=log_prob, prev_sample_mean=prev_sample_mean, std_dev_t=std_dev_t
-        )
-
-    def sample_previous_step(
-        self,
-        sample: torch.Tensor,
-        model_output: torch.Tensor,
-        timestep: torch.FloatTensor | None = None,
-        generator: torch.Generator | None = None,
-        per_token_timesteps: torch.Tensor | None = None,
-        noise_level: float = 0.7,
-        prev_sample: torch.Tensor | None = None,
-        sde_type: Literal["cps", "sde"] = "sde",
-        return_logprobs: bool = True,
-    ):
-        assert sde_type in ["sde", "cps"]
-        assert sample.dtype == torch.float32
-        if prev_sample is not None:
-            assert prev_sample.dtype == torch.float32
-
-        if per_token_timesteps is not None:
-            raise NotImplementedError("per_token_timesteps is not supported yet for FlowMatchSDEDiscreteScheduler.")
-        else:
-            if timestep is None:
-                sigma_idx = self.step_index
-                sigma = self.sigmas[sigma_idx]
-                sigma_prev = self.sigmas[sigma_idx + 1]
-            else:
-                sigma_idx = torch.tensor([self.index_for_timestep(t) for t in timestep])
-                sigma = self.sigmas[sigma_idx].view(-1, *([1] * (len(sample.shape) - 1)))
-                sigma_prev = self.sigmas[sigma_idx + 1].view(-1, *([1] * (len(sample.shape) - 1)))
-
-            sigma_max = self.sigmas[1]
-            dt = sigma_prev - sigma
-
-        if sde_type == "sde":
-            std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * noise_level
-
-            prev_sample_mean = (
-                sample * (1 + std_dev_t**2 / (2 * sigma) * dt)
-                + model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt
-            )
-
-            if prev_sample is None:
-                variance_noise = randn_tensor(
-                    model_output.shape,
-                    generator=generator,
-                    device=model_output.device,
-                    dtype=model_output.dtype,
-                )
-                prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1 * dt) * variance_noise
-
-            if return_logprobs:
-                log_prob = (
-                    -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2))
-                    - torch.log(std_dev_t * torch.sqrt(-1 * dt))
-                    - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
-                )
-            else:
-                log_prob = None
-
-        elif sde_type == "cps":
-            std_dev_t = sigma_prev * math.sin(noise_level * math.pi / 2)
-            pred_original_sample = sample - sigma * model_output
-            noise_estimate = sample + model_output * (1 - sigma)
-            prev_sample_mean = pred_original_sample * (1 - sigma_prev) + noise_estimate * torch.sqrt(
-                sigma_prev**2 - std_dev_t**2
-            )
-
-            if prev_sample is None:
-                variance_noise = randn_tensor(
-                    model_output.shape,
-                    generator=generator,
-                    device=model_output.device,
-                    dtype=model_output.dtype,
-                )
-                prev_sample = prev_sample_mean + std_dev_t * variance_noise
-
-            if return_logprobs:
-                log_prob = -((prev_sample.detach() - prev_sample_mean) ** 2)
-            else:
-                log_prob = None
-
-        # mean along all but batch dimension
-        log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) if log_prob is not None else None
-        return prev_sample, log_prob, prev_sample_mean, std_dev_t