OmniRequestOutput inherits RequestOutput

Signed-off-by: Boao Shi <aoibosh@connect.hku.hk>

.buildkite/pipeline-npu.yaml

@@ -78,89 +78,8 @@ steps:
       echo "--- Image pushed successfully"
       echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
 
-  - label: ":buildkit: Build and Push NPU Test Image - A3"
-    key: image-build-a3
-    agents:
-      queue: "ascend-a2b3"
-      resource_class: "npu-1"
-    plugins:
-      - kubernetes:
-          metadata:
-            annotations:
-              vault.hashicorp.com/agent-init-first: "true"
-              vault.hashicorp.com/agent-inject: "true"
-              vault.hashicorp.com/agent-inject-perms-ca.pem: "0400"
-              vault.hashicorp.com/agent-inject-perms-cert.pem: "0400"
-              vault.hashicorp.com/agent-inject-perms-config.json: "0400"
-              vault.hashicorp.com/agent-inject-perms-key.pem: "0400"
-              vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd
-              vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}"
-              vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}"
-              vault.hashicorp.com/agent-pre-populate-only: "true"
-              vault.hashicorp.com/agent-run-as-group: "1000"
-              vault.hashicorp.com/agent-run-as-user: "1000"
-              vault.hashicorp.com/agent-service-account-token-volume-name: token-vol
-              vault.hashicorp.com/role: ascend-gha-runners
-              vault.hashicorp.com/secret-volume-path: /home/user/.docker/
-              vault.hashicorp.com/tls-skip-verify: "true"
-          podSpecPatch:
-            volumes:
-            - name: token-vol
-              projected:
-                defaultMode: 420
-                sources:
-                - serviceAccountToken:
-                    audience: api
-                    expirationSeconds: 600
-                    path: token
-    env:
-      VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}-a3"
-      IMAGE_NAME: "vllm-omni-ci-npu"
-      IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry"
-      BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234"
-    command: |
-      set -ex
-
-      echo "--- Building and pushing NPU Test Image"
-      echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-      echo "buildkitd address: $${BUILDKITD_ADDR}"
-
-      if ! command -v buildctl &> /dev/null; then
-        echo "Installing buildctl..."
-        mkdir -p /tmp/buildkit
-        BUILDKIT_VERSION="v0.29.0"
-        wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz
-        tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit
-        cp /tmp/buildkit/bin/buildctl /usr/local/bin/
-      fi
-
-      export DOCKER_CONFIG=/home/user/.docker
-      buildctl \
-        --addr="$${BUILDKITD_ADDR}" \
-        --tlscacert=/home/user/.docker/ca.pem \
-        --tlscert=/home/user/.docker/cert.pem \
-        --tlskey=/home/user/.docker/key.pem \
-        build \
-        --frontend dockerfile.v0 \
-        --local context=. \
-        --local dockerfile=./docker \
-        --opt filename=Dockerfile.npu.ci.a3 \
-        --secret id=dockerconfig,src=/home/user/.docker/config.json \
-        --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \
-        --progress=plain
-
-      echo "--- Image pushed successfully"
-      echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}"
-
   - label: "Upload Ready Pipeline"
-    depends_on:
-      - image-build
-      - image-build-a3
+    depends_on: image-build
     key: upload-ready-pipeline
     if: >-
       build.branch != "main" && (
@@ -175,9 +94,7 @@ steps:
 
   # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild)
   - label: "Upload Nightly Pipeline"
-    depends_on:
-      - image-build
-      - image-build-a3
+    depends_on: image-build
     key: upload-nightly-pipeline
     if: >-
       (build.branch == "main" && build.env("NIGHTLY") == "1") ||

.buildkite/test-merge.yml

@@ -173,9 +173,6 @@ steps:
           - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
           - vllm_omni/deploy/qwen3_tts.yaml
         commands:
-          # 40m (was 20m): post-rebase runs take 9-12 min warm, but cold-cache
-          # runs exceed 20m and were killed without failure output (a cold run
-          # died at 25.2m wall under the old limit).
           - |
             timeout 40m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
@@ -203,9 +200,6 @@ steps:
           - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
           - vllm_omni/deploy/qwen3_tts.yaml
         commands:
-          # 40m (was 20m): post-rebase runs take 9-13 min warm, but cold-cache
-          # runs exceed 20m and were killed without failure output (a cold run
-          # died at 25.2m wall under the old limit).
           - |
             timeout 40m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
@@ -268,13 +262,10 @@ steps:
           - vllm_omni/entrypoints/openai/tts_adapters/indextts2.py
           - vllm_omni/config/pipeline_registry.py
           - vllm_omni/model_executor/models/registry.py
-        # 40m/35m (was 20m/20m): post-rebase runs take a steady ~12.5 min warm,
-        # but a cold run was killed at 20.7m under the old limit with no
-        # failure output; the rebase-time raise to 60m/45m was oversized.
-        timeout_in_minutes: 40
+        timeout_in_minutes: 60
         commands:
           - |
-            timeout 35m bash -c "
+            timeout 45m bash -c "
               export VLLM_LOGGING_LEVEL=DEBUG
               pytest -s -v tests/e2e/online_serving/test_indextts2.py tests/e2e/offline_inference/test_indextts2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
             "
@@ -299,8 +290,6 @@ steps:
           - vllm_omni/model_executor/models/qwen3_omni/
           - vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
           - vllm_omni/deploy/qwen3_omni_moe.yaml
-        # 50m (was 30m): post-rebase runs take 15-21 min warm; 30m left no
-        # headroom for cold-cache runs, which were killed without output.
         timeout_in_minutes: 50
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'

.buildkite/test-nightly.yml

@@ -12,9 +12,6 @@ steps:
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test"
     steps:
       - label: ":full_moon: Omni · Function Test with H100"
-        # 120m (was 90m): post-rebase nightlies (#2474/#2486) run 59-71 min —
-        # too close to the old 90m ceiling, which #2444 already hit (killed at
-        # 90.8m with no failure output).
         timeout_in_minutes: 120
         commands:
           - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
@@ -156,10 +153,7 @@ steps:
 
       - label: ":full_moon: Omni · Perf Test · No Async Chunk"
         key: nightly-omni-performance-no-async-chunk
-        # 180m: post-rebase nightlies (#2474/#2486) run 41-44 min; the
-        # rebase-time raise to 300m targeted a pre-rebase 179.4m outlier and
-        # is no longer needed.
-        timeout_in_minutes: 180
+        timeout_in_minutes: 300
         commands:
           - export BENCHMARK_DIR=tests/dfx/perf/results
           - |
@@ -382,10 +376,7 @@ steps:
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
     steps:
       - label: ":full_moon: TTS · Function Test with L4"
-        # 90m: post-rebase nightlies (#2474/#2486) run ~27 min; the rebase-time
-        # raise to 120m targeted pre-rebase 72-87 min runs and is no longer
-        # needed.
-        timeout_in_minutes: 90
+        timeout_in_minutes: 120
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
           - export VLLM_USE_DEEP_GEMM="0"
@@ -870,10 +861,7 @@ steps:
 
 
       - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
-        # 60m: post-rebase nightlies (#2474/#2486) run ~18 min; the rebase-time
-        # raise to 120m targeted pre-rebase 48-58 min runs and is no longer
-        # needed.
-        timeout_in_minutes: 60
+        timeout_in_minutes: 120
         commands:
           - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
         agents:

.buildkite/test-npu-nightly.yaml

@@ -22,31 +22,3 @@ steps:
           HF_TOKEN: "${HF_TOKEN}"
         commands:
           - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"
-  - group: ":card_index_dividers: TTS Model Test"
-    key: nightly-tts-test-group
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
-    steps:
-      - label: ":full_moon: TTS · Perf Test"
-        agents:
-          queue: "ascend-a3"
-          resource_class: "npu-2"
-        image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a3"
-        plugins:
-          - kubernetes:
-              podSpecPatch:
-                imagePullSecrets:
-                  - name: swr-secret
-        env:
-          VLLM_WORKER_MULTIPROC_METHOD: spawn
-          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
-          HF_TOKEN: "${HF_TOKEN}"
-        commands:
-          - export BENCHMARK_DIR=tests/dfx/perf/results
-          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          - |
-            set +e
-            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_tts.json
-            EXIT=$$?
-            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
-            exit $$EXIT

.claude/skills/add-diffusion-model/SKILL.md

@@ -334,7 +334,7 @@ Classify the model's **CI priority** first:
 | **L1** | `tests/diffusion/models/{slug}/`, `tests/diffusion/cache/`, transformer unit tests | `core_model` + `cpu` | `test-ready.yml` | Weight remap, `_sp_plan`, cache enabler registration, shape contracts |
 | **L2** | `tests/e2e/online_serving/test_{slug}.py` (and offline if the category is offline-first) | **`core_model` + `advanced_model`** (both on baseline smoke) + `diffusion` + `@hardware_test` / `hardware_marks` | `test-ready.yml` | Default deploy smoke — minimal `num_inference_steps`, single prompt |
 | **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` when offline matters | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `diffusion`) | `test-merge.yml` **or** merged into nightly diffusion function job | Real weights, streaming/API paths, LoRA/offload smoke |
-| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py` (+ offline expansion if needed) | `full_model` + `diffusion` | `test-nightly.yml` (X2I / X2V / X2A function groups) | Feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); perf → `tests/dfx/perf/tests/test_{model}_vllm_omni.json` with per-case **`mark`** (`hardware_marks` + `full_model` + `diffusion`) |
+| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py` (+ offline expansion if needed) | `full_model` + `diffusion` | `test-nightly.yml` (X2I / X2V / X2A function groups) | Feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); perf → `tests/dfx/perf/tests/test_{model}_vllm_omni.json` |
 
 **L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, minimal steps, single prompt) should carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** function so L2 (`test-ready.yml`) and L3 (`test-merge.yml`) share one smoke test. Heavier deploy variants or API paths in the same file use **`advanced_model` only**. When L3 moves to nightly, migrate those heavier cases into `test_{slug}_expansion.py` with `full_model` and remove the dedicated `test-merge.yml` job (see `test_longcat_image_expansion.py`, `test_qwen_image_expansion.py`).
 

.claude/skills/vllm-omni-test/SKILL.md

@@ -97,7 +97,7 @@ Existing references: `tests/e2e/offline_inference/test_qwen2_5_omni.py` (L2-styl
 | **Offline inference e2e** | `tests/e2e/offline_inference/` | **Module (default):** `omni_runner` + `omni_runner_handler`. **Function (isolation only):** `omni_runner_function` + `omni_runner_handler_function`. Diffusion/TTS may use `Omni(...).generate` directly | L2: `core_model` + **one of** `omni` / `tts` / `diffusion`; `@hardware_test(...)` when GPU/NPU is required |
 | **Online serving e2e** | `tests/e2e/online_serving/` | **Module (default):** `omni_server` + `openai_client`. **Function (isolation only):** `omni_server_function` + `openai_client_function`. Clients: `send_omni_request` (omni), `send_audio_speech_request` (tts), `send_diffusion_request` / `send_video_diffusion_request` / `send_images_generations_request` (diffusion) | Baseline smoke: **`core_model` + `advanced_model`**; heavier paths: `advanced_model` only; L4 expansion: `full_model` |
 | **Documentation / runnable examples** | `tests/examples/offline_inference/`, `tests/examples/online_serving/` | **Offline docs (preferred):** extract Python/Bash blocks from the doc README (e.g. `ReadmeSnippet.extract_readme_snippets`), `pytest.mark.parametrize` each snippet, run via `example_runner.run` with a stable `output_subfolder`. **Online docs:** copy client/request scripts into dedicated tests and keep them in sync with the doc page. | Usually **L4**: `advanced_model`, often `example` plus hardware marks matching the nightly docs-example job (see `.buildkite/test-nightly.yml`). Full conventions: [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) (introduced in [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910): naming, output directory layout, skip rules, avoid trimming `num_inference_steps` without a strong CI reason). |
-| **Performance / benchmark** | `tests/dfx/perf/tests/*.json` + `run_*_benchmark.py` | JSON or script-driven server + load config; assert explicit metrics / baselines | L4 Perf: JSON `mark` with `full_model` + `omni`/`tts`/`diffusion`; wire `test-nightly.yml` Perf steps |
+| **Performance / benchmark** | `tests/dfx/perf/tests/*.json` + `run_*_benchmark.py` | JSON or script-driven server + load config; assert explicit metrics / baselines | L4 Perf: `full_model` + `benchmark`; wire `test-nightly.yml` Perf steps |
 | **Invalid parameter / negative HTTP validation** | `tests/dfx/reliability/invalid_param_test/` | Live `omni_server` + low-level `send_*_http_request` with `err_code` / `err_message` | `pytest.mark.slow` + `omni` / `tts` / `diffusion` + `@hardware_marks` (`H100` or `L4`); CI in **`test-weekly.yml`** (not ready/merge/nightly) |
 
 **If the user’s test plan includes invalid parameter validation / invalid params / negative HTTP / 400 validation:** do **not** add those `test_*` functions to `tests/e2e/online_serving/test_*.py` or `*_expansion.py`. **Move or author them** under `tests/dfx/reliability/invalid_param_test/` in the **endpoint-matching script** (see **Invalid parameter validation** below). Success-path e2e and invalid-param dfx tests must stay in separate modules.
@@ -270,7 +270,7 @@ def test_text_to_video_001(omni_server, openai_client) -> None:
 
 *Documentation example tests:* follow the **Preferred Test Strategy** in [l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md): dynamic extraction for offline READMEs; explicit copied client code for online pages until extraction is justified; use the documented **naming**, **output directory** (page folder + case id), and **skipping** rules (e.g. Gradio-only scripts).
 
-*Performance tests:* add or extend entries under `tests/dfx/perf/tests/` (and JSON configs where the project uses them), with **explicit baselines**, **`mark`** on each case (`hardware_marks` + `full_model` + type marker), and the same nightly Perf step pattern as in-tree configs.
+*Performance tests:* add or extend entries under `tests/perf/` (and JSON configs where the project uses them), with **explicit baselines** and the same marker/run-level pairing as the CI step that will execute them.
 
 **3. Cross-cutting rules**
 
@@ -592,31 +592,27 @@ Implementation strategy:
 **When the user explicitly asks for L4 perf (or full L4 including perf):**
 
 1. **Do not** put throughput/latency assertions inside `test_*_expansion.py` — perf uses the **dfx benchmark harness**, not `omni_server` e2e fixtures.
-2. Add or extend a **JSON config** under `tests/dfx/perf/tests/` (mirror in-tree names: `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json`, `test_qwen3_omni_async_chunk.json`).
-3. Each JSON **case**: `test_name`, optional **`mark`** (`hardware_marks` required when present; `marks`: `full_model` + `omni`/`tts`/`diffusion`), `server_params` (`model`, `serve_args` or `stage_config_name`), `benchmark_params[]` with **`name`** and explicit **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`, …). Diffusion cases also set `server_type` (e.g. `"vllm-omni"`) and usually `benchmark_endpoint`.
+2. Add or extend a **JSON config** under `tests/dfx/perf/tests/` (mirror in-tree names: `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`).
+3. Each JSON entry: `test_name`, `server_params` (`model`, `serve_args`), `benchmark_params[]` with explicit **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`, …).
 4. Run locally via the matching script (model type → runner):
 
 | Model type | Runner script | Example config |
 |------------|---------------|----------------|
 | **Diffusion X2I/X2V** | `tests/dfx/perf/scripts/run_diffusion_benchmark.py` | `tests/dfx/perf/tests/test_qwen_image_vllm_omni.json` |
 | **TTS** | `tests/dfx/perf/scripts/run_benchmark.py` | `tests/dfx/perf/tests/test_tts.json` (shared) or `test_{slug}.json` (dedicated; use when the model must not join the shared nightly matrix before integration — e.g. `test_voxcpm2.json`, `test_coqui_tts.json`) |
-| **Omni** | `tests/dfx/perf/scripts/run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json`, `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` |
+| **Omni** | (see existing Omni Perf steps — `run_benchmark.py` / dedicated omni configs in-tree) | `tests/dfx/perf/tests/test_omni*.json` |
 
-5. Wire **`test-nightly.yml` · Perf Test** step for that model (separate from Function Test): export `DIFFUSION_BENCHMARK_DIR` / `BENCHMARK_DIR`, run pytest on the script with **`--test-config-file`** (nightly perf steps do not use `-m`), **upload artifacts** (`buildkite-agent artifact upload`), often **multi-GPU H100** for diffusion perf.
+5. Wire **`test-nightly.yml` · Perf Test** step for that model (separate from Function Test): export `DIFFUSION_BENCHMARK_DIR` / `BENCHMARK_DIR`, run pytest on the script with `--test-config-file`, **upload artifacts** (`buildkite-agent artifact upload`), often **multi-GPU H100** for diffusion perf.
 6. **One benchmark scenario → one `test_name` block** in JSON (same spirit as one `test_*` per function case). Combine server `serve_args` + workload in one entry; use multiple `benchmark_params` rows for size/step sweeps under the same server config.
 
-**Perf local commands:**
+**Perf local command (diffusion example):**
 
 ```bash
 cd tests
 export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
 export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-# Single file (CI-like)
 pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \
   --test-config-file dfx/perf/tests/test_<model>_vllm_omni.json
-# Bulk load + filter by JSON mark
-pytest -sv dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100"
-pytest -sv dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100"
 ```
 
 **Clarify in the test plan** which L4 pillars you deliver: `Function only` | `Function + Perf` | `Function + Accuracy + Perf`.

.claude/skills/vllm-omni-test/references/test-routing.md

@@ -127,12 +127,8 @@ pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and
 cd tests
 export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
 export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
-# CI-like: single JSON file (nightly Perf Test steps)
 pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \
   --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json
-# Local bulk load: all *.json under tests/dfx/perf/tests/, filter by JSON mark
-pytest -sv dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100"
-pytest -sv dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100"
 ```
 
 Broad marker sweep (when no explicit file shard):
@@ -154,26 +150,7 @@ Nightly **L4** for a model is often **multiple jobs** in `test-nightly.yml`, not
 
 **Rule:** “L4 functional cases” → **Function only** by default. State that Perf/Accuracy are separate pillars; add them only when requested.
 
-**Diffusion perf JSON** (`tests/dfx/perf/tests/test_<slug>_vllm_omni.json`): array of **case objects**, each with:
-
-| Field | Required | Notes |
-|-------|----------|-------|
-| `test_name` | Yes | Unique server/workload id |
-| `mark` | Recommended | When present: **`hardware_marks` required** + optional `marks` (`full_model`, `diffusion`, `local_model`, …). Parsed by `resolve_pytest_marks`; enables local `-m` filtering. |
-| `server_type` | Diffusion | e.g. `"vllm-omni"` — routes case to `run_diffusion_benchmark.py` |
-| `benchmark_endpoint` | Diffusion | e.g. `/v1/videos`, `/v1/images/generations` |
-| `server_params` | Yes | `model`, `serve_args`, … |
-| `benchmark_params[]` | Yes | Each row: **`name`** (pytest id suffix), workload fields, **`baseline`** |
-
-Copy structure from `test_qwen_image_vllm_omni.json` or `test_cosmos3_vllm_omni.json` (2-GPU mark example).
-
-**Omni perf JSON**: `test_qwen3_omni_*.json` — same `mark` shape with `"marks": ["full_model", "omni"]`; run via `run_benchmark.py`.
-
-**Runner split**: `is_diffusion_perf_config()` → diffusion when `server_type` is set or `"diffusion"` ∈ `mark.marks`; otherwise omni/tts → `run_benchmark.py`.
-
-**Param ids**: `{test_name}-{benchmark_params.name}` (e.g. `test_tts-p0`, not `test_tts-0`).
-
-**Result filenames**: runtime hardware from `get_runtime_resource_label()`; `H100` omitted on default CI pool.
+**Diffusion perf JSON** (`tests/dfx/perf/tests/test_<slug>_vllm_omni.json`): array of objects with `test_name`, `server_params` (`model`, `serve_args`), `benchmark_params[]`, and per-workload **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`). Copy structure from `test_qwen_image_vllm_omni.json` in-tree.
 
 **Diffusion perf nightly step** (under **Diffusion X2I(&A&T)**; copy full `kubernetes` plugins from `· Perf Test · Qwen-Image`):
 
@@ -200,9 +177,9 @@ Copy structure from `test_qwen_image_vllm_omni.json` or `test_cosmos3_vllm_omni.
 
 | Model type | Perf runner | Config example |
 |------------|-------------|----------------|
-| Diffusion X2I/X2V | `dfx/perf/scripts/run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json` |
-| TTS | `dfx/perf/scripts/run_benchmark.py` | `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json` |
-| Omni | `dfx/perf/scripts/run_benchmark.py` | `test_qwen3_omni_async_chunk.json`, `test_qwen3_omni_no_async_chunk.json`, … |
+| Diffusion X2I/X2V | `dfx/perf/scripts/run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json` |
+| TTS | `dfx/perf/scripts/run_benchmark.py` | `test_tts.json` (`BENCHMARK_DIR`) |
+| Omni | `run_benchmark.py` / omni configs in-tree | `test_omni*.json` |
 
 Do **not** put throughput/latency baselines inside `test_*_expansion.py` — that belongs in the dfx perf JSON + nightly Perf job.
 

docker/Dockerfile.npu.ci

@@ -6,8 +6,6 @@ ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 COPY . .
 
-RUN sed -i "s@^\(deb.*\)https\?://[a-z0-9.-]*\.ubuntu\.com@\1http://cache-service.nginx-pypi-cache.svc.cluster.local:8081@g" /etc/apt/sources.list
-
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
         ffmpeg \

docker/Dockerfile.npu.ci.a3

@@ -6,8 +6,6 @@ ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 COPY . .
 
-RUN sed -i "s@^\(deb.*\)https\?://[a-z0-9.-]*\.ubuntu\.com@\1http://cache-service.nginx-pypi-cache.svc.cluster.local:8081@g" /etc/apt/sources.list
-
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
         ffmpeg \

docs/contributing/ci/CI_5levels.md

@@ -126,12 +126,11 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
         <strong>Interface tests:</strong><br>
         /tests/entrypoints/test_xxx<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>
+        /tests/dfx/perf/tests/test_qwen_omni.json (Omni), test_tts.json (TTS),<br>
+        and /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>
+        tests/example/online_serving/test_{model_name}.py<br>
+        tests/example/offline_inference/test_{model_name}.py<br>
         <strong>Accuracy Test:</strong><br>
         /tests/e2e/accuracy/test_{model_name}.py
       </td>
@@ -709,16 +708,16 @@ L4 level testing is a comprehensive quality audit before a version release. It e
 ### 3.2 Testing Content and Scope
 
 -   ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
--   ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen3_omni_*.json` (Omni), `test_tts.json` / `test_voxcpm2.json` / `test_higgs_audio_v3.json` (TTS), and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` or `run_diffusion_benchmark.py` via `--test-config-file` in nightly **Perf Test** steps) to drive throughput, latency, and memory benchmarks. Each JSON **case** may declare an optional top-level **`mark`** array: exactly one ``hardware_marks`` object plus pytest marker name strings (`full_model`, `omni` / `tts` / `diffusion`, …). Runners attach those marks to each parametrized `(server, benchmark index)` pair so **local** bulk runs can filter with `-m` (for example `-m "full_model and H100 and diffusion"`). Nightly CI perf jobs select workloads by **`--test-config-file`**, not `-m`. Details are in the Performance Tests example below.
+-   ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs in the form `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` via `--test-config-file`) to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
 -   ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
 
 ### 3.3 Test Directory and Execution Files
 
 -   ***Functional Testing***: Same directories as L3.
--   ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen3_omni_*.json`, `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json`). Optional per-case `mark` for local `-m` filtering is documented in Section 3.4 Performance Tests.
+-   ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`)
 -   ***Documentation Example Tests***:
-    -   `tests/examples/online_serving/test_{model_name}.py`
-    -   `tests/examples/offline_inference/test_{model_name}.py`
+-   -   `tests/example/online_serving/test_{model_name}.py`
+    -   `tests/example/offline_inference/test_{model_name}.py`
 
 ### 3.4 Execution Method and Example
 

docs/contributing/ci/test_examples/l4_performance_tests.inc.md

@@ -1,41 +1,14 @@
-When you want to add L4-level ***performance test*** cases, add entries to JSON files under `tests/dfx/perf/tests/` and run them via `tests/dfx/perf/scripts/run_benchmark.py` (omni / TTS) or `run_diffusion_benchmark.py` (diffusion).
-
-**Config file layout (in-tree examples)**
-
-| Model type | Runner | Example JSON files |
-| ---------- | ------ | ------------------ |
-| Omni | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json`, `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` |
-| TTS | `run_benchmark.py` | `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json` |
-| Diffusion | `run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`, `test_wan22_i2v_vllm_omni.json`, `test_cosmos3_vllm_omni.json`, … |
-
-**How runners pick cases**
-
-Without **`--test-config-file`**, each runner scans all `*.json` under `tests/dfx/perf/tests/` but only keeps its own model type:
-
-- **`run_benchmark.py`**: omni and TTS cases only (skips diffusion JSON).
-- **`run_diffusion_benchmark.py`**: diffusion cases only (skips omni / TTS JSON).
-
-Diffusion cases are detected when the JSON has `server_type` (typically `"vllm-omni"`) or `"diffusion"` in the `mark` array. Omni / TTS JSON has neither.
-
-**Running perf cases**
-
-Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bulk scan above and filter with pytest **`-m`** on each case's JSON `mark` (for example `pytest … run_benchmark.py -m "full_model and tts and H100"`).
+When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, or diffusion configs such as `tests/dfx/perf/tests/test_*_vllm_omni.json` (selected via `pytest ... run_benchmark.py --test-config-file <path>`):
 
 ```JSON
 {
-    "test_name": "test_qwen3_omni_async_chunk",
-    "mark": [
-        {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
-        "full_model",
-        "omni"
-    ],
+    "test_name": "test_qwen3_omni",
     "server_params": {
         "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
         "stage_config_name": "qwen3_omni.yaml"
     },
     "benchmark_params": [
         {
-            "name": "random_10p",
             "dataset_name": "random",
             "num_prompts": [10, 20],
             "max_concurrency": [1, 4],
@@ -60,95 +33,8 @@ Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bul
 | Field            | Required | Description                                                     |
 | ---------------- | -------- | --------------------------------------------------------------- |
 | test_name        | Yes      | Unique identifier for the test case                             |
-| mark             | No       | Pytest marks for this case (see **`mark` field** below). Omit only for configs not meant to be filtered by `-m`. |
 | server_params    | Yes      | Server-side configuration parameters                            |
 | benchmark_params | Yes      | Benchmark running parameters (supports multiple configurations) |
-| server_type      | Diffusion only | When set (e.g. `"vllm-omni"`), routes the case to `run_diffusion_benchmark.py`. |
-| benchmark_endpoint | Diffusion only | API path for the benchmark client (e.g. `/v1/videos`, `/v1/images/generations`). |
-
-**`mark` field**
-
-Optional top-level field on each perf JSON **case object** (one per `test_name`). `run_benchmark.py` and `run_diffusion_benchmark.py` read it via `tests.dfx.conftest.resolve_pytest_marks` and attach the marks to the corresponding `pytest.param`, so you can filter locally with `-m`.
-
-When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same shape as `@hardware_test` / `hardware_marks()` in `tests/helpers/mark.py`), followed by registered pytest marker name strings such as `full_model`, `omni`, `tts`, `diffusion`, or `local_model`.
-
-Supported form:
-
-| Form | Example | Effect |
-| ---- | ------- | ------ |
-| Array (single platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "diffusion"]` | Hardware marks from `hardware_marks(...)`; `num_cards` > 1 adds `distributed_*` (+ CUDA `skipif_cuda` when applicable). String entries become extra pytest marks. |
-| Array (multi-platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100", "rocm": "MI325", "npu": "A2"}, "num_cards": {"cuda": 2, "rocm": 2, "npu": 1}}}, "full_model", "omni"]` | Same as above, but declares **multiple platforms** in one case. Filter examples: `-m "full_model and H100 and cuda"`, `-m "full_model and MI325 and rocm"`. |
-
-Recommended for L4 perf cases:
-
-```JSON
-{
-    "test_name": "test_bagel_single_device_single_stage_t2i",
-    "mark": [
-        {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
-        "full_model",
-        "diffusion"
-    ],
-    "server_type": "vllm-omni",
-    "server_params": { "...": "..." },
-    "benchmark_params": [ { "name": "1024x1024_steps20", "...": "..." } ]
-}
-```
-
-Multi-GPU diffusion (example: Cosmos3 with `cfg-parallel-size=2`):
-
-```JSON
-{
-    "test_name": "test_cosmos3_t2i_official_demo_2gpu",
-    "mark": [
-        {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
-        "full_model",
-        "diffusion"
-    ],
-    "server_type": "vllm-omni",
-    "benchmark_endpoint": "/v1/images/generations",
-    "server_params": { "...": "..." },
-    "benchmark_params": [ { "name": "1024x1024_steps4", "...": "..." } ]
-}
-```
-
-- Omni perf: include `"omni"` in the `mark` array
-- TTS perf: include `"tts"` in the `mark` array
-- Diffusion perf: include `"diffusion"` in the `mark` array
-- HunyuanImage local-weight cases: add `"local_model"` to the `mark` array
-
-**Parametrization IDs**
-
-Each `(server, benchmark index)` pair becomes one `pytest.param` with id `{test_name}-{suffix}`. The suffix comes from `benchmark_params[].name` when set (for example `test_tts-p0`, `test_omni-p1`); otherwise it is derived from fields like `task` / `eval_phase`.
-
-**Benchmark result filenames**
-
-Result files use the **runtime** hardware label from `get_runtime_resource_label()` (detected GPU/NPU name on the machine that ran the job), **not** `mark.hardware_marks.res`. On the default H100 CI pool, `H100` is **omitted** from filenames (`resource_label_for_filename`).
-
-Examples:
-
-- Omni/TTS: `result_{test_name}_{optional_hw}_{dataset}_....json` under `BENCHMARK_DIR`
-- Diffusion: one aggregate `diffusion_result_{config_stem}_{optional_hw}_{timestamp}.json` per source JSON file (array of all runs from that file)
-
-**Local commands**
-
-```bash
-# Bulk load + filter by JSON mark
-pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and H100 and diffusion"
-pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100"
-
-# Single file (same as nightly CI Perf steps)
-pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
-  --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
-pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \
-  --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
-
-# Optional baseline assertion (default off)
-pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --assert-baseline \
-  --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
-```
-
-See also [Markers for Tests](./tests_markers.md) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …).
 
 **`server_params` Configuration**
 
@@ -193,12 +79,10 @@ You can add any benchmark running parameters you need here. For all optional par
 1.  Change the --xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
 2.  For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
 3.  Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field.
-4.  Set `"name"` on each `benchmark_params` entry for stable pytest ids and readable result keys.
-5.  The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below:
+4.  The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below:
 
 | Parameter       | Type        | Required | Example/Values  | Description                                                                                                                                                                                                                                                          |
 | --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| name            | string      | No       | `"1024x1024_steps4"` | Stable suffix for pytest param id and logs. |
 | num_prompts     | int / array | Yes      | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
 | request_rate    | float / array | No  | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts.                          |
 | max_concurrency | int / array | No       | 1, [1, 2, 3]    | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode).                                                                                                                                                                                                             |

docs/contributing/ci/test_guide.md

@@ -149,18 +149,12 @@ Failed jobs: 1/2
     pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model
     pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model
     ```
-    If you only want to run specific test cases on a particular platform, you can use:
+
+    To run a standalone perf benchmark (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS):
+
     ```bash
-    pytest -s -v -m "full_model and distributed_cuda and L4"  --run-level=full_model
+    pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
     ```
-    Note: ``run_benchmark.py`` and ``run_diffusion_benchmark.py`` accept an optional ``--test-config-file``. If omitted, each loads every ``*.json`` under ``tests/dfx/perf/tests/`` (omni/tts vs diffusion split by ``is_diffusion_perf_config``) and pytest ``-m`` filters by each case's JSON ``mark``:
-    ```bash
-    pytest -sv tests/dfx/perf/scripts/run_benchmark.py -m "full_model and tts and H100"
-    pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100"
-    pytest -sv tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json
-    pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json
-    ```
-    Nightly **Perf Test** jobs in ``test-nightly.yml`` use ``--test-config-file`` only (no ``-m``). E2e L4 function tests still use ``full_model`` + ``--run-level full_model`` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml)). Example:
 
 === "L5 level"
 

docs/contributing/ci/tests_markers.md

@@ -33,8 +33,7 @@ Defined in `pyproject.toml`:
 | `skipif_rocm`      | Skip if the num of ROCm cards is less than the required * |
 | `skipif_npu`       | Skip if the num of NPU cards is less than the required *  |
 | `slow`             | Slow tests (may skip in quick CI)                         |
-| `benchmark`        | Benchmark tests (decorator on runner test functions; perf JSON uses `full_model` + type marker instead) |
-| `local_model`      | Tests requiring local / non-HF-hub model weights          |
+| `benchmark`        | Benchmark tests                                           |
 
 \* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
 
@@ -95,7 +94,7 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
         num_cards=2,
     )
     ```
-- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), XPU (B60), MUSA (S5000), NPU (A2/A3)
+- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), NPU (A2/A3)
 - `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing
 - Distributed markers (`distributed_cuda`, `distributed_rocm`, `distributed_npu`) are auto-added for multi-card cases
 - Filtering examples:
@@ -121,32 +120,6 @@ def test_feature(omni_server):
     ...
 ```
 
-### JSON `mark` field (L4 perf configs)
-
-Perf JSON under `tests/dfx/perf/tests/` can attach pytest marks per **case** (`test_name` block). Parsed by `tests.dfx.conftest.resolve_pytest_marks` and applied to each parametrized run in `run_benchmark.py` / `run_diffusion_benchmark.py`.
-
-When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same semantics as `hardware_marks()` above), followed by registered pytest marker name strings such as `full_model`, `omni`, `tts`, or `diffusion`.
-
-```json
-{
-  "test_name": "test_cosmos3_t2i_official_demo_2gpu",
-  "mark": [
-    {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
-    "full_model",
-    "diffusion"
-  ],
-  "server_type": "vllm-omni",
-  "server_params": { "...": "..." },
-  "benchmark_params": [{ "name": "1024x1024_steps4", "...": "..." }]
-}
-```
-
-- Local bulk load: `pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100"`
-- Nightly CI perf steps: `--test-config-file tests/dfx/perf/tests/test_<model>_vllm_omni.json` (file selects cases; no `-m`)
-- Result filenames use **runtime** GPU detection (`get_runtime_resource_label`); `H100` is omitted on the default CI pool
-
-See [L4 performance test examples](./test_examples/l4_performance_tests.inc.md) for the full schema.
-
 ## Add Support for a New Platform
 
 If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:

examples/online_serving/replica_data_parallel/README.md

@@ -1,91 +0,0 @@
-# Replica Data Parallelism for Video Generation
-
-A runnable recipe + benchmark for **replica data parallelism** (replica-DP) on a
-video DiT. Replica-DP runs `N` independent diffusion engine replicas, one per GPU,
-and routes each request to a single replica. It scales **request throughput**
-near-linearly; it does **not** change single-request latency (for that, use an
-intra-request axis such as tensor/sequence/CFG parallelism).
-
-This example uses `Wan-AI/Wan2.2-TI2V-5B-Diffusers` with one GPU per replica.
-
-## Files
-
-| File | Purpose |
-|------|---------|
-| `wan2_2_ti2v_dp.yaml` | Diffusion stage config; `runtime.num_replicas` + `runtime.devices` drive the fan-out. |
-| `run_server.sh` | Substitutes `NUM_REPLICAS` / `DEVICES` into the config and serves. |
-| `bench_replica_dp.py` | Replica-agnostic load driver; reports throughput, latency, and a per-input isolation check. |
-
-## Run
-
-Start the server with the replica count you want (one GPU per replica):
-
-```bash
-# baseline
-NUM_REPLICAS=1 DEVICES=0        ./run_server.sh
-# 2 replicas
-NUM_REPLICAS=2 DEVICES=0,1      ./run_server.sh
-# 4 replicas
-NUM_REPLICAS=4 DEVICES=0,1,2,3  ./run_server.sh
-```
-
-In another shell, drive a fixed workload and read the throughput:
-
-```bash
-python bench_replica_dp.py --url http://127.0.0.1:8098 \
-    --num-requests 28 --concurrency 8 --label "replicas=4"
-```
-
-Re-run the server at `NUM_REPLICAS=1,2,4` (same client workload each time) and
-compare `videos/min`. Pass `--baseline-tpm <N=1 value>` to print a scaling factor.
-
-### Isolation check
-
-Identical outputs under a *shared* prompt+seed prove nothing -- every request
-would be identical regardless of isolation. Instead each request gets a distinct
-input (a prompt from the pool + a unique seed `base_seed+i`), and we check that a
-request's output is unchanged whether it runs alone or concurrently amid other
-distinct requests. Capture a per-input baseline on `N=1`, then verify it under
-`N>=2` concurrency:
-
-```bash
-# 1) against NUM_REPLICAS=1 (serial), record each input's output hash
-python bench_replica_dp.py --num-requests 8 --concurrency 1 \
-    --write-baseline baseline.json
-
-# 2) against NUM_REPLICAS>=2, replay the SAME inputs concurrently and verify
-python bench_replica_dp.py --num-requests 8 --concurrency 8 \
-    --check-baseline baseline.json
-```
-
-Step 2 exits non-zero if any request's output differs from its own baseline (a
-mismatch means cross-request state bleed or misrouting). Keep `--num-requests`,
-`--base-seed`, and the prompt source identical between the two runs so the input
-sets match.
-
-## Measured scaling
-
-Wan2.2-TI2V-5B (832×480, 33 frames, 30 steps), 4× A800-80GB (NVLink), one GPU per replica:
-
-| Replicas | Throughput (videos/min) | Scaling | Efficiency |
-|----------|-------------------------|---------|------------|
-| 1 | 4.71 | 1.00× | — |
-| 2 | 9.20 | 1.95× | 98% |
-| 4 | 18.02 | 3.83× | 96% |
-
-Per-request latency stayed flat (~13–14 s) across all replica counts. Isolation
-is checked per-input via the baseline/verify protocol above: each request's
-output under `N>=2` concurrency must match its own single-replica baseline.
-
-## Notes on the config surface
-
-- `runtime.num_replicas: N` fans out `N` replicas; `runtime.devices` assigns their
-  GPUs. With `tensor_parallel_size = 1`, list one GPU per replica
-  (`num_replicas * tensor_parallel_size` entries, pool mode).
-- Replica fan-out here comes from the stage config, not a CLI flag. (The headless /
-  multi-runtime launch path uses the process-local `--omni-dp-size-local`, which
-  requires `--stage-id`.)
-- This recipe serves via `--stage-configs-path`. Expressing replica-DP under the
-  newer `--deploy-config` schema is part of the ongoing serving-config alignment
-  (see the diffusion parallelism guide and the `#4707` discussion); `num_replicas`
-  is accepted by both paths.

examples/online_serving/replica_data_parallel/bench_replica_dp.py

@@ -1,241 +0,0 @@
-#!/usr/bin/env python3
-"""Throughput-scaling + isolation benchmark for diffusion replica data parallelism.
-
-Sends `--num-requests` text-to-video requests (up to `--concurrency` in flight)
-to a running vLLM-Omni video server via `POST /v1/videos/sync` (blocks until the
-clip is produced) and reports:
-
-    completed / failed, makespan (wall clock), throughput (videos/min),
-    per-request latency p50/p90/max (nearest-rank), and -- when a baseline
-    manifest is supplied -- a per-request isolation check.
-
-Each request gets a *distinct* input (a prompt from the pool + a unique seed),
-so an output is a deterministic function of its own input and nothing else.
-
-Measuring replica scaling
--------------------------
-This client is replica-agnostic. Vary the *server's* replica count
-(`runtime.num_replicas` in the stage config; see run_server.sh) and keep the
-client workload fixed, then compare throughput:
-
-    N=1 -> baseline
-    N=2 -> expect ~2x
-    N=4 -> expect ~4x   (near-linear is the goal)
-
-Pass the single-replica videos/min via `--baseline-tpm` to print a scaling
-factor for the current run.
-
-Isolation / correctness
------------------------
-Byte-identical outputs under a *shared* prompt+seed prove nothing -- every
-request would be identical regardless of isolation. Instead, give each request a
-distinct input and check that its output is unchanged whether it runs alone or
-concurrently amid other distinct requests:
-
-    1. Against an N=1 server, capture a per-input baseline:
-         bench_replica_dp.py --num-requests 8 --concurrency 1 \
-             --write-baseline baseline.json
-    2. Against an N>=2 server, replay the *same* inputs concurrently and verify
-       each output matches its own baseline:
-         bench_replica_dp.py --num-requests 8 --concurrency 8 \
-             --check-baseline baseline.json
-
-A mismatch means a request's output changed under concurrency -- cross-request
-state bleed or misrouting -- and the run exits non-zero.
-
-Requires: requests (`pip install requests`).
-"""
-
-import argparse
-import concurrent.futures as cf
-import hashlib
-import json
-import math
-import sys
-import time
-from pathlib import Path
-
-import requests
-
-# Distinct prompts, cycled across requests; combined with a per-request unique
-# seed this makes every request a distinct input (see input_key()).
-DEFAULT_PROMPTS = [
-    "A cat playing piano, cinematic, high detail",
-    "A drone shot over a snowy mountain range at sunrise",
-    "A red sports car driving through a neon-lit city at night",
-    "Ocean waves crashing on a rocky shore, slow motion",
-    "A hot air balloon floating over green rolling hills",
-    "Timelapse of clouds moving over a desert canyon at dusk",
-    "A steam train crossing a stone viaduct in autumn",
-    "A jellyfish drifting through deep blue water, backlit",
-]
-
-
-def input_key(prompt, seed, args):
-    """Stable identity of a request's input; matches baseline<->verify by content, not order."""
-    return "\x1f".join([prompt, str(seed), args.size, str(args.num_frames), str(args.steps)])
-
-
-def nearest_rank(xs, q):
-    """Nearest-rank percentile; q in [0, 1]. Correct for tiny samples (unlike int(q*n)-1)."""
-    if not xs:
-        return float("nan")
-    s = sorted(xs)
-    rank = max(1, math.ceil(q * len(s)))
-    return s[rank - 1]
-
-
-def one_request(args, spec, idx, save):
-    """Fire one /v1/videos/sync request; return (ok, latency, nbytes, md5, key, err)."""
-    prompt, seed = spec["prompt"], spec["seed"]
-    key = input_key(prompt, seed, args)
-    # multipart form fields match POST /v1/videos/sync in the omni API server
-    form = {
-        "prompt": (None, prompt),
-        "size": (None, args.size),  # e.g. "832x480"
-        "num_frames": (None, str(args.num_frames)),
-        "num_inference_steps": (None, str(args.steps)),
-        "seed": (None, str(seed)),  # per-request unique -> distinct, reproducible input
-    }
-    if args.model:
-        form["model"] = (None, args.model)
-    if args.negative_prompt:
-        form["negative_prompt"] = (None, args.negative_prompt)
-
-    t0 = time.perf_counter()
-    try:
-        r = requests.post(f"{args.url}/v1/videos/sync", files=form, timeout=args.timeout)
-        dt = time.perf_counter() - t0
-        if r.status_code != 200:
-            return (False, dt, 0, "", key, f"HTTP {r.status_code}: {r.text[:160]}")
-        body = r.content
-        md5 = hashlib.md5(body).hexdigest()
-        if save and args.save_dir:
-            Path(args.save_dir).mkdir(parents=True, exist_ok=True)
-            (Path(args.save_dir) / f"req{idx:03d}_seed{seed}.mp4").write_bytes(body)
-        return (True, dt, len(body), md5, key, "")
-    except Exception as e:  # noqa: BLE001 - report any client-side failure
-        return (False, time.perf_counter() - t0, 0, "", key, repr(e))
-
-
-def build_specs(args):
-    """Deterministic list of distinct (prompt, seed) inputs; identical across baseline/verify runs."""
-    prompts = DEFAULT_PROMPTS
-    if args.prompts_file:
-        prompts = [ln.strip() for ln in Path(args.prompts_file).read_text().splitlines() if ln.strip()]
-    return [{"prompt": prompts[i % len(prompts)], "seed": args.base_seed + i} for i in range(args.num_requests)]
-
-
-def run_batch(args, specs, save):
-    """Fire specs concurrently; return (results, makespan). results = list of one_request tuples."""
-    results = []
-    wall0 = time.perf_counter()
-    with cf.ThreadPoolExecutor(max_workers=args.concurrency) as ex:
-        futs = {ex.submit(one_request, args, spec, i, save): i for i, spec in enumerate(specs)}
-        for done, f in enumerate(cf.as_completed(futs), 1):
-            res = f.result()
-            results.append(res)
-            good, dt, nbytes, _md5, _key, err = res
-            print(f"  [{done}/{len(specs)}] {'ok' if good else 'FAIL'} {dt:6.1f}s {(nbytes / 1e6):5.1f}MB {err}")
-    return results, time.perf_counter() - wall0
-
-
-def main():
-    ap = argparse.ArgumentParser(description="Diffusion replica-DP throughput + isolation benchmark")
-    ap.add_argument("--url", default="http://127.0.0.1:8098", help="server base URL (no path)")
-    ap.add_argument("--num-requests", type=int, default=16, help="total requests to send")
-    ap.add_argument("--concurrency", type=int, default=4, help="max in-flight (set >= replica count)")
-    ap.add_argument("--warmup", type=int, default=-1, help="warmup requests before timing (default: =concurrency)")
-    ap.add_argument("--base-seed", type=int, default=42, help="request i uses seed base_seed+i (distinct inputs)")
-    ap.add_argument("--prompts-file", default="", help="optional file, one prompt per line (else a built-in pool)")
-    ap.add_argument("--negative-prompt", default="")
-    ap.add_argument("--size", default="832x480", help="WIDTHxHEIGHT")
-    ap.add_argument("--num-frames", type=int, default=33)
-    ap.add_argument("--steps", type=int, default=30, help="num_inference_steps")
-    ap.add_argument("--model", default="", help="optional; defaults to the served model")
-    ap.add_argument("--timeout", type=float, default=1800, help="per-request timeout (s)")
-    ap.add_argument("--save-dir", default="", help="save returned clips")
-    ap.add_argument("--label", default="", help="run label, e.g. 'replicas=2'")
-    ap.add_argument(
-        "--baseline-tpm", type=float, default=0.0, help="single-replica videos/min, to print a scaling factor"
-    )
-    ap.add_argument("--write-baseline", default="", help="run against N=1: write per-input md5 manifest to this path")
-    ap.add_argument("--check-baseline", default="", help="run against N>=2: verify each output matches this manifest")
-    args = ap.parse_args()
-
-    if args.write_baseline and args.check_baseline:
-        ap.error("--write-baseline and --check-baseline are mutually exclusive")
-
-    specs = build_specs(args)
-    warmup = args.concurrency if args.warmup < 0 else args.warmup
-
-    print(
-        f"-> {args.url}/v1/videos/sync | {args.num_requests} reqs "
-        f"| concurrency {args.concurrency} | warmup {warmup} | {args.size} {args.num_frames}f {args.steps}steps "
-        f"| {args.label}"
-    )
-
-    # Warmup: absorb each replica's first-call compile/init cost so it does not
-    # skew the measured makespan. Results are discarded and never saved.
-    if warmup > 0:
-        print(f"warming up ({warmup} reqs, not timed)...")
-        run_batch(args, [specs[i % len(specs)] for i in range(warmup)], save=False)
-
-    results, makespan = run_batch(args, specs, save=True)
-
-    ok = sum(1 for r in results if r[0])
-    fail = len(results) - ok
-    lat = [r[1] for r in results if r[0]]
-    errs = [r[5] for r in results if not r[0]]
-    got = {r[4]: r[3] for r in results if r[0]}  # input_key -> md5
-
-    print(f"\n===== result {args.label} =====")
-    print(f"completed {ok} / failed {fail} | makespan {makespan:.1f}s")
-    if ok:
-        tpm = ok / makespan * 60.0
-        print(f"throughput: {tpm:.2f} videos/min  ({ok / makespan:.4f} videos/s)")
-        print(
-            f"latency (s): p50={nearest_rank(lat, 0.50):.1f} "
-            f"p90={nearest_rank(lat, 0.90):.1f} "
-            f"min={min(lat):.1f} max={max(lat):.1f}"
-        )
-        if args.baseline_tpm > 0:
-            print(f"scaling vs baseline ({args.baseline_tpm:.2f}/min): {tpm / args.baseline_tpm:.2f}x")
-    if errs:
-        print("sample failures:", errs[:3])
-
-    # Isolation check: compare each request's output against its own single-replica baseline.
-    isolation_failed = False
-    if args.write_baseline:
-        Path(args.write_baseline).write_text(json.dumps(got, indent=2))
-        print(f"baseline: wrote {len(got)} per-input hashes -> {args.write_baseline}")
-    elif args.check_baseline:
-        baseline = json.loads(Path(args.check_baseline).read_text())
-        matched = mismatched = unknown = 0
-        for key, md5 in got.items():
-            if key not in baseline:
-                unknown += 1
-                print(f"  isolation: NO BASELINE for input {key!r} (not in the baseline manifest)")
-            elif baseline[key] == md5:
-                matched += 1
-            else:
-                mismatched += 1
-                print(f"  isolation: MISMATCH for input {key!r} (output changed under concurrency)")
-        # PASS requires the run to cover the *whole* baseline: a subset (fewer requests, or
-        # some failed) must not pass just because the inputs it did run happened to match.
-        unverified = sorted(set(baseline) - set(got))
-        for key in unverified:
-            print(f"  isolation: NOT VERIFIED for baseline input {key!r} (absent from this run)")
-        isolation_failed = mismatched > 0 or unknown > 0 or len(unverified) > 0
-        verdict = "PASS (isolated)" if not isolation_failed else "FAIL"
-        print(
-            f"isolation: {matched} matched / {mismatched} mismatched / {unknown} unknown / "
-            f"{len(unverified)} unverified vs baseline({len(baseline)}) -> {verdict}"
-        )
-
-    if fail > 0 or ok == 0 or isolation_failed:
-        sys.exit(1)
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/replica_data_parallel/run_server.sh

@@ -1,33 +0,0 @@
-#!/usr/bin/env bash
-# Serve N independent Wan2.2 TI2V-5B replicas (replica data parallelism).
-#
-# Usage:
-#   NUM_REPLICAS=1 DEVICES=0        ./run_server.sh      # baseline
-#   NUM_REPLICAS=2 DEVICES=0,1      ./run_server.sh
-#   NUM_REPLICAS=4 DEVICES=0,1,2,3  ./run_server.sh
-#
-# Then drive load with bench_replica_dp.py (see README.md) and compare
-# throughput across replica counts.
-set -euo pipefail
-
-MODEL="${MODEL:-Wan-AI/Wan2.2-TI2V-5B-Diffusers}"
-HOST="${HOST:-127.0.0.1}"
-PORT="${PORT:-8098}"
-NUM_REPLICAS="${NUM_REPLICAS:-1}"
-DEVICES="${DEVICES:-0}"   # one GPU per replica for tensor_parallel_size=1
-
-HERE="$(cd "$(dirname "$0")" && pwd)"
-TEMPLATE="$HERE/wan2_2_ti2v_dp.yaml"
-CFG="$(mktemp --suffix=_wan2_2_dp.yaml)"
-trap 'rm -f "$CFG"' EXIT
-
-sed -e "s/__NUM_REPLICAS__/${NUM_REPLICAS}/" \
-    -e "s/__DEVICES__/${DEVICES}/" \
-    "$TEMPLATE" > "$CFG"
-
-echo "Serving ${MODEL}: num_replicas=${NUM_REPLICAS} devices=${DEVICES} (config ${CFG})"
-
-vllm serve "$MODEL" --omni \
-    --host "$HOST" \
-    --port "$PORT" \
-    --stage-configs-path "$CFG"

examples/online_serving/replica_data_parallel/wan2_2_ti2v_dp.yaml

@@ -1,43 +0,0 @@
-# Replica data-parallel recipe for Wan2.2 TI2V-5B (video DiT).
-#
-# Fans out `num_replicas` independent DiffusionEngine replicas, one per GPU,
-# and lets the serving layer route each request to a single replica. This
-# scales request throughput near-linearly; it does not change single-request
-# latency. See ./README.md for measured numbers and the `run_server.sh` driver.
-#
-# run_server.sh substitutes __NUM_REPLICAS__ / __DEVICES__ from env before serving.
-# For tensor_parallel_size = 1, `devices` lists one GPU per replica (pool mode):
-# num_replicas=4 -> devices "0,1,2,3".
-
-stage_args:
-  - stage_id: 0
-    stage_type: diffusion
-    runtime:
-      num_replicas: __NUM_REPLICAS__
-      devices: "__DEVICES__"
-      max_batch_size: 1
-    engine_args:
-      model_stage: dit
-      # `model_class_name` takes the diffusion registry KEY (WanPipeline), which
-      # resolves to the Wan22Pipeline class. Wan2.2 TI2V-5B is served by the base
-      # WanPipeline, which auto-detects TI2V-5B mode from model_index.json.
-      model_class_name: WanPipeline
-      max_num_seqs: 1
-      enforce_eager: true
-      trust_remote_code: true
-      distributed_executor_backend: "mp"
-      vae_use_tiling: true
-      parallel_config:
-        tensor_parallel_size: 1
-
-    final_output: true
-    final_output_type: video
-    is_comprehension: false
-    default_sampling_params:
-      seed: 42
-
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1

recipes/LTX/LTX-2.3.md

@@ -148,7 +148,7 @@ Do not copy latency or VRAM numbers between commits or machines. Record them
 only after running a latest-head sweep on the target hardware.
 
 For formal PR benchmarking, reuse
-`tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
-`tests/dfx/perf/scripts/run_diffusion_benchmark.py`. Each case in that file includes a JSON **`mark`** (`hardware_marks` + `full_model` + `diffusion`) for local `-m` filtering. That config captures the
+`tests/dfx/perf/tests/test_ltx2_3_vllm_omni.json` with
+`tests/dfx/perf/scripts/run_diffusion_benchmark.py`. That config captures the
 single-device eager baseline and CFG-parallel=2 case for a small 384x512,
 25-frame, 20-step workload.

tests/dfx/conftest.py

@@ -11,180 +11,11 @@ from typing import Any
 import pytest
 
 from tests.dfx.reliability.helpers import list_remote_process_pids_by_pattern, post_chat_completions_raw
-from tests.helpers.mark import hardware_marks
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import modify_stage_config
 from vllm_omni.platforms import current_omni_platform
 
 
-def _named_pytest_marks(names: list[str]) -> list[pytest.MarkDecorator]:
-    marks: list[pytest.MarkDecorator] = []
-    for name in names:
-        name = name.strip()
-        if not name:
-            raise ValueError("mark name must be a non-empty string")
-        marks.append(getattr(pytest.mark, name))
-    return marks
-
-
-def _hardware_marks_from_dict(hw: Any) -> list[pytest.MarkDecorator]:
-    if not isinstance(hw, dict):
-        raise ValueError(f"mark.hardware_marks must be a dict, got {type(hw).__name__}")
-    res = hw.get("res")
-    if not isinstance(res, dict):
-        raise ValueError(f"mark.hardware_marks.res must be a dict, got {type(res).__name__}")
-    num_cards = hw.get("num_cards", 1)
-    return list(hardware_marks(res=res, num_cards=num_cards))
-
-
-def resolve_pytest_marks(mark_field: Any) -> list[pytest.MarkDecorator]:
-    """Convert a JSON ``mark`` field into pytest mark decorators.
-
-    Supported form (per test-case object in perf/stability JSON)::
-
-        "mark": [
-            {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
-            "full_model",
-            "diffusion"
-        ]
-
-    Exactly one ``hardware_marks`` object is required when ``mark`` is present.
-    ``hardware_marks`` delegates to :func:`tests.helpers.mark.hardware_marks`
-    (same shape as ``@hardware_test``). Additional array entries are registered
-    pytest marker names (strings).
-    """
-    if mark_field is None:
-        return []
-
-    if isinstance(mark_field, list):
-        marks: list[pytest.MarkDecorator] = []
-        hw_seen = False
-        for item in mark_field:
-            if isinstance(item, dict) and "hardware_marks" in item:
-                if hw_seen:
-                    raise ValueError("mark array must contain at most one hardware_marks object")
-                hw_seen = True
-                marks.extend(_hardware_marks_from_dict(item["hardware_marks"]))
-                unknown_keys = set(item) - {"hardware_marks"}
-                if unknown_keys:
-                    raise ValueError(
-                        f"mark hardware_marks object only allows hardware_marks; unknown keys: {sorted(unknown_keys)}"
-                    )
-            elif isinstance(item, str):
-                item = item.strip()
-                if not item:
-                    raise ValueError("mark name must be a non-empty string")
-                marks.extend(_named_pytest_marks([item]))
-            else:
-                raise ValueError(
-                    f"mark array entries must be hardware_marks objects or marker name strings; got {type(item).__name__}"
-                )
-        if not hw_seen:
-            raise ValueError("mark array must contain a hardware_marks object")
-        return marks
-
-    raise ValueError(f"mark must be a list; got {type(mark_field).__name__}")
-
-
-def _mark_names(mark_field: Any) -> set[str]:
-    if isinstance(mark_field, list):
-        return {str(item) for item in mark_field if isinstance(item, str)}
-    return set()
-
-
-def is_diffusion_perf_config(cfg: dict[str, Any]) -> bool:
-    """True for perf JSON cases intended for ``run_diffusion_benchmark.py``."""
-    if cfg.get("server_type") is not None:
-        return True
-    return "diffusion" in _mark_names(cfg.get("mark"))
-
-
-def _marks_by_test_name(configs: list[dict[str, Any]]) -> dict[str, list[pytest.MarkDecorator]]:
-    return {str(cfg["test_name"]): resolve_pytest_marks(cfg.get("mark")) for cfg in configs}
-
-
-def create_unique_server_pytest_params(
-    configs: list[dict[str, Any]],
-    stage_configs_dir: Path,
-) -> list[Any]:
-    """Like :func:`create_unique_server_params`, but wrap each row in ``pytest.param`` with JSON marks."""
-    marks_by_name = _marks_by_test_name(configs)
-    return [
-        pytest.param(
-            row,
-            marks=marks_by_name.get(row[0], []),
-            id=row[0],
-        )
-        for row in create_unique_server_params(configs, stage_configs_dir)
-    ]
-
-
-def create_benchmark_pytest_params(
-    benchmark_configs: list[dict[str, Any]],
-    server_to_benchmark_mapping: dict[str, dict],
-) -> list[Any]:
-    """Like :func:`create_benchmark_indices`, but wrap each index in ``pytest.param`` with JSON marks."""
-    marks_by_name = _marks_by_test_name(benchmark_configs)
-    params: list[Any] = []
-    seen: set[str] = set()
-    for config in benchmark_configs:
-        test_name = config["test_name"]
-        if test_name in seen:
-            continue
-        seen.add(test_name)
-        params_list = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping)
-        id_suffixes = _unique_benchmark_param_id_suffixes(params_list)
-        for idx, id_suffix in enumerate(id_suffixes):
-            params.append(
-                pytest.param(
-                    (test_name, idx),
-                    marks=marks_by_name.get(test_name, []),
-                    id=f"{test_name}-{id_suffix}",
-                )
-            )
-    return params
-
-
-def create_paired_benchmark_pytest_params(
-    server_entries: list[tuple[Any, str]],
-    params_by_test_name: dict[str, list[Any]],
-    marks_by_name: dict[str, list[pytest.MarkDecorator]],
-) -> list[Any]:
-    """One ``pytest.param`` per ``(server entry, benchmark index)`` pair.
-
-    Pass ``server_entry`` and ``(test_name, idx)`` as separate ``pytest.param``
-    arguments for ``@pytest.mark.parametrize("server_fixture,benchmark_params", ...)``.
-    """
-    pairs: list[Any] = []
-    for server_entry, test_name in server_entries:
-        params_list = params_by_test_name.get(test_name, [])
-        id_suffixes = _unique_benchmark_param_id_suffixes(params_list)
-        for idx, id_suffix in enumerate(id_suffixes):
-            pairs.append(
-                pytest.param(
-                    server_entry,
-                    (test_name, idx),
-                    marks=marks_by_name.get(test_name, []),
-                    id=f"{test_name}-{id_suffix}",
-                )
-            )
-    return pairs
-
-
-def create_paired_omni_benchmark_pytest_params(
-    configs: list[dict[str, Any]],
-    stage_configs_dir: Path,
-) -> list[Any]:
-    """Paired params for ``run_benchmark.py`` (omni/tts)."""
-    mapping = create_test_parameter_mapping(configs)
-    marks_by_name = _marks_by_test_name(configs)
-    server_entries = [(row, row[0]) for row in create_unique_server_params(configs, stage_configs_dir)]
-    params_by_test_name = {
-        test_name: get_benchmark_params_for_server(test_name, mapping) for _, test_name in server_entries
-    }
-    return create_paired_benchmark_pytest_params(server_entries, params_by_test_name, marks_by_name)
-
-
 def load_configs(config_path: str) -> list[dict[str, Any]]:
     try:
         abs_path = Path(config_path).resolve()
@@ -201,29 +32,6 @@ def load_configs(config_path: str) -> list[dict[str, Any]]:
         raise RuntimeError(f"Failed to load configuration file: {str(e)}")
 
 
-def load_benchmark_configs(
-    config_path: str | None = None,
-    *,
-    config_dir: Path | None = None,
-) -> list[dict[str, Any]]:
-    """Load one benchmark JSON file, or merge all ``*.json`` under *config_dir*.
-
-    When *config_path* is omitted, every ``*.json`` in *config_dir* is loaded and
-    concatenated. Pytest ``-m`` expressions then select cases via each entry's
-    ``mark`` field (e.g. ``-m tts`` after bulk load).
-    """
-    if config_path is not None:
-        return load_configs(config_path)
-    if config_dir is None:
-        raise ValueError("load_benchmark_configs requires config_path or config_dir")
-    configs: list[dict[str, Any]] = []
-    for path in sorted(config_dir.glob("*.json")):
-        configs.extend(load_configs(str(path)))
-    if not configs:
-        raise ValueError(f"No benchmark JSON files found under {config_dir}")
-    return configs
-
-
 def modify_stage(default_path: str, updates: dict[str, Any] | None, deletes: dict[str, Any] | None) -> str:
     kwargs: dict[str, Any] = {}
     if updates is not None:
@@ -485,158 +293,6 @@ def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
     return s if s else default
 
 
-def _benchmark_param_id_suffix(param: dict[str, Any], *, idx: int) -> str:
-    """Derive a readable pytest id suffix from one ``benchmark_params`` entry."""
-    name = param.get("name")
-    if isinstance(name, str) and name.strip():
-        return _safe_filename_token(name.strip())
-
-    parts: list[str] = []
-    for key in ("task", "eval_phase", "dataset_name", "dataset"):
-        value = param.get(key)
-        if value is None or value == "":
-            continue
-        token = _safe_filename_token(str(value))
-        if token != "na":
-            parts.append(token)
-    if parts:
-        return "_".join(parts)
-
-    return f"case{idx}"
-
-
-def _unique_benchmark_param_id_suffixes(params_list: list[dict[str, Any]]) -> list[str]:
-    """Return unique pytest id suffixes for a server's benchmark param list."""
-    raw = [_benchmark_param_id_suffix(param, idx=idx) for idx, param in enumerate(params_list)]
-    seen: dict[str, int] = {}
-    unique: list[str] = []
-    for suffix in raw:
-        count = seen.get(suffix, 0)
-        seen[suffix] = count + 1
-        if count == 0:
-            unique.append(suffix)
-        else:
-            unique.append(f"{suffix}_{count}")
-    return unique
-
-
-def extract_mark_resource_label(mark_field: Any) -> str:
-    """Return a filename-safe hardware label from ``mark.hardware_marks.res`` values.
-
-    Example: ``{"cuda": "H100"}`` -> ``"H100"``; multiple platforms join with ``-``.
-
-    Prefer :func:`get_runtime_resource_label` for perf result filenames so labels
-    reflect the machine that actually ran the benchmark.
-    """
-    if isinstance(mark_field, list):
-        for item in mark_field:
-            if isinstance(item, dict) and "hardware_marks" in item:
-                return extract_mark_resource_label(item)
-        return "na"
-    if not isinstance(mark_field, dict):
-        return "na"
-    hw = mark_field.get("hardware_marks")
-    if not isinstance(hw, dict):
-        return "na"
-    res = hw.get("res")
-    if not isinstance(res, dict) or not res:
-        return "na"
-    labels = [_safe_filename_token(value) for value in res.values()]
-    return "-".join(labels) if labels else "na"
-
-
-_KNOWN_RUNTIME_RESOURCE_TOKENS: tuple[str, ...] = (
-    "H100",
-    "H800",
-    "H200",
-    "H20",
-    "L40S",
-    "L40",
-    "L4",
-    "A100",
-    "A800",
-    "A10G",
-    "A10",
-    "A30",
-    "MI325",
-    "MI300",
-    "MI250",
-    "B60",
-    "S5000",
-    "910B4",
-    "910B",
-    "910",
-    "310P",
-    "A2",
-    "A3",
-)
-_RUNTIME_RESOURCE_LABEL: str | None = None
-
-
-def _normalize_runtime_device_label(raw: str) -> str:
-    """Map a platform device name to a short filename-safe resource token."""
-    if not raw or not str(raw).strip():
-        return "na"
-    upper = str(raw).upper()
-    for token in _KNOWN_RUNTIME_RESOURCE_TOKENS:
-        if token.upper() in upper:
-            return _safe_filename_token(token)
-    compact = re.sub(r"[^a-zA-Z0-9]+", "", str(raw))
-    for prefix in ("NVIDIA", "AMD", "ASCEND", "HUAWEI"):
-        if compact.upper().startswith(prefix):
-            compact = compact[len(prefix) :]
-            break
-    return _safe_filename_token(compact[:48]) if compact else "na"
-
-
-def _read_runtime_device_name(*, device_id: int = 0) -> str | None:
-    """Device name from the active Omni platform."""
-    if current_omni_platform.device_count() <= device_id:
-        return None
-    get_name = getattr(current_omni_platform, "get_device_name", None)
-    if not callable(get_name):
-        return None
-    raw = get_name(device_id)
-    if isinstance(raw, str) and raw.strip():
-        return raw.strip()
-    return None
-
-
-def get_runtime_resource_label(*, device_id: int = 0, refresh: bool = False) -> str:
-    """Return a filename-safe hardware label detected on the running machine."""
-    global _RUNTIME_RESOURCE_LABEL
-    if not refresh and _RUNTIME_RESOURCE_LABEL is not None:
-        return _RUNTIME_RESOURCE_LABEL
-    raw = _read_runtime_device_name(device_id=device_id)
-    label = _normalize_runtime_device_label(raw) if raw else "na"
-    if not refresh:
-        _RUNTIME_RESOURCE_LABEL = label
-    return label
-
-
-_FILENAME_OMIT_RESOURCE_LABELS = frozenset({"H100"})
-
-
-def hardware_json_value(resource_label: str | None) -> str:
-    """Hardware token stored in perf result JSON (empty when unknown)."""
-    token = _safe_filename_token(resource_label)
-    return "" if token == "na" else token
-
-
-def resource_label_for_filename(resource_label: str | None) -> str:
-    """Hardware token embedded in result filenames (H100 omitted on default CI pool)."""
-    token = _safe_filename_token(resource_label)
-    if token in _FILENAME_OMIT_RESOURCE_LABELS:
-        return ""
-    return token
-
-
-def extract_configs_resource_label(configs: list[dict[str, Any]]) -> str:
-    """Return runtime hardware label for perf result filenames."""
-    del configs
-    return get_runtime_resource_label()
-
-
 def resolve_baseline_value(
     baseline_raw: Any,
     *,
@@ -701,7 +357,6 @@ def run_benchmark(
     max_concurrency: Any | None = None,
     random_input_len: Any | None = None,
     random_output_len: Any | None = None,
-    resource_label: str | None = None,
 ) -> dict[str, Any]:
     """Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
 
@@ -712,11 +367,7 @@ def run_benchmark(
     current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
     ri = _safe_filename_token(random_input_len)
     ro = _safe_filename_token(random_output_len)
-    hw = resource_label_for_filename(resource_label)
-    if hw:
-        result_filename = f"result_{test_name}_{hw}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json"
-    else:
-        result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json"
+    result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json"
     if "--result-filename" in args:
         print(f"The result file will be overwritten by {result_filename}")
     command = (
@@ -786,7 +437,6 @@ def run_benchmark(
         result["random_input_len"] = random_input_len
     if random_output_len is not None:
         result["random_output_len"] = random_output_len
-    result["Hardware"] = hardware_json_value(resource_label)
     with open(result_path, "w", encoding="utf-8") as f:
         json.dump(result, f, ensure_ascii=False, indent=2)
     return result

tests/dfx/perf/scripts/run_benchmark.py

@@ -7,24 +7,21 @@ from typing import Any
 import pytest
 
 from tests.dfx.conftest import (
-    create_paired_omni_benchmark_pytest_params,
+    create_benchmark_indices,
     create_test_parameter_mapping,
+    create_unique_server_params,
     get_benchmark_params_for_server,
-    get_runtime_resource_label,
-    is_diffusion_perf_config,
-    load_benchmark_configs,
+    load_configs,
     resolve_baseline_value,
     run_benchmark,
 )
 from tests.helpers.runtime import OmniServer
 
+pytestmark = [pytest.mark.full_model]
+
 # Compare metrics to each test JSON ``baseline`` block only when pytest is run with ``--assert-baseline``
 # (registered in ``tests/dfx/conftest.py``; default: off). ``run_benchmark`` and ``_resolve_baseline_value`` are
 # defined in the same module.
-#
-# Optional JSON field ``mark`` is applied as pytest marks via
-# ``create_paired_omni_benchmark_pytest_params`` (e.g. ``"mark": [{"hardware_marks":
-# {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "omni"]``).
 
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
@@ -43,22 +40,22 @@ def _get_config_file_from_argv() -> str | None:
 
 
 _PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests"
+_DEFAULT_CONFIG_FILE = str(_PERF_TESTS_DIR / "test_qwen_omni.json")
 
 CONFIG_FILE_PATH = _get_config_file_from_argv()
 if CONFIG_FILE_PATH is None:
-    _all_configs = load_benchmark_configs(config_dir=_PERF_TESTS_DIR)
-    BENCHMARK_CONFIGS = [cfg for cfg in _all_configs if not is_diffusion_perf_config(cfg)]
     print(
-        f"No --test-config-file: loaded {len(BENCHMARK_CONFIGS)} omni/tts case(s) from "
-        f"{_PERF_TESTS_DIR}/*.json (skipped {len(_all_configs) - len(BENCHMARK_CONFIGS)} diffusion; "
-        f"use -m to filter, e.g. -m tts)"
+        "No --test-config-file in argv, using default: tests/dfx/perf/tests/test_qwen_omni.json "
+        "(override with e.g. --test-config-file tests/dfx/perf/tests/test_tts.json)"
     )
-else:
-    BENCHMARK_CONFIGS = load_benchmark_configs(CONFIG_FILE_PATH)
+    CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE
+
+BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
+
 
 DEPLOY_CONFIGS_DIR = Path(__file__).parent.parent / "deploy"
+test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR)
 server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS)
-paired_benchmark_params = create_paired_omni_benchmark_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR)
 
 _omni_server_lock = threading.Lock()
 
@@ -94,11 +91,17 @@ def omni_server(request):
         print("OmniServer stopped")
 
 
+benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
+
+
 @pytest.fixture
-def benchmark_params(request):
-    """Benchmark parameters fixture; paired with ``omni_server`` via parametrization."""
+def benchmark_params(request, omni_server):
+    """Benchmark parameters fixture with proper parametrization"""
     test_name, param_index = request.param
 
+    if test_name != omni_server.test_name:
+        pytest.skip(f"Skipping parameter for {test_name} - current server is {omni_server.test_name}")
+
     all_params = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping)
 
     if not all_params:
@@ -150,11 +153,8 @@ def assert_result(
 
 
 @pytest.mark.benchmark
-@pytest.mark.parametrize(
-    "omni_server,benchmark_params",
-    paired_benchmark_params,
-    indirect=["omni_server", "benchmark_params"],
-)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
 def test_performance_benchmark(omni_server, benchmark_params, request):
     test_name = benchmark_params["test_name"]
     params = benchmark_params["params"]
@@ -168,7 +168,6 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
     print(f"Benchmark parameters: {benchmark_params}")
 
     assert_baseline = request.config.getoption("--assert-baseline", default=False)
-    resource_label = get_runtime_resource_label()
 
     def to_list(value, default=None):
         if value is None:
@@ -240,7 +239,6 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
             max_concurrency=None,
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
-            resource_label=resource_label,
         )
         assert_result(
             result,
@@ -266,7 +264,6 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
             max_concurrency=concurrency,
             random_input_len=params.get("random_input_len"),
             random_output_len=params.get("random_output_len"),
-            resource_label=resource_label,
         )
         assert_result(
             result,

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -8,21 +8,15 @@ This runner separates two concepts:
 2. ``benchmark_endpoint``: which serving API the benchmark client calls.
    Examples: ``/v1/chat/completions`` and ``/v1/videos``.
 
-A config JSON file may be passed via --test-config-file. If omitted, every ``*.json`` under
-``tests/dfx/perf/tests/`` is loaded and pytest ``-m`` filters by each case's ``mark``:
-  pytest run_diffusion_benchmark.py -m "diffusion"
+A config JSON file is REQUIRED via --test-config-file:
   pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
 
 Optional: ``--assert-baseline`` compares metrics to the ``baseline`` block in each benchmark entry (default: off).
 
-Optional JSON field ``mark`` is applied as pytest marks on that case via
-``pytest.param`` (e.g. ``"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, "full_model", "diffusion"]``).
-
-All benchmark results are written under BENCHMARK_RESULT_DIR (override via the
-DIFFUSION_BENCHMARK_DIR environment variable). Each source JSON file gets one
-aggregated ``diffusion_result_{config_stem}_{hardware}_{timestamp}.json`` (JSON array
-of all runs from cases in that file). Bulk load without ``--test-config-file`` uses
-the same per-file aggregation; ``-m`` only selects which cases run.
+All benchmark results for a session are consolidated into a single JSON file under
+BENCHMARK_RESULT_DIR (override via the DIFFUSION_BENCHMARK_DIR environment variable).
+Each entry in the file contains the test metadata (test_name, endpoint, benchmark_params,
+timestamp) together with the raw metrics returned by the benchmark script.
 """
 
 import json
@@ -41,15 +35,8 @@ import psutil
 import pytest
 
 from benchmarks.diffusion.backends import endpoint_filename_token, normalize_endpoint
-from tests.dfx.conftest import (
-    create_paired_benchmark_pytest_params,
-    get_runtime_resource_label,
-    hardware_json_value,
-    is_diffusion_perf_config,
-    resolve_pytest_marks,
-    resource_label_for_filename,
-)
-from tests.helpers.runtime import get_open_port
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model, pytest.mark.local_model]
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
@@ -139,10 +126,6 @@ _BRANCHPOINT_COMMIT_SHA: str | None = None
 DIFFUSION_RESULT_TEMPLATE_PATH = Path(__file__).parent / "diffusion_result_template.json"
 
 
-_DIFFUSION_SOURCE_CONFIG_KEY = "_source_config_file"
-_PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests"
-
-
 def _get_config_file_from_argv() -> str | None:
     """Read --test-config-file from sys.argv at import time so pytest parametrize can use it.
 
@@ -159,6 +142,9 @@ def _get_config_file_from_argv() -> str | None:
 
 
 CONFIG_FILE_PATH = _get_config_file_from_argv()
+if CONFIG_FILE_PATH is None:
+    print("No config file provided, using default config file: tests/dfx/perf/tests/test_qwen_image_vllm_omni.json")
+    CONFIG_FILE_PATH = "tests/dfx/perf/tests/test_qwen_image_vllm_omni.json"
 
 # ---------------------------------------------------------------------------
 # Config loading
@@ -198,85 +184,28 @@ def load_configs(config_path: str) -> list[dict[str, Any]]:
         raise RuntimeError(f"Failed to load configuration file: {str(e)}")
 
 
-def load_diffusion_benchmark_configs(
-    config_path: str | None = None,
-    *,
-    config_dir: Path | None = None,
-) -> list[dict[str, Any]]:
-    """Load one diffusion benchmark JSON, or merge all ``*.json`` under *config_dir*."""
-    if config_path is not None:
-        configs = load_configs(config_path)
-        source = str(Path(config_path).resolve())
-        for cfg in configs:
-            cfg.setdefault(_DIFFUSION_SOURCE_CONFIG_KEY, source)
-        return configs
-    if config_dir is None:
-        raise ValueError("load_diffusion_benchmark_configs requires config_path or config_dir")
-    configs: list[dict[str, Any]] = []
-    for path in sorted(config_dir.glob("*.json")):
-        source = str(path.resolve())
-        for cfg in load_configs(str(path)):
-            cfg[_DIFFUSION_SOURCE_CONFIG_KEY] = source
-            configs.append(cfg)
-    if not configs:
-        raise ValueError(f"No benchmark JSON files found under {config_dir}")
-    return configs
+BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
+
+_config_stem = Path(CONFIG_FILE_PATH).stem  # e.g. "test_qwen_image_vllm_omni"
+AGGREGATED_RESULT_FILE = BENCHMARK_RESULT_DIR / f"diffusion_result_{_config_stem}_{_SESSION_TIMESTAMP}.json"
 
 
-if CONFIG_FILE_PATH is None:
-    _all_configs = load_diffusion_benchmark_configs(config_dir=_PERF_TESTS_DIR)
-    BENCHMARK_CONFIGS = [cfg for cfg in _all_configs if is_diffusion_perf_config(cfg)]
-    print(
-        f"No --test-config-file: loaded {len(BENCHMARK_CONFIGS)} diffusion case(s) from "
-        f"{_PERF_TESTS_DIR}/*.json (skipped {len(_all_configs) - len(BENCHMARK_CONFIGS)} omni/tts; "
-        f"use -m to filter, e.g. -m diffusion)"
-    )
-else:
-    BENCHMARK_CONFIGS = load_diffusion_benchmark_configs(CONFIG_FILE_PATH)
+def _append_to_aggregated_file(record: dict[str, Any]) -> None:
+    """Thread-safe append of *record* to the session-level aggregated JSON file.
 
-_AGGREGATED_RESULT_FILES_BY_SOURCE: dict[str, Path] = {}
-
-
-def _normalized_source_path(source_file: str) -> str:
-    return str(Path(source_file).resolve())
-
-
-def _aggregated_result_file_for_source(source_file: str) -> Path:
-    """One session aggregate per source JSON (same naming as single ``--test-config-file``)."""
-    key = _normalized_source_path(source_file)
-    if key not in _AGGREGATED_RESULT_FILES_BY_SOURCE:
-        stem = Path(key).stem
-        resource = resource_label_for_filename(get_runtime_resource_label())
-        if resource:
-            result_name = f"diffusion_result_{stem}_{resource}_{_SESSION_TIMESTAMP}.json"
-        else:
-            result_name = f"diffusion_result_{stem}_{_SESSION_TIMESTAMP}.json"
-        _AGGREGATED_RESULT_FILES_BY_SOURCE[key] = BENCHMARK_RESULT_DIR / result_name
-    return _AGGREGATED_RESULT_FILES_BY_SOURCE[key]
-
-
-def _write_result_record(record: dict[str, Any]) -> Path:
-    """Append one benchmark record to the aggregate file for its source JSON."""
-    source_file = record.get("source_file") or CONFIG_FILE_PATH
-    if not source_file:
-        raise ValueError("benchmark record missing source_file")
-    target = _aggregated_result_file_for_source(str(source_file))
+    The file contains a JSON array; each call loads the existing array (or
+    starts a new one), appends the record, and writes the file back atomically.
+    """
     with _RESULT_LOCK:
         BENCHMARK_RESULT_DIR.mkdir(parents=True, exist_ok=True)
-        if target.exists():
-            with open(target, encoding="utf-8") as f:
+        if AGGREGATED_RESULT_FILE.exists():
+            with open(AGGREGATED_RESULT_FILE, encoding="utf-8") as f:
                 records: list[dict] = json.load(f)
         else:
             records = []
         records.append(record)
-        with open(target, "w", encoding="utf-8") as f:
+        with open(AGGREGATED_RESULT_FILE, "w", encoding="utf-8") as f:
             json.dump(records, f, indent=2, ensure_ascii=False)
-    return target
-
-
-def _append_to_aggregated_file(record: dict[str, Any]) -> None:
-    """Backward-compatible wrapper; prefer :func:`_write_result_record`."""
-    _write_result_record(record)
 
 
 _server_lock = threading.Lock()
@@ -286,6 +215,14 @@ _server_lock = threading.Lock()
 # ---------------------------------------------------------------------------
 
 
+def _get_open_port() -> int:
+    """Return an available TCP port on localhost."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
+
+
 def _wait_for_port(host: str, port: int, timeout: int = 1200, proc: subprocess.Popen | None = None) -> None:
     """Block until the given host:port accepts connections or timeout expires.
 
@@ -375,7 +312,7 @@ class DiffusionServer:
         self.model = server_cfg["model"]
         self.serve_args = server_cfg["serve_args"]
         self.host = "127.0.0.1"
-        self.port = port if port is not None else get_open_port(self.host)
+        self.port = port if port is not None else _get_open_port()
         self.proc: subprocess.Popen | None = None
         self.test_name: str = ""
 
@@ -545,7 +482,6 @@ def _unique_server_params(configs: list[dict[str, Any]]) -> list[dict[str, Any]]
                 "serve_args": _build_serve_args(serve_args_dict),
                 "benchmark_endpoint": cfg.get("benchmark_endpoint", cfg.get("benchmark_backend")),
                 "server_params": cfg["server_params"],
-                "mark": cfg.get("mark"),
             }
         )
     return result
@@ -560,29 +496,19 @@ def _test_param_mapping(configs: list[dict[str, Any]]) -> dict[str, list[dict]]:
     return mapping
 
 
-def _marks_by_test_name(configs: list[dict[str, Any]]) -> dict[str, list[pytest.MarkDecorator]]:
-    return {str(cfg["test_name"]): resolve_pytest_marks(cfg.get("mark")) for cfg in configs}
-
-
-def _paired_diffusion_benchmark_pytest_params(configs: list[dict[str, Any]]) -> list[Any]:
-    """Paired params for ``run_diffusion_benchmark.py``; same shape as omni runner."""
-    test_param_map = _test_param_mapping(configs)
-    server_entries = [(cfg, cfg["test_name"]) for cfg in _unique_server_params(configs)]
-    return create_paired_benchmark_pytest_params(server_entries, test_param_map, _marks_by_test_name(configs))
+def _make_server(server_cfg: dict[str, Any]) -> DiffusionServer:
+    """Factory: return a vLLM-Omni diffusion server instance for the config."""
+    return DiffusionServer(server_cfg=server_cfg)
 
 
 # ---------------------------------------------------------------------------
 # Parametrize data
 # ---------------------------------------------------------------------------
 
+server_params = _unique_server_params(BENCHMARK_CONFIGS)
 test_param_map = _test_param_mapping(BENCHMARK_CONFIGS)
-paired_benchmark_params = _paired_diffusion_benchmark_pytest_params(BENCHMARK_CONFIGS)
-
-
-def _make_server(server_cfg: dict[str, Any]) -> DiffusionServer:
-    """Factory: return a vLLM-Omni diffusion server instance for the config."""
-    return DiffusionServer(server_cfg=server_cfg)
 
+benchmark_indices: list[int] = list(range(max(len(v) for v in test_param_map.values())))
 
 # ---------------------------------------------------------------------------
 # Pytest fixtures
@@ -608,13 +534,16 @@ def diffusion_server(request):
 
 
 @pytest.fixture
-def benchmark_params(request):
-    """Benchmark params for the paired server/index parametrization."""
-    test_name, param_index = request.param
+def benchmark_params(request, diffusion_server):
+    """Yield the benchmark params dict for the current (server, index) pair."""
+    param_index: int = request.param
+    test_name = diffusion_server.test_name
 
     params_list = test_param_map.get(test_name, [])
     if not params_list:
         raise ValueError(f"No benchmark params for test: {test_name}")
+    if param_index >= len(params_list):
+        pytest.skip(f"Param index {param_index} out of range for {test_name} (has {len(params_list)} params)")
 
     current = param_index + 1
     total = len(params_list)
@@ -646,8 +575,8 @@ def run_benchmark(
     The raw metrics are written to a temporary file by the subprocess.  After
     the run completes the metrics are merged with full metadata (test_name,
     endpoint, benchmark_params, timestamp, flat reporting fields) and appended
-    to ``diffusion_result_{config_stem}_{hardware}_{timestamp}.json`` for the
-    source JSON file. The temporary file is removed afterwards.  Subprocess stdout/stderr are tee'd
+    to the session-wide aggregated JSON file (AGGREGATED_RESULT_FILE).  The
+    temporary file is removed afterwards.  Subprocess stdout/stderr are tee'd
     to a .log file under BENCHMARK_RESULT_DIR/logs/; its path is stored in
     the record.
     """
@@ -657,12 +586,7 @@ def run_benchmark(
     log_dir = BENCHMARK_RESULT_DIR / "logs"
     log_dir.mkdir(parents=True, exist_ok=True)
     endpoint_label = endpoint_filename_token(endpoint)
-    resource_label = get_runtime_resource_label()
-    hw_for_filename = resource_label_for_filename(resource_label)
-    if hw_for_filename:
-        log_file = log_dir / f"{test_name}_{hw_for_filename}_{endpoint_label}_{timestamp}.log"
-    else:
-        log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log"
+    log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log"
 
     with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
         tmp_result_file = Path(tmp.name)
@@ -772,7 +696,7 @@ def run_benchmark(
         "Model": model,
         "Framework": server_type,
         "API Endpoint": endpoint,
-        "Hardware": hardware_json_value(resource_label),
+        "Hardware": "",
         "Deployment": "",
         "Task": params.get("task", "t2i"),
         "Dataset": params.get("dataset", "random"),
@@ -804,8 +728,8 @@ def run_benchmark(
         "build_url": os.environ.get("BUILDKITE_BUILD_URL", ""),
         "source_file": source_file,
     }
-    result_path = _write_result_record(record)
-    print(f"\n  Result saved to: {result_path}")
+    _append_to_aggregated_file(record)
+    print(f"\n  Result appended to: {AGGREGATED_RESULT_FILE}")
     print(f"  Log saved to:       {log_file}")
 
     return metrics
@@ -1008,10 +932,12 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
 # ---------------------------------------------------------------------------
 @pytest.mark.benchmark
 @pytest.mark.parametrize(
-    "diffusion_server,benchmark_params",
-    paired_benchmark_params,
-    indirect=["diffusion_server", "benchmark_params"],
+    "diffusion_server",
+    server_params,
+    ids=[p["test_name"] for p in server_params],
+    indirect=True,
 )
+@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
 def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request):
     """Run the diffusion performance benchmark and verify request completion.
 
@@ -1036,10 +962,7 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
             test_name=test_name,
             endpoint=endpoint,
             server_cfg=server_cfg,
-            source_file=server_cfg.get(
-                _DIFFUSION_SOURCE_CONFIG_KEY,
-                CONFIG_FILE_PATH or f"{_PERF_TESTS_DIR}/*.json",
-            ),
+            source_file=cast(str, CONFIG_FILE_PATH),
         )
 
         print(f"\n{'=' * 60}")
@@ -1057,9 +980,7 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
             if key in result:
                 print(f"  {key}: {result[key]:.4f}")
 
-        source = server_cfg.get(_DIFFUSION_SOURCE_CONFIG_KEY) or CONFIG_FILE_PATH
-        if source:
-            print(f"\n  Aggregated results: {_aggregated_result_file_for_source(str(source))}")
+        print(f"\n  Aggregated results: {AGGREGATED_RESULT_FILE}")
         print("=" * 60)
 
         assert_baseline = request.config.getoption("--assert-baseline", default=False)

tests/dfx/perf/tests/test_bagel_vllm_omni.json

@@ -1,18 +1,6 @@
 [
   {
     "test_name": "test_bagel_single_device_single_stage_t2i",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
     "description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
     "server_type": "vllm-omni",
     "server_params": {
@@ -50,18 +38,6 @@
   },
   {
     "test_name": "test_bagel_single_device_single_stage_i2i",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
     "description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
     "server_type": "vllm-omni",
     "server_params": {
@@ -100,18 +76,6 @@
   },
   {
     "test_name": "test_bagel_single_device_multi_stage_t2i",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
     "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
     "server_type": "vllm-omni",
     "server_params": {
@@ -149,18 +113,6 @@
   },
   {
     "test_name": "test_bagel_single_device_multi_stage_i2i",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
     "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
     "server_type": "vllm-omni",
     "server_params": {

tests/dfx/perf/tests/test_cosmos3_vllm_omni.json

@@ -1,261 +1,213 @@
 [
-  {
-    "test_name": "test_cosmos3_t2i_official_demo_2gpu",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Cosmos3-Nano text-to-image official demo workload on 2 GPUs",
-    "server_type": "vllm-omni",
-    "benchmark_endpoint": "/v1/images/generations",
-    "server_params": {
-      "model": "nvidia/Cosmos3-Nano",
-      "serve_args": {
-        "cfg-parallel-size": 2,
-        "use-hsdp": true,
-        "hsdp-shard-size": 2,
-        "vae-patch-parallel-size": 2,
-        "vae-use-tiling": true,
-        "model-class-name": "Cosmos3OmniDiffusersPipeline",
-        "no-guardrails": true,
-        "stage-init-timeout": 900,
-        "init-timeout": 1200,
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "1024x1024_steps4",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 1024,
-        "height": 1024,
-        "num-inference-steps": 4,
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "seed": 42,
-        "extra-body": {
-          "guidance_scale": 7.0,
-          "flow_shift": 3.0,
-          "guardrails": false
+    {
+        "test_name": "test_cosmos3_t2i_official_demo_2gpu",
+        "description": "Cosmos3-Nano text-to-image official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/images/generations",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
         },
-        "baseline": {
-          "throughput_qps": 1.107,
-          "latency_mean": 0.8929,
-          "peak_memory_mb_mean": 80000.0
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_cosmos3_t2v_official_demo_2gpu",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Cosmos3-Nano text-to-video official demo workload on 2 GPUs",
-    "server_type": "vllm-omni",
-    "benchmark_endpoint": "/v1/videos",
-    "server_params": {
-      "model": "nvidia/Cosmos3-Nano",
-      "serve_args": {
-        "cfg-parallel-size": 2,
-        "use-hsdp": true,
-        "hsdp-shard-size": 2,
-        "vae-patch-parallel-size": 2,
-        "vae-use-tiling": true,
-        "model-class-name": "Cosmos3OmniDiffusersPipeline",
-        "no-guardrails": true,
-        "stage-init-timeout": 900,
-        "init-timeout": 1200,
-        "enable-diffusion-pipeline-profiler": true
-      }
+        "benchmark_params": [
+            {
+                "name": "1024x1024_steps4",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 1024,
+                "height": 1024,
+                "num-inference-steps": 4,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "seed": 42,
+                "extra-body": {
+                    "guidance_scale": 7.0,
+                    "flow_shift": 3.0,
+                    "guardrails": false
+                },
+                "baseline": {
+                    "throughput_qps": 1.107,
+                    "latency_mean": 0.8929,
+                    "peak_memory_mb_mean": 80000.0
+                }
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "1280x720_frames189_steps4",
-        "dataset": "random",
-        "task": "t2v",
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "seed": 42,
-        "random-request-config": [
-          {
-            "width": 1280,
-            "height": 720,
-            "prompt": "A robot arm is cleaning a plate in the kitchen.",
-            "num_inference_steps": 4,
-            "num_frames": 189,
-            "fps": 24,
-            "weight": 1
-          }
-        ],
-        "extra-body": {
-          "guidance_scale": 6.0,
-          "flow_shift": 10.0,
-          "max_sequence_length": 4096,
-          "guardrails": false,
-          "use_resolution_template": false,
-          "use_duration_template": false
+    {
+        "test_name": "test_cosmos3_t2v_official_demo_2gpu",
+        "description": "Cosmos3-Nano text-to-video official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/videos",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
         },
-        "baseline": {
-          "throughput_qps": 0.027,
-          "latency_mean": 36.6824,
-          "peak_memory_mb_mean": 30155.4
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_cosmos3_i2v_official_demo_2gpu",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Cosmos3-Nano image-to-video official demo workload on 2 GPUs",
-    "server_type": "vllm-omni",
-    "benchmark_endpoint": "/v1/videos",
-    "server_params": {
-      "model": "nvidia/Cosmos3-Nano",
-      "serve_args": {
-        "cfg-parallel-size": 2,
-        "use-hsdp": true,
-        "hsdp-shard-size": 2,
-        "vae-patch-parallel-size": 2,
-        "vae-use-tiling": true,
-        "model-class-name": "Cosmos3OmniDiffusersPipeline",
-        "no-guardrails": true,
-        "stage-init-timeout": 900,
-        "init-timeout": 1200,
-        "enable-diffusion-pipeline-profiler": true
-      }
+        "benchmark_params": [
+            {
+                "name": "1280x720_frames189_steps4",
+                "dataset": "random",
+                "task": "t2v",
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "seed": 42,
+                "random-request-config": [
+                    {
+                        "width": 1280,
+                        "height": 720,
+                        "prompt": "A robot arm is cleaning a plate in the kitchen.",
+                        "num_inference_steps": 4,
+                        "num_frames": 189,
+                        "fps": 24,
+                        "weight": 1
+                    }
+                ],
+                "extra-body": {
+                    "guidance_scale": 6.0,
+                    "flow_shift": 10.0,
+                    "max_sequence_length": 4096,
+                    "guardrails": false,
+                    "use_resolution_template": false,
+                    "use_duration_template": false
+                },
+                "baseline": {
+                    "throughput_qps": 0.027,
+                    "latency_mean": 36.6824,
+                    "peak_memory_mb_mean": 30155.4
+                }
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "1280x720_frames189_steps4",
-        "dataset": "random",
-        "task": "i2v",
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "num-input-images": 1,
-        "seed": 42,
-        "random-request-config": [
-          {
-            "width": 1280,
-            "height": 720,
-            "prompt": "The scene comes to life with smooth, natural motion.",
-            "num_inference_steps": 4,
-            "num_frames": 189,
-            "fps": 24,
-            "weight": 1
-          }
-        ],
-        "extra-body": {
-          "guidance_scale": 6.0,
-          "flow_shift": 10.0,
-          "max_sequence_length": 4096,
-          "guardrails": false
+    {
+        "test_name": "test_cosmos3_i2v_official_demo_2gpu",
+        "description": "Cosmos3-Nano image-to-video official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/videos",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
         },
-        "baseline": {
-          "throughput_qps": 0.027,
-          "latency_mean": 36.968,
-          "peak_memory_mb_mean": 30443.6
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_cosmos3_v2v_official_demo_2gpu",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Cosmos3-Nano video-to-video official demo workload on 2 GPUs",
-    "server_type": "vllm-omni",
-    "benchmark_endpoint": "/v1/videos",
-    "server_params": {
-      "model": "nvidia/Cosmos3-Nano",
-      "serve_args": {
-        "cfg-parallel-size": 2,
-        "use-hsdp": true,
-        "hsdp-shard-size": 2,
-        "vae-patch-parallel-size": 2,
-        "vae-use-tiling": true,
-        "model-class-name": "Cosmos3OmniDiffusersPipeline",
-        "no-guardrails": true,
-        "stage-init-timeout": 900,
-        "init-timeout": 1200,
-        "enable-diffusion-pipeline-profiler": true
-      }
+        "benchmark_params": [
+            {
+                "name": "1280x720_frames189_steps4",
+                "dataset": "random",
+                "task": "i2v",
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "num-input-images": 1,
+                "seed": 42,
+                "random-request-config": [
+                    {
+                        "width": 1280,
+                        "height": 720,
+                        "prompt": "The scene comes to life with smooth, natural motion.",
+                        "num_inference_steps": 4,
+                        "num_frames": 189,
+                        "fps": 24,
+                        "weight": 1
+                    }
+                ],
+                "extra-body": {
+                    "guidance_scale": 6.0,
+                    "flow_shift": 10.0,
+                    "max_sequence_length": 4096,
+                    "guardrails": false
+                },
+                "baseline": {
+                    "throughput_qps": 0.027,
+                    "latency_mean": 36.968,
+                    "peak_memory_mb_mean": 30443.6
+                }
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "1280x720_frames189_steps4",
-        "dataset": "random",
-        "task": "v2v",
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "seed": 42,
-        "random-request-config": [
-          {
-            "width": 1280,
-            "height": 720,
-            "prompt": "Continue the same scene with smooth natural motion.",
-            "num_inference_steps": 4,
-            "num_frames": 189,
-            "fps": 24,
-            "weight": 1
-          }
-        ],
-        "extra-body": {
-          "guidance_scale": 6.0,
-          "flow_shift": 10.0,
-          "max_sequence_length": 4096,
-          "guardrails": false,
-          "condition_frame_indexes_vision": [
-            0,
-            1
-          ],
-          "condition_video_keep": "first"
+    {
+        "test_name": "test_cosmos3_v2v_official_demo_2gpu",
+        "description": "Cosmos3-Nano video-to-video official demo workload on 2 GPUs",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/videos",
+        "server_params": {
+            "model": "nvidia/Cosmos3-Nano",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "use-hsdp": true,
+                "hsdp-shard-size": 2,
+                "vae-patch-parallel-size": 2,
+                "vae-use-tiling": true,
+                "model-class-name": "Cosmos3OmniDiffusersPipeline",
+                "no-guardrails": true,
+                "stage-init-timeout": 900,
+                "init-timeout": 1200,
+                "enable-diffusion-pipeline-profiler": true
+            }
         },
-        "baseline": {
-          "throughput_qps": 0.027,
-          "latency_mean": 37.2117,
-          "peak_memory_mb_mean": 30135.6
-        }
-      }
-    ]
-  }
+        "benchmark_params": [
+            {
+                "name": "1280x720_frames189_steps4",
+                "dataset": "random",
+                "task": "v2v",
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "seed": 42,
+                "random-request-config": [
+                    {
+                        "width": 1280,
+                        "height": 720,
+                        "prompt": "Continue the same scene with smooth natural motion.",
+                        "num_inference_steps": 4,
+                        "num_frames": 189,
+                        "fps": 24,
+                        "weight": 1
+                    }
+                ],
+                "extra-body": {
+                    "guidance_scale": 6.0,
+                    "flow_shift": 10.0,
+                    "max_sequence_length": 4096,
+                    "guardrails": false,
+                    "condition_frame_indexes_vision": [
+                        0,
+                        1
+                    ],
+                    "condition_video_keep": "first"
+                },
+                "baseline": {
+                    "throughput_qps": 0.027,
+                    "latency_mean": 37.2117,
+                    "peak_memory_mb_mean": 30135.6
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_higgs_audio_v3.json

@@ -1,18 +1,6 @@
 [
   {
     "test_name": "test_higgs_audio_v3",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "tts"
-    ],
     "server_params": {
       "model": "bosonai/higgs-audio-v3-tts-4b",
       "stage_overrides": {

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -1,152 +1,140 @@
 [
-  {
-    "test_name": "test_hunyuan_image3_ti2i_inline",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 4
-        }
-      },
-      "diffusion",
-      "local_model"
-    ],
-    "description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config",
-    "server_type": "vllm-omni",
-    "benchmark_endpoint": "/v1/images/edits",
-    "server_params": {
-      "model": "tencent/HunyuanImage-3.0-Instruct",
-      "serve_args": {
-        "trust-remote-code": true,
-        "deploy-config-inline": {
-          "type": "yaml",
-          "content": {
-            "pipeline": "hunyuan_image_3_moe",
-            "async_chunk": false,
-            "trust_remote_code": true,
-            "connectors": {
-              "rdma_connector": {
-                "name": "MooncakeTransferEngineConnector",
-                "extra": {
-                  "host": "auto",
-                  "zmq_port": 50051,
-                  "protocol": "rdma",
-                  "device_name": "",
-                  "memory_pool_size": 17179869184,
-                  "memory_pool_device": "cpu"
-                }
-              }
-            },
-            "stages": [
-              {
-                "stage_id": 0,
-                "is_comprehension": true,
-                "final_output": true,
-                "final_output_type": "text",
-                "max_num_seqs": 16,
-                "gpu_memory_utilization": 0.95,
-                "enforce_eager": false,
-                "enable_prefix_caching": false,
-                "max_num_batched_tokens": 32768,
-                "devices": "0,1",
-                "tensor_parallel_size": 2,
-                "hf_overrides": {
-                  "rope_parameters": {
-                    "mrope_section": [
-                      0,
-                      32,
-                      32
-                    ],
-                    "rope_type": "default"
-                  }
+    {
+        "test_name": "test_hunyuan_image3_ti2i_inline",
+        "description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config",
+        "server_type": "vllm-omni",
+        "benchmark_endpoint": "/v1/images/edits",
+        "server_params": {
+            "model": "tencent/HunyuanImage-3.0-Instruct",
+            "serve_args": {
+                "trust-remote-code": true,
+                "deploy-config-inline": {
+                    "type": "yaml",
+                    "content": {
+                        "pipeline": "hunyuan_image_3_moe",
+                        "async_chunk": false,
+                        "trust_remote_code": true,
+                        "connectors": {
+                            "rdma_connector": {
+                                "name": "MooncakeTransferEngineConnector",
+                                "extra": {
+                                    "host": "auto",
+                                    "zmq_port": 50051,
+                                    "protocol": "rdma",
+                                    "device_name": "",
+                                    "memory_pool_size": 17179869184,
+                                    "memory_pool_device": "cpu"
+                                }
+                            }
+                        },
+                        "stages": [
+                            {
+                                "stage_id": 0,
+                                "is_comprehension": true,
+                                "final_output": true,
+                                "final_output_type": "text",
+                                "max_num_seqs": 16,
+                                "gpu_memory_utilization": 0.95,
+                                "enforce_eager": false,
+                                "enable_prefix_caching": false,
+                                "max_num_batched_tokens": 32768,
+                                "devices": "0,1",
+                                "tensor_parallel_size": 2,
+                                "hf_overrides": {
+                                    "rope_parameters": {
+                                        "mrope_section": [
+                                            0,
+                                            32,
+                                            32
+                                        ],
+                                        "rope_type": "default"
+                                    }
+                                },
+                                "omni_kv_config": {
+                                    "need_send_cache": true
+                                },
+                                "output_connectors": {
+                                    "to_stage_1": "rdma_connector"
+                                },
+                                "default_sampling_params": {
+                                    "temperature": 0.0,
+                                    "top_p": 1,
+                                    "top_k": -1,
+                                    "max_tokens": 8192,
+                                    "detokenize": true,
+                                    "skip_special_tokens": false,
+                                    "stop_token_ids": [
+                                        128025
+                                    ]
+                                }
+                            },
+                            {
+                                "stage_id": 1,
+                                "max_num_seqs": 1,
+                                "enforce_eager": true,
+                                "devices": "2,3",
+                                "distributed_executor_backend": "mp",
+                                "omni_kv_config": {
+                                    "need_recv_cache": true
+                                },
+                                "parallel_config": {
+                                    "tensor_parallel_size": 2,
+                                    "vae_patch_parallel_size": 2,
+                                    "enable_expert_parallel": true
+                                },
+                                "input_connectors": {
+                                    "from_stage_0": "rdma_connector"
+                                },
+                                "default_sampling_params": {
+                                    "num_inference_steps": 50,
+                                    "guidance_scale": 0
+                                }
+                            }
+                        ],
+                        "edges": [
+                            {
+                                "from": 0,
+                                "to": 1,
+                                "window_size": -1,
+                                "max_inflight": 1
+                            }
+                        ]
+                    }
                 },
-                "omni_kv_config": {
-                  "need_send_cache": true
-                },
-                "output_connectors": {
-                  "to_stage_1": "rdma_connector"
-                },
-                "default_sampling_params": {
-                  "temperature": 0.0,
-                  "top_p": 1,
-                  "top_k": -1,
-                  "max_tokens": 8192,
-                  "detokenize": true,
-                  "skip_special_tokens": false,
-                  "stop_token_ids": [
-                    128025
-                  ]
-                }
-              },
-              {
-                "stage_id": 1,
-                "max_num_seqs": 1,
-                "enforce_eager": true,
-                "devices": "2,3",
-                "distributed_executor_backend": "mp",
-                "omni_kv_config": {
-                  "need_recv_cache": true
-                },
-                "parallel_config": {
-                  "tensor_parallel_size": 2,
-                  "vae_patch_parallel_size": 2,
-                  "enable_expert_parallel": true
-                },
-                "input_connectors": {
-                  "from_stage_0": "rdma_connector"
-                },
-                "default_sampling_params": {
-                  "num_inference_steps": 50,
-                  "guidance_scale": 0
-                }
-              }
-            ],
-            "edges": [
-              {
-                "from": 0,
-                "to": 1,
-                "window_size": -1,
-                "max_inflight": 1
-              }
-            ]
-          }
-        },
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "inline_dataset_5_requests",
-        "dataset": "custom",
-        "dataset-path-inline": {
-          "type": "jsonl",
-          "content": [
-            {
-              "prompt": "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。",
-              "image_urls": [
-                "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_0_0.png"
-              ],
-              "num_inference_steps": 8,
-              "seed": 42,
-              "guidance_scale": 1.0
+                "enable-diffusion-pipeline-profiler": true
             }
-          ]
         },
-        "task": "ti2i",
-        "bot-task": "think_recaption",
-        "max-concurrency": 16,
-        "num-prompts": 96,
-        "skip-performance-assertion": true,
-        "baseline": {
-          "throughput_qps": 0.1874,
-          "latency_p99": 101.852,
-          "latency_mean": 79.0385,
-          "peak_memory_mb_max": 70000,
-          "peak_memory_mb_mean": 70000
-        }
-      }
-    ]
-  }
+        "benchmark_params": [
+            {
+                "name": "inline_dataset_5_requests",
+                "dataset": "custom",
+                "dataset-path-inline": {
+                    "type": "jsonl",
+                    "content": [
+                        {
+                            "prompt": "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。",
+                            "image_urls": [
+                                "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_0_0.png"
+                            ],
+                            "num_inference_steps": 8,
+                            "seed": 42,
+                            "guidance_scale": 1.0
+                        }
+                    ]
+                },
+                "task": "ti2i",
+                "bot-task": "think_recaption",
+                "max-concurrency": 16,
+                "num-prompts": 96,
+                "skip-performance-assertion": true,
+                "baseline": {
+                    "throughput_qps": 0.1874,
+                    "latency_p99": 101.852,
+                    "latency_mean": 79.0385,
+                    "peak_memory_mb_max": 70000,
+                    "peak_memory_mb_mean": 70000
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json

@@ -1,47 +1,35 @@
 [
-  {
-    "test_name": "test_hunyuan_image_tp2_cfgp2",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 4
-        }
-      },
-      "diffusion",
-      "local_model"
-    ],
-    "description": "TP=2 CfgP=2 baseline",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "tencent/HunyuanImage-3.0-Instruct",
-      "serve_args": {
-        "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
-        "tensor-parallel-size": 2,
-        "cfg-parallel-size": 2,
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "1024x1024_steps8",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 1024,
-        "height": 1024,
-        "num-inference-steps": 8,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "skip-performance-assertion": true,
-        "baseline": {
-          "throughput_qps": 0.21,
-          "latency_p99": 4.7469,
-          "peak_memory_mb_max": 101100,
-          "peak_memory_mb_mean": 101100
-        }
-      }
-    ]
-  }
+    {
+        "test_name": "test_hunyuan_image_tp2_cfgp2",
+        "description": "TP=2 CfgP=2 baseline",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "tencent/HunyuanImage-3.0-Instruct",
+            "serve_args": {
+                "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
+                "tensor-parallel-size": 2,
+                "cfg-parallel-size": 2,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1024x1024_steps8",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 1024,
+                "height": 1024,
+                "num-inference-steps": 8,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "skip-performance-assertion": true,
+                "baseline": {
+                    "throughput_qps": 0.21,
+                    "latency_p99": 4.7469,
+                    "peak_memory_mb_max": 101100,
+                    "peak_memory_mb_mean": 101100
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json

@@ -1,47 +1,35 @@
 [
-  {
-    "test_name": "test_hunyuan_image_tp2_sp2",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 4
-        }
-      },
-      "diffusion",
-      "local_model"
-    ],
-    "description": "TP=2 SP=2 baseline",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "tencent/HunyuanImage-3.0-Instruct",
-      "serve_args": {
-        "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
-        "tensor-parallel-size": 2,
-        "ulysses-degree": 2,
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "1024x1024_steps8",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 1024,
-        "height": 1024,
-        "num-inference-steps": 8,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "skip-performance-assertion": true,
-        "baseline": {
-          "throughput_qps": 0.20,
-          "latency_p99": 5.1025,
-          "peak_memory_mb_max": 97402,
-          "peak_memory_mb_mean": 97402
-        }
-      }
-    ]
-  }
+    {
+        "test_name": "test_hunyuan_image_tp2_sp2",
+        "description": "TP=2 SP=2 baseline",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "tencent/HunyuanImage-3.0-Instruct",
+            "serve_args": {
+                "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
+                "tensor-parallel-size": 2,
+                "ulysses-degree": 2,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1024x1024_steps8",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 1024,
+                "height": 1024,
+                "num-inference-steps": 8,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "skip-performance-assertion": true,
+                "baseline": {
+                    "throughput_qps": 0.20,
+                    "latency_p99": 5.1025,
+                    "peak_memory_mb_max": 97402,
+                    "peak_memory_mb_mean": 97402
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_hunyuan_image_tp4.json

@@ -1,47 +1,35 @@
 [
-  {
-    "test_name": "test_hunyuan_image_tp4",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 4
-        }
-      },
-      "diffusion",
-      "local_model"
-    ],
-    "description": "TP=4 baseline",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "tencent/HunyuanImage-3.0-Instruct",
-      "serve_args": {
-        "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
-        "tensor-parallel-size": 4,
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "1024x1024_steps8",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 1024,
-        "height": 1024,
-        "num-inference-steps": 8,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "skip-performance-assertion": true,
-        "baseline": {
-          "throughput_qps": 0.213,
-          "latency_mean": 4.6953,
-          "latency_p99": 4.8601,
-          "peak_memory_mb_max": 56912.0,
-          "peak_memory_mb_mean": 56912.0
-        }
-      }
-    ]
-  }
+    {
+        "test_name": "test_hunyuan_image_tp4",
+        "description": "TP=4 baseline",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "tencent/HunyuanImage-3.0-Instruct",
+            "serve_args": {
+                "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml",
+                "tensor-parallel-size": 4,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1024x1024_steps8",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 1024,
+                "height": 1024,
+                "num-inference-steps": 8,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "skip-performance-assertion": true,
+                "baseline": {
+                    "throughput_qps": 0.213,
+                    "latency_mean": 4.6953,
+                    "latency_p99": 4.8601,
+                    "peak_memory_mb_max": 56912.0,
+                    "peak_memory_mb_mean": 56912.0
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_ltx2_vllm_omni.json

@@ -1,273 +1,217 @@
 [
-  {
-    "test_name": "test_ltx2_baseline_eager",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Single-device baseline with enforce-eager (no torch.compile)",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Lightricks/LTX-2",
-      "serve_args": {
-        "enforce-eager": true,
-        "enable-diffusion-pipeline-profiler": true
-      }
+    {
+        "test_name": "test_ltx2_baseline_eager",
+        "description": "Single-device baseline with enforce-eager (no torch.compile)",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "enforce-eager": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "256x256_145f_steps6",
-        "dataset": "random",
-        "task": "t2v",
-        "backend": "v1/videos",
-        "width": 256,
-        "height": 256,
-        "num-frames": 145,
-        "fps": 24,
-        "num-inference-steps": 6,
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true
-      },
-      {
-        "name": "480x768_41f_steps20",
-        "dataset": "random",
-        "task": "t2v",
-        "backend": "v1/videos",
-        "width": 768,
-        "height": 480,
-        "num-frames": 41,
-        "fps": 24,
-        "num-inference-steps": 20,
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true
-      }
-    ]
-  },
-  {
-    "test_name": "test_ltx2_torch_compile",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Single-device with torch.compile (default, no enforce-eager)",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Lightricks/LTX-2",
-      "serve_args": {
-        "enable-diffusion-pipeline-profiler": true
-      }
+
+    {
+        "test_name": "test_ltx2_torch_compile",
+        "description": "Single-device with torch.compile (default, no enforce-eager)",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "256x256_145f_steps6",
-        "dataset": "random",
-        "task": "t2v",
-        "backend": "v1/videos",
-        "width": 256,
-        "height": 256,
-        "num-frames": 145,
-        "fps": 24,
-        "num-inference-steps": 6,
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true
-      },
-      {
-        "name": "480x768_41f_steps20",
-        "dataset": "random",
-        "task": "t2v",
-        "backend": "v1/videos",
-        "width": 768,
-        "height": 480,
-        "num-frames": 41,
-        "fps": 24,
-        "num-inference-steps": 20,
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true
-      }
-    ]
-  },
-  {
-    "test_name": "test_ltx2_cfg2_eager",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "CFG-parallel=2 with enforce-eager",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Lightricks/LTX-2",
-      "serve_args": {
-        "cfg-parallel-size": 2,
-        "enforce-eager": true,
-        "enable-diffusion-pipeline-profiler": true
-      }
+
+    {
+        "test_name": "test_ltx2_cfg2_eager",
+        "description": "CFG-parallel=2 with enforce-eager",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enforce-eager": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "256x256_145f_steps6",
-        "dataset": "random",
-        "task": "t2v",
-        "backend": "v1/videos",
-        "width": 256,
-        "height": 256,
-        "num-frames": 145,
-        "fps": 24,
-        "num-inference-steps": 6,
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true
-      },
-      {
-        "name": "480x768_41f_steps20",
-        "dataset": "random",
-        "task": "t2v",
-        "backend": "v1/videos",
-        "width": 768,
-        "height": 480,
-        "num-frames": 41,
-        "fps": 24,
-        "num-inference-steps": 20,
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true
-      }
-    ]
-  },
-  {
-    "test_name": "test_ltx2_cfg2_compile",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "CFG-parallel=2 with torch.compile",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Lightricks/LTX-2",
-      "serve_args": {
-        "cfg-parallel-size": 2,
-        "enable-diffusion-pipeline-profiler": true
-      }
+
+    {
+        "test_name": "test_ltx2_cfg2_compile",
+        "description": "CFG-parallel=2 with torch.compile",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "cfg-parallel-size": 2,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "256x256_145f_steps6",
-        "dataset": "random",
-        "task": "t2v",
-        "backend": "v1/videos",
-        "width": 256,
-        "height": 256,
-        "num-frames": 145,
-        "fps": 24,
-        "num-inference-steps": 6,
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true
-      },
-      {
-        "name": "480x768_41f_steps20",
-        "dataset": "random",
-        "task": "t2v",
-        "backend": "v1/videos",
-        "width": 768,
-        "height": 480,
-        "num-frames": 41,
-        "fps": 24,
-        "num-inference-steps": 20,
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true
-      }
-    ]
-  },
-  {
-    "test_name": "test_ltx2_cache_dit_eager",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "CacheDiT with enforce-eager",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Lightricks/LTX-2",
-      "serve_args": {
-        "cache-backend": "cache_dit",
-        "enforce-eager": true,
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "256x256_145f_steps6",
-        "dataset": "random",
-        "task": "t2v",
-        "backend": "v1/videos",
-        "width": 256,
-        "height": 256,
-        "num-frames": 145,
-        "fps": 24,
-        "num-inference-steps": 6,
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true
-      },
-      {
-        "name": "480x768_41f_steps20",
-        "dataset": "random",
-        "task": "t2v",
-        "backend": "v1/videos",
-        "width": 768,
-        "height": 480,
-        "num-frames": 41,
-        "fps": 24,
-        "num-inference-steps": 20,
-        "num-prompts": 3,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true
-      }
-    ]
-  }
+
+    {
+        "test_name": "test_ltx2_cache_dit_eager",
+        "description": "CacheDiT with enforce-eager",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Lightricks/LTX-2",
+            "serve_args": {
+                "cache-backend": "cache_dit",
+                "enforce-eager": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "256x256_145f_steps6",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 256,
+                "height": 256,
+                "num-frames": 145,
+                "fps": 24,
+                "num-inference-steps": 6,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            },
+            {
+                "name": "480x768_41f_steps20",
+                "dataset": "random",
+                "task": "t2v",
+                "backend": "v1/videos",
+                "width": 768,
+                "height": 480,
+                "num-frames": 41,
+                "fps": 24,
+                "num-inference-steps": 20,
+                "num-prompts": 3,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json

@@ -1,343 +1,331 @@
 [
-  {
-    "test_name": "test_qwen3_omni_chunk",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "omni"
-    ],
-    "server_params": {
-      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-      "extra_cli_args": [
-        "--async-chunk"
-      ]
-    },
-    "benchmark_params": [
-      {
-        "dataset_name": "random",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          4,
-          16,
-          32,
-          64,
-          128
-        ],
-        "max_concurrency": [
-          1,
-          4,
-          8,
-          16,
-          32
-        ],
-        "random_input_len": 2500,
-        "random_output_len": 900,
-        "ignore_eos": true,
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "mean_ttft_ms": [
-            135.4379,
-            187.0896,
-            400,
-            435.675,
-            724.6836
-          ],
-          "mean_e2el_ms": [
-            17619.9851,
-            34811.8335,
-            43656.9501,
-            56657.1548,
-            85097.6016
-          ],
-          "mean_audio_rtf": [
-            0.1279,
-            0.2238,
-            0.2735,
-            0.3712,
-            0.5634
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          10
-        ],
-        "request_rate": [
-          0.1
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "random_mm_base_items_per_request": 1,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "audio": 1
+    {
+        "test_name": "test_qwen3_omni_chunk",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "extra_cli_args": [
+                "--async-chunk"
+            ]
         },
-        "random_mm_bucket_config": {
-          "(0, 60, 3)": 1.0
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "mean_ttft_ms": [
-            340.329
-          ],
-          "mean_e2el_ms": [
-            4305.477
-          ],
-          "mean_audio_rtf": [
-            0.1305
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          40
-        ],
-        "request_rate": [
-          0.5
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "random_mm_base_items_per_request": 2,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "image": 1,
-          "video": 1
-        },
-        "random_mm_bucket_config": {
-          "(256, 256, 1)": 0.5,
-          "(720, 1280, 2)": 0.5
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "mean_ttft_ms": [
-            220.0909
-          ],
-          "mean_e2el_ms": [
-            5475.9565
-          ],
-          "mean_audio_rtf": [
-            0.1818
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          100
-        ],
-        "request_rate": [
-          1.0
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "random_mm_base_items_per_request": 3,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "image": 1,
-          "video": 1,
-          "audio": 1
-        },
-        "random_mm_bucket_config": {
-          "(256, 256, 1)": 0.34,
-          "(720, 1280, 2)": 0.33,
-          "(0, 60, 3)": 0.33
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "mean_ttft_ms": [
-            450.0155
-          ],
-          "mean_e2el_ms": [
-            8516.8316
-          ],
-          "mean_audio_rtf": [
-            0.2717
-          ]
-        }
-      },
-      {
-        "dataset_name": "random",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          4,
-          16,
-          32,
-          64,
-          128
-        ],
-        "max_concurrency": [
-          1,
-          4,
-          8,
-          16,
-          32
-        ],
-        "random_input_len": 2500,
-        "random_output_len": 900,
-        "ignore_eos": true,
-        "extra_body": {
-          "modalities": [
-            "text"
-          ]
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el",
-        "baseline": {
-          "mean_ttft_ms": [
-            78.4278,
-            195.8037,
-            333.7318,
-            561.805,
-            1038.043
-          ],
-          "mean_e2el_ms": [
-            4288.3029,
-            7429.9429,
-            9451.6184,
-            12607.5292,
-            18153.2547
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          10
-        ],
-        "request_rate": [
-          0.1
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "extra_body": {
-          "modalities": [
-            "text"
-          ]
-        },
-        "random_mm_base_items_per_request": 1,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "audio": 1
-        },
-        "random_mm_bucket_config": {
-          "(0, 60, 3)": 1.0
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el",
-        "baseline": {
-          "mean_ttft_ms": [
-            275.3875
-          ],
-          "mean_e2el_ms": [
-            730.6995
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          40
-        ],
-        "request_rate": [
-          0.5
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "extra_body": {
-          "modalities": [
-            "text"
-          ]
-        },
-        "random_mm_base_items_per_request": 2,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "image": 1,
-          "video": 1
-        },
-        "random_mm_bucket_config": {
-          "(256, 256, 1)": 0.5,
-          "(720, 1280, 2)": 0.5
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el",
-        "baseline": {
-          "mean_ttft_ms": [
-            166.2201
-          ],
-          "mean_e2el_ms": [
-            673.5789
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          100
-        ],
-        "request_rate": [
-          1.0
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "extra_body": {
-          "modalities": [
-            "text"
-          ]
-        },
-        "random_mm_base_items_per_request": 3,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "image": 1,
-          "video": 1,
-          "audio": 1
-        },
-        "random_mm_bucket_config": {
-          "(256, 256, 1)": 0.34,
-          "(720, 1280, 2)": 0.33,
-          "(0, 60, 3)": 0.33
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el",
-        "baseline": {
-          "mean_ttft_ms": [
-            357.7781
-          ],
-          "mean_e2el_ms": [
-            953.3493
-          ]
-        }
-      }
-    ]
-  }
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    4,
+                    16,
+                    32,
+                    64,
+                    128
+                ],
+                "max_concurrency": [
+                    1,
+                    4,
+                    8,
+                    16,
+                    32
+                ],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        135.4379,
+                        187.0896,
+                        400,
+                        435.675,
+                        724.6836
+                    ],
+                    "mean_e2el_ms": [
+                        17619.9851,
+                        34811.8335,
+                        43656.9501,
+                        56657.1548,
+                        85097.6016
+                    ],
+                    "mean_audio_rtf": [
+                        0.1279,
+                        0.2238,
+                        0.2735,
+                        0.3712,
+                        0.5634
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    10
+                ],
+                "request_rate": [
+                    0.1
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 1,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(0, 60, 3)": 1.0
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        340.329
+                    ],
+                    "mean_e2el_ms": [
+                        4305.477
+                    ],
+                    "mean_audio_rtf": [
+                        0.1305
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    40
+                ],
+                "request_rate": [
+                    0.5
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 2,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.5,
+                    "(720, 1280, 2)": 0.5
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        220.0909
+                    ],
+                    "mean_e2el_ms": [
+                        5475.9565
+                    ],
+                    "mean_audio_rtf": [
+                        0.1818
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    100
+                ],
+                "request_rate": [
+                    1.0
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 3,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1,
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.34,
+                    "(720, 1280, 2)": 0.33,
+                    "(0, 60, 3)": 0.33
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        450.0155
+                    ],
+                    "mean_e2el_ms": [
+                        8516.8316
+                    ],
+                    "mean_audio_rtf": [
+                        0.2717
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    4,
+                    16,
+                    32,
+                    64,
+                    128
+                ],
+                "max_concurrency": [
+                    1,
+                    4,
+                    8,
+                    16,
+                    32
+                ],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": [
+                        "text"
+                    ]
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        78.4278,
+                        195.8037,
+                        333.7318,
+                        561.805,
+                        1038.043
+                    ],
+                    "mean_e2el_ms": [
+                        4288.3029,
+                        7429.9429,
+                        9451.6184,
+                        12607.5292,
+                        18153.2547
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    10
+                ],
+                "request_rate": [
+                    0.1
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": [
+                        "text"
+                    ]
+                },
+                "random_mm_base_items_per_request": 1,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(0, 60, 3)": 1.0
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        275.3875
+                    ],
+                    "mean_e2el_ms": [
+                        730.6995
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    40
+                ],
+                "request_rate": [
+                    0.5
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": [
+                        "text"
+                    ]
+                },
+                "random_mm_base_items_per_request": 2,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.5,
+                    "(720, 1280, 2)": 0.5
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        166.2201
+                    ],
+                    "mean_e2el_ms": [
+                        673.5789
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    100
+                ],
+                "request_rate": [
+                    1.0
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": [
+                        "text"
+                    ]
+                },
+                "random_mm_base_items_per_request": 3,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1,
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.34,
+                    "(720, 1280, 2)": 0.33,
+                    "(0, 60, 3)": 0.33
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        357.7781
+                    ],
+                    "mean_e2el_ms": [
+                        953.3493
+                    ]
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json

@@ -1,127 +1,115 @@
 [
-  {
-    "test_name": "test_qwen3_omni_3gpu_replica2",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 3
-        }
-      },
-      "full_model",
-      "omni"
-    ],
-    "server_params": {
-      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-      "stage_overrides": {
-          "1": {"num_replicas": 2, "devices": "1,2"},
-          "2": {"num_replicas": 2, "devices": "1,2"}
-      },
-      "extra_cli_args": [
-        "--async-chunk"
-      ]
-    },
-    "benchmark_params": [
-      {
-        "dataset_name": "random",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [128, 128, 128, 128],
-        "max_concurrency": [8, 16, 24, 32],
-        "random_input_len": 2500,
-        "random_output_len": 900,
-        "ignore_eos": true,
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-            "mean_ttft_ms": [203, 283, 470, 869],
-            "mean_tpot_ms": [28, 36, 41, 45],
-            "mean_audio_ttfp_ms": [861, 1200, 1528, 2100],
-            "mean_audio_rtf": [0.24, 0.30, 0.35, 0.39]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [10],
-        "request_rate": [0.1],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "random_mm_base_items_per_request": 1,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "audio": 1
+    {
+        "test_name": "test_qwen3_omni_3gpu_replica2",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "stage_overrides": {
+                "1": {"num_replicas": 2, "devices": "1,2"},
+                "2": {"num_replicas": 2, "devices": "1,2"}
+            },
+            "extra_cli_args": [
+                "--async-chunk"
+            ]
         },
-        "random_mm_bucket_config": {
-          "(0, 60, 3)": 1.0
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-            "mean_ttft_ms": [377],
-            "mean_audio_ttfp_ms": [684],
-            "mean_audio_rtf": [0.14]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [40],
-        "request_rate": [0.5],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "random_mm_base_items_per_request": 2,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "image": 1,
-          "video": 1
-        },
-        "random_mm_bucket_config": {
-          "(256, 256, 1)": 0.5,
-          "(720, 1280, 2)": 0.5
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-            "mean_ttft_ms": [228],
-            "mean_audio_ttfp_ms": [686],
-            "mean_audio_rtf": [0.16]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [100],
-        "request_rate": [1.0],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "random_mm_base_items_per_request": 3,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "image": 1,
-          "video": 1,
-          "audio": 1
-        },
-        "random_mm_bucket_config": {
-          "(256, 256, 1)": 0.34,
-          "(720, 1280, 2)": 0.33,
-          "(0, 60, 3)": 0.33
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-            "mean_ttft_ms": [522],
-            "mean_audio_ttfp_ms": [1119],
-            "mean_audio_rtf": [0.21]
-        }
-      }
-    ]
-  }
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [128, 128, 128, 128],
+                "max_concurrency": [8, 16, 24, 32],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [203, 283, 470, 869],
+                    "mean_tpot_ms": [28, 36, 41, 45],
+                    "mean_audio_ttfp_ms": [861, 1200, 1528, 2100],
+                    "mean_audio_rtf": [0.24, 0.30, 0.35, 0.39]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [10],
+                "request_rate": [0.1],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 1,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(0, 60, 3)": 1.0
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [377],
+                    "mean_audio_ttfp_ms": [684],
+                    "mean_audio_rtf": [0.14]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [40],
+                "request_rate": [0.5],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 2,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.5,
+                    "(720, 1280, 2)": 0.5
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [228],
+                    "mean_audio_ttfp_ms": [686],
+                    "mean_audio_rtf": [0.16]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [100],
+                "request_rate": [1.0],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 3,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1,
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.34,
+                    "(720, 1280, 2)": 0.33,
+                    "(0, 60, 3)": 0.33
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [522],
+                    "mean_audio_ttfp_ms": [1119],
+                    "mean_audio_rtf": [0.21]
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json

@@ -1,182 +1,170 @@
 [
-  {
-    "test_name": "test_qwen3_omni",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "omni"
-    ],
-    "server_params": {
-      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-      "extra_cli_args": [
-        "--no-async-chunk"
-      ]
-    },
-    "benchmark_params": [
-      {
-        "dataset_name": "random",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          4,
-          16,
-          32,
-          64,
-          128
-        ],
-        "max_concurrency": [
-          1,
-          4,
-          8,
-          16,
-          32
-        ],
-        "random_input_len": 2500,
-        "random_output_len": 900,
-        "ignore_eos": true,
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "mean_ttft_ms": [
-            109.1218,
-            247.4411,
-            354.9384,
-            447.7885,
-            694.7487
-          ],
-          "mean_e2el_ms": [
-            22413.667,
-            30808.9715,
-            35677.844,
-            46345.8978,
-            68350.4853
-          ],
-          "mean_audio_rtf": [
-            0.1695,
-            0.2033,
-            0.2467,
-            0.3234,
-            0.4438
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          10
-        ],
-        "request_rate": [
-          0.1
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "random_mm_base_items_per_request": 1,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "audio": 1
+    {
+        "test_name": "test_qwen3_omni",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "extra_cli_args": [
+                "--no-async-chunk"
+            ]
         },
-        "random_mm_bucket_config": {
-          "(0, 60, 3)": 1.0
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "mean_ttft_ms": [
-            323.3573
-          ],
-          "mean_e2el_ms": [
-            4623.4173
-          ],
-          "mean_audio_rtf": [
-            0.1406
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          40
-        ],
-        "request_rate": [
-          0.5
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "random_mm_base_items_per_request": 2,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "image": 1,
-          "video": 1
-        },
-        "random_mm_bucket_config": {
-          "(256, 256, 1)": 0.5,
-          "(720, 1280, 2)": 0.5
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "mean_ttft_ms": [
-            215.4502
-          ],
-          "mean_e2el_ms": [
-            4924.149
-          ],
-          "mean_audio_rtf": [
-            0.1637
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "openai-chat-omni",
-        "endpoint": "/v1/chat/completions",
-        "num_prompts": [
-          100
-        ],
-        "request_rate": [
-          1.0
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "random_mm_base_items_per_request": 3,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "image": 1,
-          "video": 1,
-          "audio": 1
-        },
-        "random_mm_bucket_config": {
-          "(256, 256, 1)": 0.34,
-          "(720, 1280, 2)": 0.33,
-          "(0, 60, 3)": 0.33
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
-        "baseline": {
-          "mean_ttft_ms": [
-            442.5815
-          ],
-          "mean_e2el_ms": [
-            6441.0986
-          ],
-          "mean_audio_rtf": [
-            0.2091
-          ]
-        }
-      }
-    ]
-  }
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    4,
+                    16,
+                    32,
+                    64,
+                    128
+                ],
+                "max_concurrency": [
+                    1,
+                    4,
+                    8,
+                    16,
+                    32
+                ],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        109.1218,
+                        247.4411,
+                        354.9384,
+                        447.7885,
+                        694.7487
+                    ],
+                    "mean_e2el_ms": [
+                        22413.667,
+                        30808.9715,
+                        35677.844,
+                        46345.8978,
+                        68350.4853
+                    ],
+                    "mean_audio_rtf": [
+                        0.1695,
+                        0.2033,
+                        0.2467,
+                        0.3234,
+                        0.4438
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    10
+                ],
+                "request_rate": [
+                    0.1
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 1,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(0, 60, 3)": 1.0
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        323.3573
+                    ],
+                    "mean_e2el_ms": [
+                        4623.4173
+                    ],
+                    "mean_audio_rtf": [
+                        0.1406
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    40
+                ],
+                "request_rate": [
+                    0.5
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 2,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.5,
+                    "(720, 1280, 2)": 0.5
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        215.4502
+                    ],
+                    "mean_e2el_ms": [
+                        4924.149
+                    ],
+                    "mean_audio_rtf": [
+                        0.1637
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "openai-chat-omni",
+                "endpoint": "/v1/chat/completions",
+                "num_prompts": [
+                    100
+                ],
+                "request_rate": [
+                    1.0
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "random_mm_base_items_per_request": 3,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1,
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.34,
+                    "(720, 1280, 2)": 0.33,
+                    "(0, 60, 3)": 0.33
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        442.5815
+                    ],
+                    "mean_e2el_ms": [
+                        6441.0986
+                    ],
+                    "mean_audio_rtf": [
+                        0.2091
+                    ]
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

@@ -1,192 +1,180 @@
 [
-  {
-    "test_name": "test_qwen3_omni_vllm_text",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "omni"
-    ],
-    "server_params": {
-      "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-      "use_omni": false,
-      "extra_cli_args": [
-        "--no-enable-prefix-caching",
-        "--skip-mm-profiling"
-      ]
-    },
-    "benchmark_params": [
-      {
-        "dataset_name": "random",
-        "backend": "vllm",
-        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-        "endpoint": "/v1/completions",
-        "num_prompts": [
-          4,
-          16,
-          32,
-          64,
-          128
-        ],
-        "max_concurrency": [
-          1,
-          4,
-          8,
-          16,
-          32
-        ],
-        "random_input_len": 2500,
-        "random_output_len": 900,
-        "ignore_eos": true,
-        "extra_body": {
-          "modalities": [
-            "text"
-          ]
+    {
+        "test_name": "test_qwen3_omni_vllm_text",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "use_omni": false,
+            "extra_cli_args": [
+                "--no-enable-prefix-caching",
+                "--skip-mm-profiling"
+            ]
         },
-        "percentile-metrics": "ttft,tpot,itl,e2el",
-        "baseline": {
-          "mean_ttft_ms": [
-            69.1276,
-            255.1882,
-            310.639,
-            422.0337,
-            532.0582
-          ],
-          "mean_e2el_ms": [
-            4382.7211,
-            7510.3815,
-            9440.4854,
-            12553.0186,
-            16932.6065
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "vllm",
-        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-        "endpoint": "/v1/completions",
-        "num_prompts": [
-          10
-        ],
-        "request_rate": [
-          0.1
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "extra_body": {
-          "modalities": [
-            "text"
-          ]
-        },
-        "random_mm_base_items_per_request": 1,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "audio": 1
-        },
-        "random_mm_bucket_config": {
-          "(0, 60, 3)": 1.0
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el",
-        "baseline": {
-          "mean_ttft_ms": [
-            26.3381
-          ],
-          "mean_e2el_ms": [
-            485.2436
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "vllm",
-        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-        "endpoint": "/v1/completions",
-        "num_prompts": [
-          40
-        ],
-        "request_rate": [
-          0.5
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "extra_body": {
-          "modalities": [
-            "text"
-          ]
-        },
-        "random_mm_base_items_per_request": 2,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "image": 1,
-          "video": 1
-        },
-        "random_mm_bucket_config": {
-          "(256, 256, 1)": 0.5,
-          "(720, 1280, 2)": 0.5
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el",
-        "baseline": {
-          "mean_ttft_ms": [
-            26.4517
-          ],
-          "mean_e2el_ms": [
-            517.2017
-          ]
-        }
-      },
-      {
-        "dataset_name": "random-mm",
-        "backend": "vllm",
-        "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-        "endpoint": "/v1/completions",
-        "num_prompts": [
-          100
-        ],
-        "request_rate": [
-          1.0
-        ],
-        "random_input_len": 100,
-        "random_output_len": 100,
-        "random_range_ratio": 0.0,
-        "ignore_eos": true,
-        "extra_body": {
-          "modalities": [
-            "text"
-          ]
-        },
-        "random_mm_base_items_per_request": 3,
-        "random_mm_num_mm_items_range_ratio": 0.5,
-        "random_mm_limit_mm_per_prompt": {
-          "image": 1,
-          "video": 1,
-          "audio": 1
-        },
-        "random_mm_bucket_config": {
-          "(256, 256, 1)": 0.34,
-          "(720, 1280, 2)": 0.33,
-          "(0, 60, 3)": 0.33
-        },
-        "percentile-metrics": "ttft,tpot,itl,e2el",
-        "baseline": {
-          "mean_ttft_ms": [
-            27.6972
-          ],
-          "mean_e2el_ms": [
-            561.3684
-          ]
-        }
-      }
-    ]
-  }
+        "benchmark_params": [
+            {
+                "dataset_name": "random",
+                "backend": "vllm",
+                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+                "endpoint": "/v1/completions",
+                "num_prompts": [
+                    4,
+                    16,
+                    32,
+                    64,
+                    128
+                ],
+                "max_concurrency": [
+                    1,
+                    4,
+                    8,
+                    16,
+                    32
+                ],
+                "random_input_len": 2500,
+                "random_output_len": 900,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": [
+                        "text"
+                    ]
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        69.1276,
+                        255.1882,
+                        310.639,
+                        422.0337,
+                        532.0582
+                    ],
+                    "mean_e2el_ms": [
+                        4382.7211,
+                        7510.3815,
+                        9440.4854,
+                        12553.0186,
+                        16932.6065
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "vllm",
+                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+                "endpoint": "/v1/completions",
+                "num_prompts": [
+                    10
+                ],
+                "request_rate": [
+                    0.1
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": [
+                        "text"
+                    ]
+                },
+                "random_mm_base_items_per_request": 1,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(0, 60, 3)": 1.0
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        26.3381
+                    ],
+                    "mean_e2el_ms": [
+                        485.2436
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "vllm",
+                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+                "endpoint": "/v1/completions",
+                "num_prompts": [
+                    40
+                ],
+                "request_rate": [
+                    0.5
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": [
+                        "text"
+                    ]
+                },
+                "random_mm_base_items_per_request": 2,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.5,
+                    "(720, 1280, 2)": 0.5
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        26.4517
+                    ],
+                    "mean_e2el_ms": [
+                        517.2017
+                    ]
+                }
+            },
+            {
+                "dataset_name": "random-mm",
+                "backend": "vllm",
+                "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+                "endpoint": "/v1/completions",
+                "num_prompts": [
+                    100
+                ],
+                "request_rate": [
+                    1.0
+                ],
+                "random_input_len": 100,
+                "random_output_len": 100,
+                "random_range_ratio": 0.0,
+                "ignore_eos": true,
+                "extra_body": {
+                    "modalities": [
+                        "text"
+                    ]
+                },
+                "random_mm_base_items_per_request": 3,
+                "random_mm_num_mm_items_range_ratio": 0.5,
+                "random_mm_limit_mm_per_prompt": {
+                    "image": 1,
+                    "video": 1,
+                    "audio": 1
+                },
+                "random_mm_bucket_config": {
+                    "(256, 256, 1)": 0.34,
+                    "(720, 1280, 2)": 0.33,
+                    "(0, 60, 3)": 0.33
+                },
+                "percentile-metrics": "ttft,tpot,itl,e2el",
+                "baseline": {
+                    "mean_ttft_ms": [
+                        27.6972
+                    ],
+                    "mean_e2el_ms": [
+                        561.3684
+                    ]
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json

@@ -1,185 +1,149 @@
 [
-  {
-    "test_name": "test_qwen_image_edit_2511_single_device",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Single-device baseline (two input images)",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image-Edit-2511",
-      "serve_args": {
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "512x512_steps20_i2i_2img",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 512,
-        "height": 512,
-        "num-inference-steps": 20,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "num-input-images": 2,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.0721,
-          "latency_mean": 13.8398,
-          "peak_memory_mb_max": 54731.4,
-          "peak_memory_mb_mean": 54731.4
-        }
-      },
-      {
-        "name": "1536x1536_steps35_i2i_2img",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 1536,
-        "height": 1536,
-        "num-inference-steps": 35,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "num-input-images": 2,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.0186,
-          "latency_mean": 53.6942,
-          "peak_memory_mb_max": 64402.4,
-          "peak_memory_mb_mean": 64402.4
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_vae_patch4",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 4
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image-Edit-2511",
-      "serve_args": {
-        "ulysses-degree": 2,
-        "cfg-parallel-size": 2,
-        "vae-patch-parallel-size": 4,
-        "vae-use-tiling": true,
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "1536x1536_steps35_i2i_2img",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 1536,
-        "height": 1536,
-        "num-inference-steps": 35,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "num-input-images": 2,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.0538,
-          "latency_mean": 18.5591,
-          "peak_memory_mb_max": 53485.0,
-          "peak_memory_mb_mean": 53485.0
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_cache_dit",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 4
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image-Edit-2511",
-      "serve_args": {
-        "ulysses-degree": 2,
-        "cfg-parallel-size": 2,
-        "cache-backend": "cache_dit",
-        "cache-config": {
-          "Fn_compute_blocks": 1,
-          "Bn_compute_blocks": 0,
-          "max_warmup_steps": 4,
-          "residual_diff_threshold": 0.24,
-          "max_continuous_cached_steps": 3,
-          "enable_taylorseer": false,
-          "taylorseer_order": 1,
-          "scm_steps_mask_policy": null,
-          "scm_steps_policy": "dynamic"
+    {
+        "test_name": "test_qwen_image_edit_2511_single_device",
+        "description": "Single-device baseline (two input images)",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image-Edit-2511",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true
+            }
         },
-        "enable-diffusion-pipeline-profiler": true
-      }
+        "benchmark_params": [
+            {
+                "name": "512x512_steps20_i2i_2img",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "num-input-images": 2,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.0721,
+                    "latency_mean": 13.8398,
+                    "peak_memory_mb_max": 54731.4,
+                    "peak_memory_mb_mean": 54731.4
+                }
+            },
+            {
+                "name": "1536x1536_steps35_i2i_2img",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 1536,
+                "height": 1536,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "num-input-images": 2,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.0186,
+                    "latency_mean": 53.6942,
+                    "peak_memory_mb_max": 64402.4,
+                    "peak_memory_mb_mean": 64402.4
+                }
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "512x512_steps20_i2i_2img",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 512,
-        "height": 512,
-        "num-inference-steps": 20,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "num-input-images": 2,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.2776,
-          "latency_mean": 3.5938,
-          "peak_memory_mb_max": 55214.0,
-          "peak_memory_mb_mean": 55214.0
-        }
-      },
-      {
-        "name": "1536x1536_steps35_i2i_2img",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 1536,
-        "height": 1536,
-        "num-inference-steps": 35,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "num-input-images": 2,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.1051,
-          "latency_mean": 9.4918,
-          "peak_memory_mb_max": 64265.6,
-          "peak_memory_mb_mean": 64265.6
-        }
-      }
-    ]
-  }
+    {
+        "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_vae_patch4",
+        "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image-Edit-2511",
+            "serve_args": {
+                "ulysses-degree": 2,
+                "cfg-parallel-size": 2,
+                "vae-patch-parallel-size": 4,
+                "vae-use-tiling": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1536x1536_steps35_i2i_2img",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 1536,
+                "height": 1536,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "num-input-images": 2,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.0538,
+                    "latency_mean": 18.5591,
+                    "peak_memory_mb_max": 53485.0,
+                    "peak_memory_mb_mean": 53485.0
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_cache_dit",
+        "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image-Edit-2511",
+            "serve_args": {
+                "ulysses-degree": 2,
+                "cfg-parallel-size": 2,
+                "cache-backend": "cache_dit",
+                "cache-config": {
+                    "Fn_compute_blocks": 1,
+                    "Bn_compute_blocks": 0,
+                    "max_warmup_steps": 4,
+                    "residual_diff_threshold": 0.24,
+                    "max_continuous_cached_steps": 3,
+                    "enable_taylorseer": false,
+                    "taylorseer_order": 1,
+                    "scm_steps_mask_policy": null,
+                    "scm_steps_policy": "dynamic"
+                },
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps20_i2i_2img",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "num-input-images": 2,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.2776,
+                    "latency_mean": 3.5938,
+                    "peak_memory_mb_max": 55214.0,
+                    "peak_memory_mb_mean": 55214.0
+                }
+            },
+            {
+                "name": "1536x1536_steps35_i2i_2img",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 1536,
+                "height": 1536,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "num-input-images": 2,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.1051,
+                    "latency_mean": 9.4918,
+                    "peak_memory_mb_max": 64265.6,
+                    "peak_memory_mb_mean": 64265.6
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json

@@ -1,180 +1,144 @@
 [
-  {
-    "test_name": "test_qwen_image_edit_single_device",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Single-device baseline",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image-Edit",
-      "serve_args": {
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "512x512_steps20_i2i",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 512,
-        "height": 512,
-        "num-inference-steps": 20,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.1129,
-          "latency_mean": 8.8352,
-          "peak_memory_mb_max": 55356.5,
-          "peak_memory_mb_mean": 55356.5
-        }
-      },
-      {
-        "name": "1536x1536_steps35_i2i",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 1536,
-        "height": 1536,
-        "num-inference-steps": 35,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.0240,
-          "latency_mean": 41.5686,
-          "peak_memory_mb_max": 62762.7,
-          "peak_memory_mb_mean": 62762.7
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_qwen_image_edit_ulysses2_cfg2_vae_patch4",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 4
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image-Edit",
-      "serve_args": {
-        "ulysses-degree": 2,
-        "cfg-parallel-size": 2,
-        "vae-patch-parallel-size": 4,
-        "vae-use-tiling": true,
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "1536x1536_steps35_i2i",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 1536,
-        "height": 1536,
-        "num-inference-steps": 35,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.0672,
-          "latency_mean": 14.8517,
-          "peak_memory_mb_max": 53931.5,
-          "peak_memory_mb_mean": 53931.5
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_qwen_image_edit_ulysses2_cfg2_cache_dit",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 4
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image-Edit",
-      "serve_args": {
-        "ulysses-degree": 2,
-        "cfg-parallel-size": 2,
-        "cache-backend": "cache_dit",
-        "cache-config": {
-          "Fn_compute_blocks": 1,
-          "Bn_compute_blocks": 0,
-          "max_warmup_steps": 4,
-          "residual_diff_threshold": 0.24,
-          "max_continuous_cached_steps": 3,
-          "enable_taylorseer": false,
-          "taylorseer_order": 1,
-          "scm_steps_mask_policy": null,
-          "scm_steps_policy": "dynamic"
+    {
+        "test_name": "test_qwen_image_edit_single_device",
+        "description": "Single-device baseline",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image-Edit",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true
+            }
         },
-        "enable-diffusion-pipeline-profiler": true
-      }
+        "benchmark_params": [
+            {
+                "name": "512x512_steps20_i2i",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.1129,
+                    "latency_mean": 8.8352,
+                    "peak_memory_mb_max": 55356.5,
+                    "peak_memory_mb_mean": 55356.5
+                }
+            },
+            {
+                "name": "1536x1536_steps35_i2i",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 1536,
+                "height": 1536,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.0240,
+                    "latency_mean": 41.5686,
+                    "peak_memory_mb_max": 62762.7,
+                    "peak_memory_mb_mean": 62762.7
+                }
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "512x512_steps20_i2i",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 512,
-        "height": 512,
-        "num-inference-steps": 20,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.2879,
-          "latency_mean": 3.4701,
-          "peak_memory_mb_max": 55525.6,
-          "peak_memory_mb_mean": 55525.6
-        }
-      },
-      {
-        "name": "1536x1536_steps35_i2i",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 1536,
-        "height": 1536,
-        "num-inference-steps": 35,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.1113,
-          "latency_mean": 8.9656,
-          "peak_memory_mb_max": 63758.3,
-          "peak_memory_mb_mean": 63758.3
-        }
-      }
-    ]
-  }
+    {
+        "test_name": "test_qwen_image_edit_ulysses2_cfg2_vae_patch4",
+        "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image-Edit",
+            "serve_args": {
+                "ulysses-degree": 2,
+                "cfg-parallel-size": 2,
+                "vae-patch-parallel-size": 4,
+                "vae-use-tiling": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1536x1536_steps35_i2i",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 1536,
+                "height": 1536,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.0672,
+                    "latency_mean": 14.8517,
+                    "peak_memory_mb_max": 53931.5,
+                    "peak_memory_mb_mean": 53931.5
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_qwen_image_edit_ulysses2_cfg2_cache_dit",
+        "description": "Ulysses SP=2 + CFG=2 + CacheDiT",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image-Edit",
+            "serve_args": {
+                "ulysses-degree": 2,
+                "cfg-parallel-size": 2,
+                "cache-backend": "cache_dit",
+                "cache-config": {
+                    "Fn_compute_blocks": 1,
+                    "Bn_compute_blocks": 0,
+                    "max_warmup_steps": 4,
+                    "residual_diff_threshold": 0.24,
+                    "max_continuous_cached_steps": 3,
+                    "enable_taylorseer": false,
+                    "taylorseer_order": 1,
+                    "scm_steps_mask_policy": null,
+                    "scm_steps_policy": "dynamic"
+                },
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps20_i2i",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.2879,
+                    "latency_mean": 3.4701,
+                    "peak_memory_mb_max": 55525.6,
+                    "peak_memory_mb_mean": 55525.6
+                }
+            },
+            {
+                "name": "1536x1536_steps35_i2i",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 1536,
+                "height": 1536,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.1113,
+                    "latency_mean": 8.9656,
+                    "peak_memory_mb_max": 63758.3,
+                    "peak_memory_mb_mean": 63758.3
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json

@@ -1,61 +1,49 @@
 [
-  {
-    "test_name": "test_qwen_image_layered_single_device",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Single-device baseline",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image-Layered",
-      "serve_args": {
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "640x640_steps20_i2i",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 640,
-        "height": 640,
-        "num-inference-steps": 20,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.0703,
-          "latency_mean": 14.208,
-          "peak_memory_mb_max": 60543.5,
-          "peak_memory_mb_mean": 60543.5
-        }
-      },
-      {
-        "name": "1024x1024_steps35_i2i",
-        "dataset": "random",
-        "task": "i2i",
-        "width": 1024,
-        "height": 1024,
-        "num-inference-steps": 35,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.0417,
-          "latency_mean": 23.914,
-          "peak_memory_mb_max": 60543.5,
-          "peak_memory_mb_mean": 60543.5
-        }
-      }
-    ]
-  }
+    {
+        "test_name": "test_qwen_image_layered_single_device",
+        "description": "Single-device baseline",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image-Layered",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "640x640_steps20_i2i",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 640,
+                "height": 640,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.0703,
+                    "latency_mean": 14.208,
+                    "peak_memory_mb_max": 60543.5,
+                    "peak_memory_mb_mean": 60543.5
+                }
+            },
+            {
+                "name": "1024x1024_steps35_i2i",
+                "dataset": "random",
+                "task": "i2i",
+                "width": 1024,
+                "height": 1024,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.0417,
+                    "latency_mean": 23.914,
+                    "peak_memory_mb_max": 60543.5,
+                    "peak_memory_mb_mean": 60543.5
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_qwen_image_vllm_omni.json

@@ -1,279 +1,219 @@
 [
-  {
-    "test_name": "test_qwen_image_single_device",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Single-device baseline (no parallelism)",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image",
-      "serve_args": {
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "512x512_steps20",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 512,
-        "height": 512,
-        "num-inference-steps": 20,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.4153,
-          "latency_mean": 2.4065,
-          "peak_memory_mb_mean": 52289.9
-        }
-      },
-      {
-        "name": "1536x1536_steps35",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 1536,
-        "height": 1536,
-        "num-inference-steps": 35,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.0423,
-          "latency_mean": 23.5991,
-          "peak_memory_mb_mean": 61360.5
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_qwen_image_single_device_step_execution",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Single-device baseline (no parallelism) with step execution",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image",
-      "serve_args": {
-        "enable-diffusion-pipeline-profiler": true,
-        "step-execution": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "512x512_steps20",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 512,
-        "height": 512,
-        "num-inference-steps": 20,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.4175,
-          "latency_mean": 2.3953,
-          "peak_memory_mb_mean": 56000
-        }
-      },
-      {
-        "name": "1536x1536_steps35",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 1536,
-        "height": 1536,
-        "num-inference-steps": 35,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.0423,
-          "latency_mean": 23.5768,
-          "peak_memory_mb_mean": 65000
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 4
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image",
-      "serve_args": {
-        "ulysses-degree": 2,
-        "cfg-parallel-size": 2,
-        "vae-patch-parallel-size": 4,
-        "vae-use-tiling": true,
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "1536x1536_steps35",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 1536,
-        "height": 1536,
-        "num-inference-steps": 35,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.1214,
-          "latency_mean": 8.2211,
-          "peak_memory_mb_mean": 51704.7
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_qwen_image_ulysses2_cfg2_cache_dit",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 4
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Ulysses SP=2 + CFG-parallel=2 + CacheDiT acceleration",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image",
-      "serve_args": {
-        "ulysses-degree": 2,
-        "cfg-parallel-size": 2,
-        "cache-backend": "cache_dit",
-        "cache-config": {
-          "Fn_compute_blocks": 1,
-          "Bn_compute_blocks": 0,
-          "max_warmup_steps": 4,
-          "residual_diff_threshold": 0.24,
-          "max_continuous_cached_steps": 3,
-          "enable_taylorseer": false,
-          "taylorseer_order": 1,
-          "scm_steps_mask_policy": null,
-          "scm_steps_policy": "dynamic"
+    {
+        "test_name": "test_qwen_image_single_device",
+        "description": "Single-device baseline (no parallelism)",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true
+            }
         },
-        "enable-diffusion-pipeline-profiler": true
-      }
+        "benchmark_params": [
+            {
+                "name": "512x512_steps20",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.4153,
+                    "latency_mean": 2.4065,
+                    "peak_memory_mb_mean": 52289.9
+                }
+            },
+            {
+                "name": "1536x1536_steps35",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 1536,
+                "height": 1536,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.0423,
+                    "latency_mean": 23.5991,
+                    "peak_memory_mb_mean": 61360.5
+                }
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "512x512_steps20",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 512,
-        "height": 512,
-        "num-inference-steps": 20,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.559,
-          "latency_mean": 1.794,
-          "peak_memory_mb_mean": 52397.6775
-        }
-      },
-      {
-        "name": "1536x1536_steps35",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 1536,
-        "height": 1536,
-        "num-inference-steps": 35,
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "enable-negative-prompt": true,
-        "baseline": {
-          "throughput_qps": 0.1918,
-          "latency_mean": 5.2022,
-          "peak_memory_mb_mean": 61468.8
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Qwen/Qwen-Image",
-      "serve_args": {
-        "enable-diffusion-pipeline-profiler": true,
-        "step-execution": true,
-        "max-num-seqs": 8
-      }
+    {
+        "test_name": "test_qwen_image_single_device_step_execution",
+        "description": "Single-device baseline (no parallelism) with step execution",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true,
+                "step-execution": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps20",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.4175,
+                    "latency_mean": 2.3953,
+                    "peak_memory_mb_mean": 56000
+                }
+            },
+            {
+                "name": "1536x1536_steps35",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 1536,
+                "height": 1536,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.0423,
+                    "latency_mean": 23.5768,
+                    "peak_memory_mb_mean": 65000
+                }
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "512x512_steps20_high_concurrency",
-        "dataset": "random",
-        "task": "t2i",
-        "width": 512,
-        "height": 512,
-        "num-inference-steps": 20,
-        "num-prompts": [20, 40, 80, 160],
-        "max-concurrency": [1, 2, 4, 8],
-        "warmup-requests": 8,
-        "warmup-concurrency": 8,
-        "warmup-num-inference-steps": 20,
-        "enable-negative-prompt": true,
-        "baseline": {
-            "throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
-            "latency_mean": [2.4215, 2.992, 5.9293, 9.656],
-            "peak_memory_mb_mean": [67000, 67000, 67000, 67000]
-        }
-      }
-    ]
-  }
+    {
+        "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4",
+        "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image",
+            "serve_args": {
+                "ulysses-degree": 2,
+                "cfg-parallel-size": 2,
+                "vae-patch-parallel-size": 4,
+                "vae-use-tiling": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "1536x1536_steps35",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 1536,
+                "height": 1536,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.1214,
+                    "latency_mean": 8.2211,
+                    "peak_memory_mb_mean": 51704.7
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_qwen_image_ulysses2_cfg2_cache_dit",
+        "description": "Ulysses SP=2 + CFG-parallel=2 + CacheDiT acceleration",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image",
+            "serve_args": {
+                "ulysses-degree": 2,
+                "cfg-parallel-size": 2,
+                "cache-backend": "cache_dit",
+                "cache-config": {
+                    "Fn_compute_blocks": 1,
+                    "Bn_compute_blocks": 0,
+                    "max_warmup_steps": 4,
+                    "residual_diff_threshold": 0.24,
+                    "max_continuous_cached_steps": 3,
+                    "enable_taylorseer": false,
+                    "taylorseer_order": 1,
+                    "scm_steps_mask_policy": null,
+                    "scm_steps_policy": "dynamic"
+                },
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps20",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.559,
+                    "latency_mean": 1.794,
+                    "peak_memory_mb_mean": 52397.6775
+                }
+            },
+            {
+                "name": "1536x1536_steps35",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 1536,
+                "height": 1536,
+                "num-inference-steps": 35,
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": 0.1918,
+                    "latency_mean": 5.2022,
+                    "peak_memory_mb_mean": 61468.8
+                }
+            }
+        ]
+    },
+    {
+        "test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
+        "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Qwen/Qwen-Image",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true,
+                "step-execution": true,
+                "max-num-seqs": 8
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "512x512_steps20_high_concurrency",
+                "dataset": "random",
+                "task": "t2i",
+                "width": 512,
+                "height": 512,
+                "num-inference-steps": 20,
+                "num-prompts": [20, 40, 80, 160],
+                "max-concurrency": [1, 2, 4, 8],
+                "warmup-requests": 8,
+                "warmup-concurrency": 8,
+                "warmup-num-inference-steps": 20,
+                "enable-negative-prompt": true,
+                "baseline": {
+                    "throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
+                    "latency_mean": [2.4215, 2.992, 5.9293, 9.656],
+                    "peak_memory_mb_mean": [67000, 67000, 67000, 67000]
+                }
+            }
+        ]
+    }
 ]

tests/dfx/perf/tests/test_runner_metadata.py

@@ -86,202 +86,3 @@ def test_enabled_false_entry_is_skipped():
     # Only the enabled=True entry should appear
     assert len(params) == 1
     assert params[0].get("task") == "voice_clone"
-
-
-def test_resolve_pytest_marks_hardware_dict_with_extra():
-    from tests.dfx.conftest import resolve_pytest_marks
-
-    marks = resolve_pytest_marks(
-        [
-            {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}},
-            "full_model",
-            "diffusion",
-            "local_model",
-        ]
-    )
-    names = {m.name for m in marks}
-    assert "H100" in names
-    assert "cuda" in names
-    assert "gpu" in names
-    assert "distributed_cuda" in names
-    assert "full_model" in names
-    assert "diffusion" in names
-    assert "local_model" in names
-
-
-def test_resolve_pytest_marks_rejects_legacy_object_format():
-    from tests.dfx.conftest import resolve_pytest_marks
-
-    with pytest.raises(ValueError, match="mark must be a list"):
-        resolve_pytest_marks(
-            {
-                "hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1},
-                "marks": ["full_model"],
-            }
-        )
-
-
-def test_extract_mark_resource_label():
-    from tests.dfx.conftest import extract_mark_resource_label
-
-    assert extract_mark_resource_label(None) == "na"
-    assert (
-        extract_mark_resource_label(
-            [
-                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
-                "full_model",
-            ]
-        )
-        == "H100"
-    )
-    assert (
-        extract_mark_resource_label(
-            [
-                {
-                    "hardware_marks": {
-                        "res": {"cuda": "H100", "rocm": "MI325"},
-                        "num_cards": 2,
-                    }
-                },
-                "diffusion",
-            ]
-        )
-        == "H100-MI325"
-    )
-
-
-def test_resource_label_for_filename():
-    from tests.dfx.conftest import resource_label_for_filename
-
-    assert resource_label_for_filename("H100") == ""
-    assert resource_label_for_filename("L4") == "L4"
-    assert resource_label_for_filename("910B") == "910B"
-    assert resource_label_for_filename("na") == "na"
-
-
-def test_hardware_json_value():
-    from tests.dfx.conftest import hardware_json_value
-
-    assert hardware_json_value("H100") == "H100"
-    assert hardware_json_value("na") == ""
-    assert hardware_json_value(None) == ""
-
-
-def test_extract_configs_resource_label(monkeypatch):
-    from tests.dfx.conftest import extract_configs_resource_label, get_runtime_resource_label
-
-    monkeypatch.setattr(
-        "tests.dfx.conftest._read_runtime_device_name",
-        lambda *, device_id=0: "NVIDIA H100 80GB HBM3",
-    )
-    get_runtime_resource_label(refresh=True)
-    assert extract_configs_resource_label([]) == "H100"
-    get_runtime_resource_label(refresh=True)
-    monkeypatch.setattr(
-        "tests.dfx.conftest._read_runtime_device_name",
-        lambda *, device_id=0: "Ascend910B2",
-    )
-    assert get_runtime_resource_label(refresh=True) == "910B"
-
-
-def test_load_benchmark_configs_from_dir(tmp_path):
-    from tests.dfx.conftest import load_benchmark_configs
-
-    (tmp_path / "a.json").write_text(
-        json.dumps([{"test_name": "test_a", "server_params": {"model": "m/a"}, "benchmark_params": []}]),
-        encoding="utf-8",
-    )
-    (tmp_path / "b.json").write_text(
-        json.dumps([{"test_name": "test_b", "server_params": {"model": "m/b"}, "benchmark_params": []}]),
-        encoding="utf-8",
-    )
-    configs = load_benchmark_configs(config_dir=tmp_path)
-    assert [c["test_name"] for c in configs] == ["test_a", "test_b"]
-
-
-def test_create_unique_server_pytest_params_applies_marks(tmp_path):
-    from tests.dfx.conftest import create_unique_server_pytest_params
-
-    configs = [
-        {
-            "test_name": "test_with_mark",
-            "mark": [
-                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
-                "full_model",
-            ],
-            "server_params": {"model": "some/model"},
-            "benchmark_params": [{"name": "p0", "num_prompts": 1}],
-        },
-        {
-            "test_name": "test_without_mark",
-            "server_params": {"model": "other/model"},
-            "benchmark_params": [{"name": "p0", "num_prompts": 1}],
-        },
-    ]
-    params = create_unique_server_pytest_params(configs, tmp_path)
-    by_id = {p.id: p for p in params}
-    assert len(by_id["test_with_mark"].values) == 1
-    assert isinstance(by_id["test_with_mark"].values[0], tuple)
-    assert any(m.name == "H100" for m in by_id["test_with_mark"].marks)
-    assert not any(m.name == "H100" for m in by_id["test_without_mark"].marks)
-
-
-def test_is_diffusion_perf_config():
-    from tests.dfx.conftest import is_diffusion_perf_config
-
-    assert not is_diffusion_perf_config(
-        {"test_name": "omni_a", "mark": [{"hardware_marks": {"res": {"cuda": "H100"}}}, "omni"]}
-    )
-    assert is_diffusion_perf_config(
-        {
-            "test_name": "diff_a",
-            "server_type": "vllm-omni",
-            "mark": [{"hardware_marks": {"res": {"cuda": "H100"}}}, "diffusion"],
-        }
-    )
-
-
-def test_benchmark_param_id_suffix_from_task_eval_phase():
-    from tests.dfx.conftest import _unique_benchmark_param_id_suffixes
-
-    params = [
-        {"task": "default_voice", "eval_phase": "latency"},
-        {"task": "default_voice", "eval_phase": "throughput"},
-    ]
-    assert _unique_benchmark_param_id_suffixes(params) == [
-        "default_voice_latency",
-        "default_voice_throughput",
-    ]
-
-
-def test_create_paired_omni_benchmark_pytest_params(tmp_path):
-    from tests.dfx.conftest import create_paired_omni_benchmark_pytest_params
-
-    configs = [
-        {
-            "test_name": "test_omni",
-            "mark": [
-                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
-                "omni",
-            ],
-            "server_params": {"model": "m/omni"},
-            "benchmark_params": [{"name": "p0", "num_prompts": 1}, {"name": "p1", "num_prompts": 2}],
-        },
-        {
-            "test_name": "test_tts",
-            "mark": [
-                {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}},
-                "tts",
-            ],
-            "server_params": {"model": "m/tts"},
-            "benchmark_params": [{"name": "p0", "num_prompts": 1}],
-        },
-    ]
-    params = create_paired_omni_benchmark_pytest_params(configs, tmp_path)
-    by_id = {p.id: p for p in params}
-    assert set(by_id) == {"test_omni-p0", "test_omni-p1", "test_tts-p0"}
-    omni_row, bench_row = by_id["test_tts-p0"].values
-    assert omni_row[0] == "test_tts"
-    assert bench_row == ("test_tts", 0)
-    assert any(m.name == "tts" for m in by_id["test_tts-p0"].marks)
-    assert not any(m.name == "tts" for m in by_id["test_omni-p0"].marks)

tests/dfx/perf/tests/test_tts.json

@@ -1,19 +1,6 @@
 [
   {
     "test_name": "test_qwen3_tts_base",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100",
-            "npu": "A3"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "tts"
-    ],
     "server_params": {
       "model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
       "stage_overrides": {
@@ -111,18 +98,6 @@
   },
   {
     "test_name": "test_qwen3_tts_customvoice",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "tts"
-    ],
     "server_params": {
       "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
       "stage_overrides": {
@@ -367,18 +342,6 @@
   },
   {
     "test_name": "test_voxtral_tts",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "tts"
-    ],
     "server_params": {
       "model": "mistralai/Voxtral-4B-TTS-2603"
     },

tests/dfx/perf/tests/test_voxcpm2.json

@@ -1,18 +1,6 @@
 [
   {
     "test_name": "test_voxcpm2",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "tts"
-    ],
     "server_params": {
       "model": "openbmb/VoxCPM2",
       "trust_remote_code": true,

tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json

@@ -1,131 +1,107 @@
 [
-  {
-    "test_name": "test_wan22_i2v_single_device",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 1
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "Single-device baseline",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-      "serve_args": {
-        "enable-diffusion-pipeline-profiler": true
-      }
+    {
+        "test_name": "test_wan22_i2v_single_device",
+        "description": "Single-device baseline",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+            "serve_args": {
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "832x480_frames81_steps4",
+                "dataset": "random",
+                "task": "i2v",
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "num-input-images": 1,
+                "seed": 42,
+                "enable-negative-prompt": true,
+                "random-request-config": [
+                    {
+                        "width": 832,
+                        "height": 480,
+                        "num_inference_steps": 4,
+                        "num_frames": 81,
+                        "fps": 16,
+                        "weight": 1
+                    }
+                ],
+                "baseline": {
+                    "throughput_qps": 0.0465,
+                    "latency_mean": 21.4354,
+                    "peak_memory_mb_mean": 72557.58
+                }
+            }
+        ]
     },
-    "benchmark_params": [
-      {
-        "name": "832x480_frames81_steps4",
-        "dataset": "random",
-        "task": "i2v",
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "num-input-images": 1,
-        "seed": 42,
-        "enable-negative-prompt": true,
-        "random-request-config": [
-          {
-            "width": 832,
-            "height": 480,
-            "num_inference_steps": 4,
-            "num_frames": 81,
-            "fps": 16,
-            "weight": 1
-          }
-        ],
-        "baseline": {
-          "throughput_qps": 0.0465,
-          "latency_mean": 21.4354,
-          "peak_memory_mb_mean": 72557.58
-        }
-      }
-    ]
-  },
-  {
-    "test_name": "test_wan22_i2v_usp2_vae_patch2_hsdp_slicing",
-    "mark": [
-      {
-        "hardware_marks": {
-          "res": {
-            "cuda": "H100"
-          },
-          "num_cards": 2
-        }
-      },
-      "full_model",
-      "diffusion"
-    ],
-    "description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing",
-    "server_type": "vllm-omni",
-    "server_params": {
-      "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-      "serve_args": {
-        "usp": 2,
-        "vae-patch-parallel-size": 2,
-        "use-hsdp": true,
-        "vae-use-slicing": true,
-        "enable-diffusion-pipeline-profiler": true
-      }
-    },
-    "benchmark_params": [
-      {
-        "name": "832x480_frames81_steps4",
-        "dataset": "random",
-        "task": "i2v",
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "num-input-images": 1,
-        "seed": 42,
-        "enable-negative-prompt": true,
-        "random-request-config": [
-          {
-            "width": 832,
-            "height": 480,
-            "num_inference_steps": 4,
-            "num_frames": 81,
-            "fps": 16,
-            "weight": 1
-          }
-        ],
-        "baseline": {
-          "throughput_qps": 0.0642,
-          "latency_mean": 15.5589,
-          "peak_memory_mb_mean": 45742.88
-        }
-      },
-      {
-        "name": "1280x720_frames121_steps4",
-        "dataset": "random",
-        "task": "i2v",
-        "num-prompts": 10,
-        "max-concurrency": 1,
-        "num-input-images": 1,
-        "seed": 42,
-        "enable-negative-prompt": true,
-        "random-request-config": [
-          {
-            "width": 1280,
-            "height": 720,
-            "num_inference_steps": 4,
-            "num_frames": 121,
-            "fps": 16,
-            "weight": 1
-          }
-        ],
-        "baseline": {
-          "throughput_qps": 0.0128,
-          "latency_mean": 77.8565,
-          "peak_memory_mb_mean": 53947.84
-        }
-      }
-    ]
-  }
+    {
+        "test_name": "test_wan22_i2v_usp2_vae_patch2_hsdp_slicing",
+        "description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing",
+        "server_type": "vllm-omni",
+        "server_params": {
+            "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+            "serve_args": {
+                "usp": 2,
+                "vae-patch-parallel-size": 2,
+                "use-hsdp": true,
+                "vae-use-slicing": true,
+                "enable-diffusion-pipeline-profiler": true
+            }
+        },
+        "benchmark_params": [
+            {
+                "name": "832x480_frames81_steps4",
+                "dataset": "random",
+                "task": "i2v",
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "num-input-images": 1,
+                "seed": 42,
+                "enable-negative-prompt": true,
+                "random-request-config": [
+                    {
+                        "width": 832,
+                        "height": 480,
+                        "num_inference_steps": 4,
+                        "num_frames": 81,
+                        "fps": 16,
+                        "weight": 1
+                    }
+                ],
+                "baseline": {
+                    "throughput_qps": 0.0642,
+                    "latency_mean": 15.5589,
+                    "peak_memory_mb_mean": 45742.88
+                }
+            },
+            {
+                "name": "1280x720_frames121_steps4",
+                "dataset": "random",
+                "task": "i2v",
+                "num-prompts": 10,
+                "max-concurrency": 1,
+                "num-input-images": 1,
+                "seed": 42,
+                "enable-negative-prompt": true,
+                "random-request-config": [
+                    {
+                        "width": 1280,
+                        "height": 720,
+                        "num_inference_steps": 4,
+                        "num_frames": 121,
+                        "fps": 16,
+                        "weight": 1
+                    }
+                ],
+                "baseline": {
+                    "throughput_qps": 0.0128,
+                    "latency_mean": 77.8565,
+                    "peak_memory_mb_mean": 53947.84
+                }
+            }
+        ]
+    }
 ]

tests/diffusion/model_loader/test_modelopt_fp8_adapter.py

@@ -9,7 +9,6 @@ import torch.nn as nn
 
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     ModelOptFp8CheckpointAdapter,
-    ModelOptNvFp4CheckpointAdapter,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -90,23 +89,6 @@ class _QuantizedRemappedModelOptModel(nn.Module):
         return name
 
 
-class _RemappedNvFp4Model(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-        self.runtime = nn.Module()
-        self.runtime.proj = nn.Module()
-        self.runtime.proj.register_parameter(
-            "weight",
-            nn.Parameter(torch.empty(1, 1, dtype=torch.uint8), requires_grad=False),
-        )
-        for name in ("input_scale", "weight_scale", "weight_scale_2"):
-            self.runtime.proj.register_parameter(name, nn.Parameter(torch.empty(1), requires_grad=False))
-
-    @staticmethod
-    def remap_checkpoint_key(name: str) -> str:
-        return name.replace("transformer.orig.proj.", "runtime.proj.")
-
-
 def _make_source() -> SimpleNamespace:
     return SimpleNamespace(
         subfolder="transformer",
@@ -206,36 +188,13 @@ def test_modelopt_adapter_remaps_and_keeps_scales_for_quantized_target():
     emitted = dict(adapted)
     assert len(adapted) == len(emitted) == 3
 
-    # Both scales are retained and emitted under the remapped runtime names.
-    assert torch.equal(emitted["runtime.proj.weight_scale"], weight_scale)
-    assert torch.equal(emitted["runtime.proj.input_scale"], input_scale)
+    # Both scales are retained (the dropped-scale regression). They keep their
+    # source-checkpoint names; the model's load_weights remaps them downstream.
+    assert torch.equal(emitted["transformer.orig.proj.weight_scale"], weight_scale)
+    assert torch.equal(emitted["transformer.orig.proj.input_scale"], input_scale)
 
     # The FP8 weight is passed through unchanged under the remapped target name
     # (the quantized param is FP8, so it must NOT be dequantized). Compare via
     # float32 since torch.equal has no FP8 CPU kernel.
     assert emitted["runtime.proj.weight"].dtype == torch.float8_e4m3fn
     assert torch.equal(emitted["runtime.proj.weight"].to(torch.float32), fp8_weight.to(torch.float32))
-
-
-def test_modelopt_nvfp4_adapter_remaps_quantized_weights_and_scales():
-    adapter = ModelOptNvFp4CheckpointAdapter(_RemappedNvFp4Model(), _make_source())
-    prefix = "transformer.orig.proj"
-    checkpoint_tensors = [
-        (f"{prefix}.input_scale", torch.tensor([1.0])),
-        (f"{prefix}.weight_scale", torch.tensor([0.5])),
-        (f"{prefix}.weight_scale_2", torch.tensor([0.25])),
-        (f"{prefix}.weight", torch.tensor([[1]], dtype=torch.uint8)),
-    ]
-
-    adapted = list(adapter.adapt(iter(checkpoint_tensors)))
-
-    assert [name for name, _ in adapted] == [
-        name.replace(f"{prefix}.", "runtime.proj.") for name, _ in checkpoint_tensors
-    ]
-
-
-def test_modelopt_nvfp4_adapter_rejects_unconsumed_pre_quant_scale():
-    adapter = ModelOptNvFp4CheckpointAdapter(nn.Module(), _make_source())
-
-    with pytest.raises(ValueError, match="does not consume pre_quant_scale"):
-        list(adapter.adapt(iter([("transformer.proj.pre_quant_scale", torch.tensor([1.0]))])))

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -15,25 +15,6 @@ from vllm_omni.platforms import current_omni_platform
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
-@pytest.fixture(autouse=True)
-def _single_rank_tensor_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Provide the TP metadata required by vLLM parallel linear layers."""
-    from vllm.model_executor import parameter
-    from vllm.model_executor.layers import linear
-
-    from vllm_omni.diffusion.models.cosmos3 import (
-        transformer_cosmos3,
-        transformer_cosmos3_edge,
-    )
-
-    monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0)
-    monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1)
-    monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
-    monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
-    monkeypatch.setattr(transformer_cosmos3, "get_tensor_model_parallel_world_size", lambda: 1)
-    monkeypatch.setattr(transformer_cosmos3_edge, "get_tensor_model_parallel_world_size", lambda: 1)
-
-
 @pytest.fixture
 def accelerator_device() -> torch.device:
     """Provide an accelerator device, skipping when none is available."""
@@ -59,22 +40,6 @@ def _tiny_cosmos3_config(**overrides):
     return config
 
 
-def _tiny_cosmos3_edge_config(**overrides):
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
-        COSMOS3_EDGE_BACKBONE_TYPE,
-    )
-
-    config = _tiny_cosmos3_config(
-        latent_patch_size=2,
-        latent_channel=48,
-        temporal_compression_factor=4,
-        backbone_type=COSMOS3_EDGE_BACKBONE_TYPE,
-        qk_norm_for_text=False,
-    )
-    config.update(overrides)
-    return config
-
-
 def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
         compute_mrope_position_ids_action,
@@ -113,16 +78,6 @@ def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
     assert action_offset == 14
 
 
-def test_timestep_embedder_stores_frequencies_in_fp32() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
-        TimestepEmbedder,
-    )
-
-    embedder = TimestepEmbedder(hidden_size=8, frequency_embedding_size=16)
-
-    assert embedder.freqs.dtype == torch.float32
-
-
 @pytest.mark.parametrize(
     ("key", "value"),
     [
@@ -142,245 +97,30 @@ def test_validate_supported_config_rejects_unsupported_flags(key: str, value) ->
     Cosmos3VFMTransformer._validate_supported_config(None)
 
 
-def test_edge_config_resolves_nemotron_defaults() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
-        Cosmos3EdgeVFMTransformer,
-    )
-
-    model = Cosmos3EdgeVFMTransformer(
-        SimpleNamespace(
-            tf_model_config=_tiny_cosmos3_edge_config(
-                layer_norm_epsilon=1e-5,
-                rope_scaling={},
-                rope_parameters={"mrope_section": [24, 20, 20]},
-            ),
-            dtype=torch.float32,
-        )
-    )
-
-    assert model.rms_norm_eps == 1e-5
-    assert model.rope_theta == 100_000_000
-    assert model.mrope_section == [24, 20, 20]
-    assert model.latent_channel_size == 48
-    assert model.latent_patch_size == 2
-    assert model.temporal_compression_factor == 4
-
-
-def test_edge_config_requires_backbone_type() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
-        Cosmos3EdgeVFMTransformer,
-    )
-
-    config = _tiny_cosmos3_edge_config()
-    config.pop("backbone_type")
-
-    with pytest.raises(ValueError, match="must declare backbone_type"):
-        Cosmos3EdgeVFMTransformer(
-            SimpleNamespace(
-                tf_model_config=config,
-                dtype=torch.float32,
-            )
-        )
-
-
-def test_edge_und_and_gen_use_relu2_and_no_und_qk_norm() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
-        Cosmos3EdgeVFMTransformer,
-        Cosmos3Relu2MLP,
-    )
-
-    model = Cosmos3EdgeVFMTransformer(
-        SimpleNamespace(
-            tf_model_config=_tiny_cosmos3_edge_config(
-                num_hidden_layers=1,
-                use_k_norm_und_for_gen=False,
-            ),
-            dtype=torch.float32,
-        )
-    )
-
-    layer = model.language_model.layers[0]
-    gen_layer = model.gen_layers[0]
-    assert isinstance(layer.mlp, Cosmos3Relu2MLP)
-    assert not hasattr(layer.mlp, "gate_proj")
-    assert isinstance(gen_layer.mlp, Cosmos3Relu2MLP)
-    assert not hasattr(gen_layer.mlp, "gate_proj")
-    assert not hasattr(layer.self_attn, "norm_q")
-    assert not hasattr(layer.self_attn, "norm_k")
-    assert layer.self_attn.k_norm_und_for_gen is None
-
-
-def test_edge_creates_gen_k_norm_when_enabled() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
-        Cosmos3EdgeVFMTransformer,
-    )
-
-    model = Cosmos3EdgeVFMTransformer(
-        SimpleNamespace(
-            tf_model_config=_tiny_cosmos3_edge_config(
-                num_hidden_layers=1,
-                qk_norm_for_diffusion=True,
-                use_k_norm_und_for_gen=True,
-            ),
-            dtype=torch.float32,
-        )
-    )
-
-    k_norm = model.language_model.layers[0].self_attn.k_norm_und_for_gen
-    assert model.use_k_norm_und_for_gen is True
-    assert isinstance(k_norm, RMSNorm)
-    assert isinstance(model.gen_layers[0].cross_attention.norm_q, RMSNorm)
-    assert isinstance(model.gen_layers[0].cross_attention.norm_k, RMSNorm)
-
-
-def test_edge_supports_no_gen_qk_norm_variant() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
-        Cosmos3EdgeVFMTransformer,
-    )
-
-    model = Cosmos3EdgeVFMTransformer(
-        SimpleNamespace(
-            tf_model_config=_tiny_cosmos3_edge_config(
-                num_hidden_layers=1,
-                qk_norm_for_diffusion=False,
-                use_k_norm_und_for_gen=False,
-            ),
-            dtype=torch.float32,
-        )
-    )
-
-    assert model.qk_norm_for_diffusion is False
-    assert model.language_model.layers[0].self_attn.k_norm_und_for_gen is None
-    assert model.gen_layers[0].cross_attention.qk_norm is False
-    assert not hasattr(model.gen_layers[0].cross_attention, "norm_q")
-    assert not hasattr(model.gen_layers[0].cross_attention, "norm_k")
-
-
-def test_edge_validates_required_relu2_weights() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
-        Cosmos3EdgeVFMTransformer,
-    )
-
-    model = object.__new__(Cosmos3EdgeVFMTransformer)
-    nn.Module.__init__(model)
-    model.num_hidden_layers = 1
-    model.qk_norm_for_diffusion = True
-    model.use_k_norm_und_for_gen = True
-
-    complete = {
-        "transformer.language_model.layers.0.mlp.up_proj.weight",
-        "transformer.language_model.layers.0.mlp.down_proj.weight",
-        "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight",
-        "transformer.gen_layers.0.mlp.up_proj.weight",
-        "transformer.gen_layers.0.mlp.down_proj.weight",
-    }
-    model.validate_loaded_weights(complete)
-
-    mlp_weights = complete - {
-        "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight",
-    }
-    for missing_weight in mlp_weights:
-        with pytest.raises(ValueError, match="missing required weights"):
-            model.validate_loaded_weights(complete - {missing_weight})
-
-    missing_k_norm = complete - {"transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight"}
-    with pytest.raises(ValueError, match=r"self_attn\.k_norm_und_for_gen"):
-        model.validate_loaded_weights(missing_k_norm)
-
-    model.use_k_norm_und_for_gen = False
-    model.validate_loaded_weights(missing_k_norm)
-
-
-def test_edge_gen_cached_k_is_normalized_but_reasoner_uses_raw_k() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
-        Cosmos3EdgeCausalAttention,
-    )
-
-    class ScaleProjection(nn.Module):
-        def __init__(self, scale: float) -> None:
-            super().__init__()
-            self.scale = scale
-
-        def forward(self, x: torch.Tensor) -> torch.Tensor:
-            return x * self.scale
-
-    class IdentityOutput(nn.Module):
-        def forward(self, x: torch.Tensor) -> torch.Tensor:
-            return x
-
-    class RawKAttention(nn.Module):
-        def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-            del v
-            return q + k
-
-    attn = object.__new__(Cosmos3EdgeCausalAttention)
-    nn.Module.__init__(attn)
-    attn.hidden_size = 4
-    attn.num_heads = 2
-    attn.num_kv_heads = 2
-    attn.head_dim = 2
-    attn.num_heads_local = 2
-    attn.num_kv_heads_local = 2
-    attn.to_q = ScaleProjection(1.0)
-    attn.to_k = ScaleProjection(1.0)
-    attn.to_v = ScaleProjection(1.0)
-    attn.to_out = IdentityOutput()
-    attn.k_norm_und_for_gen = RMSNorm(2, eps=1e-5)
-    attn.attn = RawKAttention()
-
-    hidden_states = torch.tensor(
-        [[[1.0, 2.0, 3.0, 4.0], [2.0, 1.0, 4.0, 3.0]]],
-        dtype=torch.float32,
-    )
-    freqs_cos = torch.ones(1, 2, 1, 2)
-    freqs_sin = torch.zeros(1, 2, 1, 2)
-
-    reasoner_out, gen_k, gen_v = attn(hidden_states, freqs_cos, freqs_sin)
-    attn.to_k.scale = 3.0
-    scaled_reasoner_out, scaled_gen_k, scaled_gen_v = attn(hidden_states, freqs_cos, freqs_sin)
-
-    torch.testing.assert_close(gen_k, scaled_gen_k)
-    torch.testing.assert_close(gen_v, scaled_gen_v)
-    assert not torch.allclose(reasoner_out, scaled_reasoner_out)
-
-
 def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3LanguageModel, Cosmos3VFMTransformer
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import (
-        Cosmos3EdgeLanguageModel,
-        Cosmos3EdgeVFMTransformer,
-    )
 
-    for transformer_cls, language_model_cls in (
-        (Cosmos3VFMTransformer, Cosmos3LanguageModel),
-        (Cosmos3EdgeVFMTransformer, Cosmos3EdgeLanguageModel),
-    ):
-        model = object.__new__(transformer_cls)
-        nn.Module.__init__(model)
-        model.language_model = nn.Module()
-        model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)])
-        model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
-        model.norm_moe_gen = nn.LayerNorm(2)
+    model = object.__new__(Cosmos3VFMTransformer)
+    nn.Module.__init__(model)
+    model.language_model = nn.Module()
+    model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)])
+    model.gen_layers = nn.ModuleList([nn.Linear(2, 2)])
+    model.norm_moe_gen = nn.LayerNorm(2)
 
-        matched = [
-            name
-            for name, module in model.named_modules()
-            if any(condition(name, module) for condition in model._hsdp_shard_conditions)
-        ]
-        assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"]
-        assert transformer_cls._layerwise_offload_blocks_attrs == ["gen_layers"]
-        assert language_model_cls._layerwise_offload_blocks_attrs == ["layers"]
-        assert transformer_cls._repeated_blocks == ["Cosmos3GenDecoderLayer"]
+    matched = [
+        name
+        for name, module in model.named_modules()
+        if any(condition(name, module) for condition in model._hsdp_shard_conditions)
+    ]
+    assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"]
+    assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"]
+    assert Cosmos3LanguageModel._layerwise_offload_blocks_attrs == ["layers"]
+    assert Cosmos3VFMTransformer._repeated_blocks == ["Cosmos3GenDecoderLayer"]
 
-        model.latent_patch_size = 2
-        model.latent_channel_size = 3
-        latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5)
-        torch.testing.assert_close(
-            model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5),
-            latents,
-        )
+    model.latent_patch_size = 2
+    model.latent_channel_size = 3
+    latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5)
+    torch.testing.assert_close(model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5), latents)
 
 
 def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -402,17 +142,6 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
-def test_qwen_gen_mlp_remains_gated() -> None:
-    from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer
-
-    model = Cosmos3VFMTransformer(
-        SimpleNamespace(tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1), dtype=torch.float32)
-    )
-
-    assert isinstance(model.gen_layers[0].mlp, Cosmos3GatedMLP)
-    assert hasattr(model.gen_layers[0].mlp, "gate_proj")
-
-
 def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 

tests/e2e/accuracy/test_cosmos3_similarity.py

@@ -42,6 +42,8 @@ def _server_args() -> list[str]:
     return [
         "--num-gpus",
         "1",
+        "--model-class-name",
+        "Cosmos3OmniDiffusersPipeline",
         "--stage-init-timeout",
         "900",
         "--init-timeout",

tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

@@ -48,7 +48,6 @@ _DEPLOY_CONFIG = {
             "trust_remote_code": True,
             "devices": "0,1,2,3",  # set dynamically by _write_deploy_config
             "vae_use_slicing": False,
-            "moe_backend": "flashinfer_cutlass",
             "vae_use_tiling": False,
             "parallel_config": {
                 "pipeline_parallel_size": 1,

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -34,6 +34,7 @@ def _make_minimal_talker(
     """
     model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
     model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16)
+    model._embedding_dtype = torch.bfloat16
     if tts_pad_embed is None:
         tts_pad_embed = torch.zeros((1, 4), dtype=torch.bfloat16)
     model._tts_pad_embed = tts_pad_embed
@@ -41,6 +42,7 @@ def _make_minimal_talker(
     def _default_raise(**_kwargs):
         raise AssertionError("build_prompt_embeds was not stubbed in this test")
 
+    model._embedding_dtype = torch.bfloat16
     model._prompt_builder = SimpleNamespace(
         build_prompt_embeds=build_prompt_embeds if build_prompt_embeds is not None else _default_raise,
     )
@@ -81,6 +83,7 @@ def _make_minimal_builder(
     builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
     builder._residual_code_embeddings = lambda: []
     builder._speaker_encoder = None
+    builder._embedding_dtype = torch.bfloat16
     builder._tts_pad_embed_buffer = (
         tts_pad_embed if tts_pad_embed is not None else torch.zeros((1, 4), dtype=torch.bfloat16)
     )
@@ -89,6 +92,7 @@ def _make_minimal_builder(
     )
     builder._speaker_cache = None
     builder._text_tokenizer = None
+    builder._embedding_dtype = torch.bfloat16
     builder._ref_audio_artifact_cache_max_entries = 256
     builder._ref_audio_artifact_cache = OrderedDict()
     builder._resampler_cache = OrderedDict()

tests/model_extras/test_model_extras.py

@@ -60,9 +60,8 @@ def test_sensenova_extra_registry_declares_request_and_response_params() -> None
 
 @pytest.mark.core_model
 @pytest.mark.cpu
-@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
-def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_name: str) -> None:
-    assert get_extra_body_params(pipeline_name) == frozenset(
+def test_cosmos3_extra_registry_declares_request_and_response_params() -> None:
+    assert get_extra_body_params("Cosmos3OmniDiffusersPipeline") == frozenset(
         {
             "flow_shift",
             "max_sequence_length",
@@ -99,7 +98,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na
             "session_id",
         }
     )
-    assert get_extra_output_params(pipeline_name) == frozenset(
+    assert get_extra_output_params("Cosmos3OmniDiffusersPipeline") == frozenset(
         {
             "action",
             "raw_action_dim",
@@ -107,7 +106,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na
             "action_mode",
         }
     )
-    assert should_init_extra_args_for_non_diffusion_stages(pipeline_name) is False
+    assert should_init_extra_args_for_non_diffusion_stages("Cosmos3OmniDiffusersPipeline") is False
 
 
 @pytest.mark.core_model
@@ -225,10 +224,9 @@ def test_ming_flash_omni_image_to_image_prompt_builder() -> None:
 
 @pytest.mark.core_model
 @pytest.mark.cpu
-@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"])
-def test_cosmos3_text_to_image_prompt_builder_selects_image_modality(pipeline_name: str) -> None:
+def test_cosmos3_text_to_image_prompt_builder_selects_image_modality() -> None:
     assert build_text_to_image_prompt(
-        pipeline_name,
+        "Cosmos3OmniDiffusersPipeline",
         prompt="a red sports car at golden hour",
         negative_prompt="blurry, distorted",
         height=1024,
@@ -239,7 +237,7 @@ def test_cosmos3_text_to_image_prompt_builder_selects_image_modality(pipeline_na
         "negative_prompt": "blurry, distorted",
     }
     assert build_text_to_image_prompt(
-        pipeline_name,
+        "Cosmos3OmniDiffusersPipeline",
         prompt="a red sports car",
         negative_prompt=None,
     ) == {"prompt": "a red sports car", "modalities": ["image"]}

tests/model_tests/diffusion/test_alignment.py

@@ -71,7 +71,6 @@ EXCLUDED_MODELS = [
     "OmniVoicePipeline",
     "OmniVoice",
     "Cosmos3OmniDiffusersPipeline",
-    "Cosmos3OmniPipeline",
     "DiffusersAdapterPipeline",
     "HiDreamImagePipeline",
     "DreamZeroPipeline",

vllm_omni/config/omni_config.py

@@ -469,7 +469,6 @@ class _DiffusionConfigProjection:
     diffusion_kv_cache_skip_layers: str | list[int] | tuple[int, ...] | set[int] | None = None
     diffusion_kv_cache_skip_step_indices: set[int] | None = None
     diffusion_kv_cache_skip_layer_indices: set[int] | None = None
-    moe_backend: str = "auto"
     force_cutlass_fp8: bool = False
     enable_diffusion_pipeline_profiler: bool = False
     step_execution: bool = False

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import time
 from collections.abc import Iterable
 from typing import Any
@@ -19,12 +20,21 @@ _STATS_INTERVAL_S = 1.0
 # (the state ``OmniSchedulingCoordinator`` records via ``_waiting_since``)
 # before the scheduler force-fails it.  Defends against stuck consumer-side
 # requests when the producer drops a full-payload, send fails, or recv
-# never arrives.  A value <= 0 disables the safety net.
+# never arrives.  Override per-deployment via
+# VLLM_OMNI_INPUT_WAIT_TIMEOUT_S; set <=0 to disable the safety net.
 #
 # Scope: this constant only covers the full-payload coordinator path
 # (``input_coordinator``).  The async-chunk path uses
 # ``chunk_transfer_adapter`` and is not affected by this constant.
-DEFAULT_INPUT_WAIT_TIMEOUT_S: float = 600.0
+_INPUT_WAIT_TIMEOUT_RAW = os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600")
+try:
+    DEFAULT_INPUT_WAIT_TIMEOUT_S: float = float(_INPUT_WAIT_TIMEOUT_RAW)
+except ValueError:
+    logger.warning(
+        "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to 600 seconds.",
+        _INPUT_WAIT_TIMEOUT_RAW,
+    )
+    DEFAULT_INPUT_WAIT_TIMEOUT_S = 600.0
 
 
 class OmniSchedulerMixin:

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -1016,7 +1016,6 @@ CUSTOM_DIT_ENABLERS.update(
         "Wan22VACEPipeline": enable_cache_for_wan22,
         "Wan22S2VPipeline": enable_cache_for_wan22_s2v,
         "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3,
-        "Cosmos3OmniPipeline": enable_cache_for_cosmos3,
         "Krea2Pipeline": enable_cache_for_krea2,
     }
 )

vllm_omni/diffusion/data.py

@@ -664,9 +664,6 @@ class OmniDiffusionConfig:
     # STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
     skip_time_steps: int = 15
 
-    # MoE kernel backend selection
-    moe_backend: str = "auto"
-
     # Compilation
     enforce_eager: bool = False
 

vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py

@@ -11,12 +11,7 @@ from vllm.model_executor.utils import get_packed_modules_mapping
 
 logger = init_logger(__name__)
 
-MODEL_OPT_SCALE_SUFFIXES = (
-    ".input_scale",
-    ".weight_scale",
-    ".weight_scale_2",
-    ".weight_scale_inv",
-)
+MODEL_OPT_SCALE_SUFFIXES = (".input_scale", ".weight_scale", ".weight_scale_2", ".weight_scale_inv")
 DEFAULT_PACKED_MODULES_MAPPING = {
     "to_qkv": ("to_q", "to_k", "to_v"),
     "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"),
@@ -127,7 +122,6 @@ class ModelOptFp8CheckpointAdapter:
 
         for candidate in self._weights_mapper.apply_list([name]):
             if candidate != name and candidate in self._loadable_tensors:
-                # Preserve shard names so downstream packed-weight loaders can route them.
                 return candidate, name
         return None, name
 
@@ -184,7 +178,6 @@ class ModelOptFp8CheckpointAdapter:
     def _handle_scale_tensor(
         self,
         name: str,
-        output_name: str,
         tensor: torch.Tensor,
         target_name: str | None,
         state: _AdaptState,
@@ -193,7 +186,7 @@ class ModelOptFp8CheckpointAdapter:
         if target_name is None:
             state.skipped_scales += 1
         else:
-            yield output_name, tensor
+            yield name, tensor
         yield from self._flush_pending_weights(name, state)
 
     def _target_dtype_for_dequantization(
@@ -256,7 +249,7 @@ class ModelOptFp8CheckpointAdapter:
         for name, tensor in weights:
             target_name, output_name = self._resolve_target_and_output_names(name)
             if self._is_scale(name):
-                yield from self._handle_scale_tensor(name, output_name, tensor, target_name, state)
+                yield from self._handle_scale_tensor(name, tensor, target_name, state)
                 continue
 
             target_dtype = self._target_dtype_for_dequantization(tensor, target_name)
@@ -271,8 +264,6 @@ class ModelOptFp8CheckpointAdapter:
 
 
 class ModelOptNvFp4CheckpointAdapter(ModelOptFp8CheckpointAdapter):
-    _PRE_QUANT_SCALE_SUFFIX = ".pre_quant_scale"
-
     @staticmethod
     def _is_checkpoint_quant_config(quant_config: object | None) -> bool:
         return (
@@ -282,22 +273,6 @@ class ModelOptNvFp4CheckpointAdapter(ModelOptFp8CheckpointAdapter):
             and bool(getattr(quant_config, "is_checkpoint_nvfp4_serialized", False))
         )
 
-    def adapt(
-        self,
-        weights: Iterable[tuple[str, torch.Tensor]],
-    ) -> Generator[tuple[str, torch.Tensor], None, None]:
-        def validated_weights() -> Generator[tuple[str, torch.Tensor], None, None]:
-            for name, tensor in weights:
-                if name.endswith(self._PRE_QUANT_SCALE_SUFFIX):
-                    raise ValueError(
-                        f"ModelOpt NVFP4 checkpoint tensor {name!r} is unsupported: "
-                        "vLLM 0.25.0 does not consume pre_quant_scale. Export the checkpoint "
-                        "with pre-quant scales folded into the weights."
-                    )
-                yield name, tensor
-
-        yield from super().adapt(validated_weights())
-
 
 class ModelOptMixedPrecisionCheckpointAdapter(ModelOptFp8CheckpointAdapter):
     @staticmethod

vllm_omni/diffusion/models/cosmos3/__init__.py

@@ -7,12 +7,10 @@ from .pipeline_cosmos3 import (
     get_cosmos3_pre_process_func,
 )
 from .transformer_cosmos3 import Cosmos3VFMTransformer
-from .transformer_cosmos3_edge import Cosmos3EdgeVFMTransformer
 
 __all__ = [
     "Cosmos3OmniDiffusersPipeline",
     "get_cosmos3_post_process_func",
     "get_cosmos3_pre_process_func",
     "Cosmos3VFMTransformer",
-    "Cosmos3EdgeVFMTransformer",
 ]

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -29,7 +29,6 @@ is produced from sound latents rather than from ``multi_modal_data["audio"]``.
 
 from __future__ import annotations
 
-import json
 import math
 import os
 import time
@@ -40,6 +39,7 @@ from typing import Any, ClassVar
 import numpy as np
 import PIL.Image
 import torch
+from diffusers import UniPCMultistepScheduler
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
 from torch import nn
@@ -61,12 +61,6 @@ from vllm_omni.diffusion.models.interface import (
     SupportsComponentDiscovery,
 )
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
-from vllm_omni.diffusion.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler,
-)
-from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import (
-    FlowUniPCMultistepScheduler,
-)
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -100,8 +94,7 @@ from .transfer import (
     transfer_max_frames_from_extra_args,
     uint8_cthw_to_normalized_5d,
 )
-from .transformer_cosmos3 import Cosmos3VFMTransformer, _tf_config_get, resolve_sound_gen
-from .transformer_cosmos3_edge import COSMOS3_EDGE_BACKBONE_TYPE, Cosmos3EdgeVFMTransformer
+from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen
 from .utils import (
     COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION,
     COSMOS3_VAE_TEMPORAL_COMPRESSION,
@@ -161,7 +154,6 @@ COSMOS3_T2V_DEFAULT_WIDTH = 1280
 COSMOS3_T2V_DEFAULT_NUM_FRAMES = 189
 COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS = 35
 COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE = 6.0
-COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT = 10.0
 
 COSMOS3_T2I_DEFAULT_HEIGHT = 1024
 COSMOS3_T2I_DEFAULT_WIDTH = 1024
@@ -170,73 +162,12 @@ COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE = 7.0
 COSMOS3_T2I_DEFAULT_FLOW_SHIFT = 3.0
 COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0)
 
-COSMOS3_EDGE_T2V_DEFAULT_HEIGHT = 480
-COSMOS3_EDGE_T2V_DEFAULT_WIDTH = 832
-COSMOS3_EDGE_T2V_DEFAULT_GUIDANCE_SCALE = 5.0
-COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT = 3.0
-COSMOS3_EDGE_T2I_DEFAULT_HEIGHT = 640
-COSMOS3_EDGE_T2I_DEFAULT_WIDTH = 640
-
-COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS = "FlowMatchEulerDiscreteScheduler"
-
 # Truncation cap on the prompt token count (shared by T2I and T2V).  Prompts
 # are tokenized to their natural length (no padding); this only bounds the
 # UND pathway / GEN cross-attention cost for pathologically long prompts.
 COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096
 
 
-def _ceil_video_num_frames(num_frames: int, temporal_compression_factor: int) -> int:
-    """Round a video length up to the causal VAE's ``factor * k + 1`` grid."""
-    if temporal_compression_factor <= 0:
-        raise ValueError(f"Cosmos3 temporal_compression_factor must be positive, got {temporal_compression_factor}.")
-    if num_frames <= 1:
-        return num_frames
-    latent_frames = math.ceil((num_frames - 1) / temporal_compression_factor)
-    return latent_frames * temporal_compression_factor + 1
-
-
-def _format_json_object_prompt(
-    prompt: str,
-    *,
-    num_frames: int,
-    frame_rate: float,
-    height: int,
-    width: int,
-    aspect_ratio: str | None,
-) -> str | None:
-    """Match Imaginaire4 metadata injection for JSON-object prompts."""
-    try:
-        prompt_obj = json.loads(prompt)
-    except (json.JSONDecodeError, TypeError, ValueError):
-        return None
-    if not isinstance(prompt_obj, dict):
-        return None
-
-    metadata: dict[str, Any] = {}
-    if num_frames > 1:
-        duration_seconds = int(num_frames / frame_rate) if frame_rate > 0 else 0
-        metadata.update({"duration": f"{duration_seconds}s", "fps": float(frame_rate)})
-    else:
-        prompt_obj.pop("duration", None)
-        prompt_obj.pop("fps", None)
-    metadata["resolution"] = {"H": int(height), "W": int(width)}
-    if aspect_ratio is not None:
-        metadata["aspect_ratio"] = aspect_ratio
-
-    prompt_obj.update(metadata)
-    return json.dumps(prompt_obj)
-
-
-def resolve_cosmos3_transformer_cls(model_config: Any) -> type[Cosmos3VFMTransformer]:
-    """Select the Cosmos3 transformer implementation from transformer/config.json."""
-    backbone_type = _tf_config_get(model_config, "backbone_type", None)
-    if backbone_type is None:
-        return Cosmos3VFMTransformer
-    if backbone_type == COSMOS3_EDGE_BACKBONE_TYPE:
-        return Cosmos3EdgeVFMTransformer
-    raise ValueError(f"Unsupported Cosmos3 transformer backbone_type={backbone_type!r}.")
-
-
 # ---------------------------------------------------------------------------
 # Post-process function (registered in registry.py)
 # ---------------------------------------------------------------------------
@@ -259,9 +190,6 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
     from .guardrails import check_text_safety, ensure_initialized, is_guardrails_enabled
 
     video_processor = VideoProcessor(vae_scale_factor=16)
-    is_edge_model = (
-        _tf_config_get(getattr(od_config, "tf_model_config", None), "backbone_type", None) == COSMOS3_EDGE_BACKBONE_TYPE
-    )
     # Eager-load guardrail models at pipeline build time when the server-level
     # gate is on. Per-request overrides only decide whether the loaded models
     # are *invoked* — they cannot turn checks on without a server-side preload.
@@ -482,17 +410,12 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
         extra = _extra_args(request)
         transfer_requested = action_mode is None and has_transfer_hints(extra)
 
-        # Resolve missing H/W.
+        # Auto-calculate H/W from aspect ratio (720p max area)
         if transfer_requested:
             _set_transfer_size_from_image(request, image)
         elif request.sampling_params.height is None or request.sampling_params.width is None:
             if action_mode is not None:
                 _set_action_size_from_image(request, image)
-            elif is_edge_model:
-                if request.sampling_params.height is None:
-                    request.sampling_params.height = COSMOS3_EDGE_T2V_DEFAULT_HEIGHT
-                if request.sampling_params.width is None:
-                    request.sampling_params.width = COSMOS3_EDGE_T2V_DEFAULT_WIDTH
             else:
                 max_area = 720 * 1280
                 aspect_ratio = image.height / image.width
@@ -900,52 +823,27 @@ class Cosmos3OmniDiffusersPipeline(
             sound_latent_fps = self._sound_tokenizer.latent_fps
 
         # --- Transformer (weights loaded later via weights_sources) ---
-        transformer_cls = resolve_cosmos3_transformer_cls(od_config.tf_model_config)
-        self.transformer = transformer_cls(
+        self.transformer = Cosmos3VFMTransformer(
             od_config=od_config,
             temporal_compression_factor=self.vae_scale_factor_temporal,
             sound_gen=sound_gen,
             sound_dim=sound_dim,
             sound_latent_fps=sound_latent_fps,
         )
-        self.is_edge_model = transformer_cls is Cosmos3EdgeVFMTransformer
 
         # --- Scheduler ---
-        # Distilled model differs from regular one only by scheduler,
-        # distilled one uses FlowMatchEulerDiscreteScheduler, while
-        # regular should use FlowUniPCMultistepScheduler
-
-        scheduler_config = FlowUniPCMultistepScheduler.load_config(
+        # Load from checkpoint to preserve solver_order, timestep_spacing,
+        # beta_schedule, sigma bounds, flow_shift, etc. Only override
+        # flow_shift when explicitly requested by the user.
+        self.scheduler = UniPCMultistepScheduler.from_pretrained(
             model_path,
             subfolder="scheduler",
             local_files_only=local_files_only,
         )
-
-        scheduler_class_name = scheduler_config.get("_class_name")
-
-        self.is_distilled_model = False
-        if scheduler_class_name == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS:
-            fixed_step_config = scheduler_config.get("fixed_step_sampler_config")
-            if not isinstance(fixed_step_config, dict) or fixed_step_config.get("sample_type") != "sde":
-                raise ValueError("Cosmos3 distilled scheduler requires fixed_step_sampler_config.sample_type=sde.")
-            t_list = fixed_step_config.get("t_list")
-            if not isinstance(t_list, list) or not t_list:
-                raise ValueError("Cosmos3 distilled scheduler requires a non-empty fixed_step_sampler_config.t_list.")
-            self.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
-                scheduler_config,
-                stochastic_sampling=True,
-            )
-            self._scheduler_init_t_list = list(t_list)
-            self.is_distilled_model = True
-        else:
-            # Preserve compatible solver settings from the checkpoint, but keep
-            # the base shift neutral. The concrete request shift is applied when
-            # FlowUniPC builds its timesteps.
-            self.scheduler = FlowUniPCMultistepScheduler.from_config(
-                scheduler_config,
-                shift=1.0,
-                use_dynamic_shifting=False,
-                prediction_type="flow_prediction",
+        if od_config.flow_shift is not None:
+            self.scheduler = UniPCMultistepScheduler.from_config(
+                self.scheduler.config,
+                flow_shift=od_config.flow_shift,
             )
         self._cpu_scheduler_state()
 
@@ -964,16 +862,14 @@ class Cosmos3OmniDiffusersPipeline(
             ),
         ]
 
-        # An engine-level override becomes the video default; otherwise use
-        # the matching regular/Edge default. Per-request values still take
-        # precedence.
-        default_video_flow_shift = (
-            COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT if self.is_edge_model else COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT
-        )
-        self._engine_init_flow_shift = float(
-            od_config.flow_shift if od_config.flow_shift is not None else default_video_flow_shift
-        )
+        # Snapshot the loaded scheduler config so we can rebuild the
+        # scheduler at request time when a per-request flow_shift override
+        # is supplied (T2I uses shift=3.0; T2V/I2V use the engine default).
+        self._base_scheduler_config = self.scheduler.config
+        self._engine_init_flow_shift = float(getattr(self.scheduler.config, "flow_shift", 1.0) or 1.0)
         self._current_flow_shift = self._engine_init_flow_shift
+        self._base_scheduler_use_karras_sigmas = self._scheduler_use_karras_sigmas(self.scheduler.config)
+        self._current_scheduler_use_karras_sigmas = self._base_scheduler_use_karras_sigmas
 
         self._guidance_scale = None
         self._num_timesteps = None
@@ -1093,7 +989,6 @@ class Cosmos3OmniDiffusersPipeline(
             "self_attn.to_out.": f"{und_lp}.self_attn.to_out.",
             "self_attn.norm_q.": f"{und_lp}.self_attn.norm_q.",
             "self_attn.norm_k.": f"{und_lp}.self_attn.norm_k.",
-            "self_attn.k_norm_und_for_gen.": f"{und_lp}.self_attn.k_norm_und_for_gen.",
             # GEN attention
             "self_attn.add_q_proj.": f"{gen_lp}.cross_attention.to_q.",
             "self_attn.add_k_proj.": f"{gen_lp}.cross_attention.to_k.",
@@ -1159,7 +1054,6 @@ class Cosmos3OmniDiffusersPipeline(
         loaded = loader.load_weights(_remapped_weights())
         self.transformer.post_load_weights()
         self.transformer.eval()
-        self.transformer.validate_loaded_weights(loaded)
         if getattr(self.transformer, "sound_gen", False):
             sound_markers = ("audio_proj_in.", "audio_proj_out.", "audio_modality_embed")
             missing = [marker.rstrip(".") for marker in sound_markers if not any(marker in name for name in loaded)]
@@ -1465,8 +1359,6 @@ class Cosmos3OmniDiffusersPipeline(
         inputs: RoboLabPolicyInputs,
         pipeline_start: float,
     ) -> DiffusionOutput:
-        if self.is_distilled_model:
-            raise self._distilled_unsupported_error("RoboLab/action policy requests are unsupported.")
         if not getattr(self.transformer, "action_gen", False):
             raise ValueError(
                 "Cosmos3 RoboLab policy serving was requested, but the transformer "
@@ -1577,7 +1469,6 @@ class Cosmos3OmniDiffusersPipeline(
             guidance_interval=None,
             raw_action_dim=raw_action_dim,
             scheduler=scheduler,
-            generator=generator,
         )
 
         if _is_rank_zero():
@@ -1669,34 +1560,38 @@ class Cosmos3OmniDiffusersPipeline(
             raise ValueError(f"Incorrect modality value in {modalities}, expected one of {accepted_modalities}.")
         return "image" in modalities
 
-    def _set_timesteps(self, num_inference_steps: int, device: str | torch.device, shift: float) -> None:
-        if self.is_distilled_model:
-            self.scheduler.set_timesteps(sigmas=self._scheduler_init_t_list, device=device)
-        else:
-            self.scheduler.set_timesteps(
-                num_inference_steps,
-                device=device,
-                shift=shift,
-            )
-
-    def _set_flow_shift(self, target_shift: float) -> None:
-        """Select the shift applied by FlowUniPC when timesteps are built."""
-        self._current_flow_shift = float(target_shift)
-
     @staticmethod
-    def _resolve_seed(sp: OmniDiffusionSamplingParams, generator: torch.Generator | None, default: int = 42) -> int:
-        if sp.seed is not None:
-            return int(sp.seed)
-        if isinstance(generator, torch.Generator):
-            return int(generator.initial_seed())
-        return int(default)
+    def _scheduler_use_karras_sigmas(config: Any) -> bool | None:
+        value = getattr(config, "use_karras_sigmas", None)
+        return None if value is None else bool(value)
 
-    def _resolve_guidance_scale(self, sp: OmniDiffusionSamplingParams, default: float) -> float:
-        if self.is_distilled_model:
-            return 1.0
-        if sp.guidance_scale_provided:
-            return float(sp.guidance_scale)
-        return float(default)
+    def _set_flow_shift(self, target_shift: float, *, use_karras_sigmas: bool | None = None) -> None:
+        """Set UniPC scheduler mode for a concrete request.
+
+        The scheduler is rebuilt from the saved base config if
+        the target differs from the current shift or Karras-sigma mode.
+        Tracking explicit scheduler state is required because the previous
+        mode may have rebuilt the scheduler - we cannot rely on
+        ``self.scheduler.config`` reflecting the last requested target if a
+        rebuild was skipped via the equality check.
+        """
+        target = float(target_shift)
+        target_use_karras_sigmas = (
+            self._base_scheduler_use_karras_sigmas if use_karras_sigmas is None else bool(use_karras_sigmas)
+        )
+        if (
+            target == float(self._current_flow_shift)
+            and target_use_karras_sigmas == self._current_scheduler_use_karras_sigmas
+        ):
+            return
+
+        scheduler_kwargs: dict[str, Any] = {"flow_shift": target}
+        if use_karras_sigmas is not None:
+            scheduler_kwargs["use_karras_sigmas"] = bool(use_karras_sigmas)
+        self.scheduler = UniPCMultistepScheduler.from_config(self._base_scheduler_config, **scheduler_kwargs)
+        self._cpu_scheduler_state()
+        self._current_flow_shift = target
+        self._current_scheduler_use_karras_sigmas = self._scheduler_use_karras_sigmas(self.scheduler.config)
 
     def _cpu_scheduler_state(self) -> None:
         # We need to move scheduler tensors to CPU, as unipc from diffusers assumes they are on CPU.
@@ -1717,53 +1612,6 @@ class Cosmos3OmniDiffusersPipeline(
     def num_timesteps(self):
         return self._num_timesteps
 
-    @staticmethod
-    def _distilled_unsupported_error(detail: str) -> ValueError:
-        return ValueError(
-            f"Cosmos3 distilled checkpoints support only text-to-image and image-to-video generation; {detail}"
-        )
-
-    def _validate_distilled_generation_mode(
-        self,
-        *,
-        is_t2i: bool,
-        image_tensor: torch.Tensor | None,
-        action_enabled: bool,
-        transfer_config: Cosmos3TransferConfig | None,
-        is_v2v: bool,
-        sound_enabled: bool,
-    ) -> None:
-        if not self.is_distilled_model:
-            return
-        if action_enabled:
-            raise self._distilled_unsupported_error("action requests are unsupported.")
-        if transfer_config is not None:
-            raise self._distilled_unsupported_error("transfer requests are unsupported.")
-        if is_v2v:
-            raise self._distilled_unsupported_error("video-to-video requests are unsupported.")
-        if sound_enabled:
-            raise self._distilled_unsupported_error("sound generation is unsupported.")
-        if not is_t2i and image_tensor is None:
-            raise self._distilled_unsupported_error("text-to-video requests are unsupported.")
-        if is_t2i and image_tensor is not None:
-            raise self._distilled_unsupported_error("image-conditioned image generation is unsupported.")
-
-    def _validate_edge_generation_mode(
-        self,
-        *,
-        transfer_config: Cosmos3TransferConfig | None,
-        is_v2v: bool,
-        sound_enabled: bool,
-    ) -> None:
-        if not self.is_edge_model:
-            return
-        if sound_enabled:
-            raise ValueError("Cosmos3 Edge checkpoints do not support sound generation.")
-        if transfer_config is not None:
-            raise ValueError("Cosmos3 Edge checkpoints do not support transfer inference.")
-        if is_v2v:
-            raise ValueError("Cosmos3 Edge checkpoints do not support video-to-video generation.")
-
     # -- Prompt formatting -----------------------------------------------------
 
     @staticmethod
@@ -1780,10 +1628,6 @@ class Cosmos3OmniDiffusersPipeline(
         """
         Append duration and resolution metadata to a prompt.
         """
-        prompt = prompt.strip()
-        if duration_template is None and resolution_template is None:
-            return prompt
-
         parts: list[str] = []
         head = prompt.rstrip(".").strip()
         if head:
@@ -2059,26 +1903,15 @@ class Cosmos3OmniDiffusersPipeline(
             res_tmpl = COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_RESOLUTION_TEMPLATE
         else:
             res_tmpl = None
-        json_prompt = _format_json_object_prompt(
+        prompt = self._apply_metadata_templates(
             prompt,
-            num_frames=num_frames,
-            frame_rate=frame_rate,
-            height=height,
-            width=width,
-            aspect_ratio=self._get_sp_param(sp, "aspect_ratio", None),
+            num_frames,
+            frame_rate,
+            height,
+            width,
+            duration_template=dur_tmpl,
+            resolution_template=res_tmpl,
         )
-        if json_prompt is not None:
-            prompt = json_prompt
-        else:
-            prompt = self._apply_metadata_templates(
-                prompt,
-                num_frames,
-                frame_rate,
-                height,
-                width,
-                duration_template=dur_tmpl,
-                resolution_template=res_tmpl,
-            )
         if _is_rank_zero():
             logger.info("Final prompt: '%s'", prompt)
 
@@ -2439,7 +2272,6 @@ class Cosmos3OmniDiffusersPipeline(
         guidance_interval: tuple[float, float] | None = None,
         raw_action_dim: int | None = None,
         scheduler: Any | None = None,
-        generator: torch.Generator | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Denoising loop with 3-mode CFG support (parallel, sequential, none).
 
@@ -2551,23 +2383,11 @@ class Cosmos3OmniDiffusersPipeline(
                 if raw_action_dim is not None and 0 < raw_action_dim < action_pred.shape[-1]:
                     action_pred[..., raw_action_dim:] = 0
             if action_latents is None and sound_latents is None:
-                latents = step_scheduler.step(
-                    video_pred,
-                    t,
-                    latents,
-                    generator=generator,
-                    return_dict=False,
-                )[0]
+                latents = step_scheduler.step(video_pred, t, latents, return_dict=False)[0]
             else:
                 packed_noise, shapes, numels = _pack_joint(video_pred, action_pred, sound_pred)
                 packed_latents, _, _ = _pack_joint(latents, action_latents, sound_latents)
-                packed_next = step_scheduler.step(
-                    packed_noise,
-                    t,
-                    packed_latents,
-                    generator=generator,
-                    return_dict=False,
-                )[0]
+                packed_next = step_scheduler.step(packed_noise, t, packed_latents, return_dict=False)[0]
                 unpacked = _unpack_joint(packed_next, shapes, numels)
                 latents = unpacked[0]
                 idx = 1
@@ -2793,7 +2613,6 @@ class Cosmos3OmniDiffusersPipeline(
         velocity_mask: torch.Tensor,
         condition_latents: torch.Tensor,
         guidance_interval: tuple[float, float] | None = None,
-        generator: torch.Generator | None = None,
     ) -> torch.Tensor:
         def _active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool:
             if interval is None:
@@ -2905,13 +2724,7 @@ class Cosmos3OmniDiffusersPipeline(
                 if isinstance(noise_pred, tuple):
                     raise ValueError("Cosmos3 transfer diffusion expects video-only tensor predictions.")
                 noise_pred = noise_pred * velocity_mask
-                latents = self.scheduler.step(
-                    noise_pred,
-                    t,
-                    latents,
-                    generator=generator,
-                    return_dict=False,
-                )[0]
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
                 latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents
         finally:
             self._cosmos3_branch_caches = None
@@ -2928,8 +2741,6 @@ class Cosmos3OmniDiffusersPipeline(
         transfer_video_tensor: torch.Tensor | None,
         transfer_input_fps: float | None,
     ) -> DiffusionOutput:
-        if self.is_distilled_model:
-            raise self._distilled_unsupported_error("transfer requests are unsupported.")
         input_frames = None
         if transfer_video_tensor is not None:
             input_frames = normalized_video_to_uint8_cthw(transfer_video_tensor)
@@ -2995,7 +2806,7 @@ class Cosmos3OmniDiffusersPipeline(
         guidance_scale = (
             float(transfer_config.guidance_scale)
             if transfer_config.guidance_scale is not None
-            else self._resolve_guidance_scale(sp, COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE)
+            else float(sp.guidance_scale or COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE)
         )
         flow_shift_target = float(
             transfer_config.flow_shift
@@ -3010,11 +2821,11 @@ class Cosmos3OmniDiffusersPipeline(
 
         self._guidance_scale = guidance_scale
         self._num_timesteps = num_inference_steps
-        self._set_flow_shift(flow_shift_target)
+        self._set_flow_shift(flow_shift_target, use_karras_sigmas=False)
 
         generator = sp.generator
-        seed = self._resolve_seed(sp, generator)
         if generator is None:
+            seed = sp.seed if sp.seed is not None else 42
             generator = torch.Generator(device=self.device).manual_seed(seed)
 
         cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts(
@@ -3104,11 +2915,7 @@ class Cosmos3OmniDiffusersPipeline(
                 transfer_share_vision_temporal_positions=transfer_config.share_vision_temporal_positions,
             )
 
-            self._set_timesteps(
-                num_inference_steps,
-                device=self.device,
-                shift=self._current_flow_shift,
-            )
+            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
             latents = self.diffuse_transfer(
                 latents=latents,
                 timesteps=self.scheduler.timesteps,
@@ -3123,7 +2930,6 @@ class Cosmos3OmniDiffusersPipeline(
                 shared_kwargs=shared_kwargs,
                 velocity_mask=velocity_mask,
                 condition_latents=condition_latents,
-                generator=generator,
             )
             output_video = self._decode_latents(latents).clamp(-1, 1)
             previous_output = output_video
@@ -3203,20 +3009,6 @@ class Cosmos3OmniDiffusersPipeline(
         action_enabled = action_mode is not None
         transfer_config = resolve_transfer_config(sp, prompt_data)
         action_video_tensor = video_tensor if action_enabled else None
-        is_v2v = video_tensor is not None and not is_t2i and not action_enabled
-        self._validate_distilled_generation_mode(
-            is_t2i=is_t2i,
-            image_tensor=image_tensor,
-            action_enabled=action_enabled,
-            transfer_config=transfer_config,
-            is_v2v=is_v2v,
-            sound_enabled=sound_enabled,
-        )
-        self._validate_edge_generation_mode(
-            transfer_config=transfer_config,
-            is_v2v=is_v2v,
-            sound_enabled=sound_enabled,
-        )
         if transfer_config is not None:
             if is_t2i:
                 raise ValueError("Cosmos3 transfer inference is supported only for video outputs.")
@@ -3264,36 +3056,31 @@ class Cosmos3OmniDiffusersPipeline(
             raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.")
         if video_tensor is not None and is_t2i:
             raise ValueError("Cosmos3 video-to-video generation is supported only for video outputs.")
+        is_v2v = video_tensor is not None and not is_t2i and not action_enabled
 
-        # T2I and T2V share the same model and forward path; their defaults are:
-        #   T2I: regular 1024x1024, Edge 640x640; 50 steps, shift=3.0,
-        #        guidance_interval=[400, 1000]
-        #   T2V: 189 frames and 35 steps; regular guidance=6, shift=10;
-        #        Edge guidance=5, shift=3;
-        #        no guidance interval
+        # T2I and T2V share the same model + forward path; only defaults
+        # differ:
+        #   T2I: 1024x1024, 50 steps, shift=3.0, guidance_interval=[400, 1000]
+        #   T2V: 720x1280,  35 steps, shift=engine-init, no interval
         if is_t2i:
-            height = sp.height or (
-                COSMOS3_EDGE_T2I_DEFAULT_HEIGHT if self.is_edge_model else COSMOS3_T2I_DEFAULT_HEIGHT
-            )
-            width = sp.width or (COSMOS3_EDGE_T2I_DEFAULT_WIDTH if self.is_edge_model else COSMOS3_T2I_DEFAULT_WIDTH)
+            height = sp.height or COSMOS3_T2I_DEFAULT_HEIGHT
+            width = sp.width or COSMOS3_T2I_DEFAULT_WIDTH
             num_frames = 1
             num_inference_steps = sp.num_inference_steps or COSMOS3_T2I_DEFAULT_NUM_INFERENCE_STEPS
-            guidance_scale = self._resolve_guidance_scale(sp, COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE)
+            guidance_scale = sp.guidance_scale if sp.guidance_scale else COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE
             default_flow_shift = COSMOS3_T2I_DEFAULT_FLOW_SHIFT
             default_guidance_interval: tuple[float, float] | None = COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL
             batch_size = max(1, int(sp.num_outputs_per_prompt or 1))
         else:
-            height = sp.height or (
-                COSMOS3_EDGE_T2V_DEFAULT_HEIGHT if self.is_edge_model else COSMOS3_T2V_DEFAULT_HEIGHT
-            )
-            width = sp.width or (COSMOS3_EDGE_T2V_DEFAULT_WIDTH if self.is_edge_model else COSMOS3_T2V_DEFAULT_WIDTH)
-            default_guidance_scale = (
-                COSMOS3_EDGE_T2V_DEFAULT_GUIDANCE_SCALE if self.is_edge_model else COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE
-            )
+            height = sp.height or COSMOS3_T2V_DEFAULT_HEIGHT
+            width = sp.width or COSMOS3_T2V_DEFAULT_WIDTH
             num_frames = sp.num_frames or COSMOS3_T2V_DEFAULT_NUM_FRAMES
             num_inference_steps = sp.num_inference_steps or COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS
-            guidance_scale = self._resolve_guidance_scale(sp, default_guidance_scale)
-            default_flow_shift = self._engine_init_flow_shift
+            guidance_scale = sp.guidance_scale if sp.guidance_scale else COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE
+            # Fall back to the engine-init shift, NOT None: passing None
+            # to ``_set_flow_shift`` would leak a prior T2I rebuild
+            # (shift=3.0) into a subsequent video request.
+            default_flow_shift = COSMOS3_V2V_DEFAULT_FLOW_SHIFT if is_v2v else self._engine_init_flow_shift
             default_guidance_interval = None
             batch_size = 1  # Existing video pipeline assumes B=1.
 
@@ -3316,23 +3103,9 @@ class Cosmos3OmniDiffusersPipeline(
                     f"or action_chunk_size + 1; got num_frames={num_frames}, action_chunk_size={action_chunk_size}."
                 )
             num_inference_steps = sp.num_inference_steps or 30
-            guidance_scale = self._resolve_guidance_scale(sp, 1.0)
+            guidance_scale = sp.guidance_scale if sp.guidance_scale is not None else 1.0
             default_flow_shift = 5.0
 
-        if not is_t2i and not action_enabled:
-            requested_num_frames = int(num_frames)
-            num_frames = _ceil_video_num_frames(
-                requested_num_frames,
-                self.vae_scale_factor_temporal,
-            )
-            if num_frames != requested_num_frames and _is_rank_zero():
-                logger.info(
-                    "Rounded Cosmos3 num_frames from %d to %d for temporal compression factor %d.",
-                    requested_num_frames,
-                    num_frames,
-                    self.vae_scale_factor_temporal,
-                )
-
         domain_id = None
         if action_enabled:
             domain_id = resolve_domain_id(
@@ -3373,12 +3146,13 @@ class Cosmos3OmniDiffusersPipeline(
         self._guidance_scale = guidance_scale
         self._num_timesteps = num_inference_steps
 
-        # Always resolve to a concrete target shift for this request.
-        self._set_flow_shift(flow_shift_target)
+        # Always resolve to a concrete target shift for this request, then
+        # update the shared Diffusers scheduler.
+        self._set_flow_shift(flow_shift_target, use_karras_sigmas=False if is_v2v else None)
 
         generator = sp.generator
-        seed = self._resolve_seed(sp, generator)
         if generator is None:
+            seed = sp.seed if sp.seed is not None else 42
             generator = torch.Generator(device=self.device).manual_seed(seed)
 
         # --- Format prompts & tokenize (B=1; reused across loop iterations
@@ -3519,11 +3293,7 @@ class Cosmos3OmniDiffusersPipeline(
             )
 
         def _run_diffusion(start_latents):
-            self._set_timesteps(
-                num_inference_steps,
-                device=self.device,
-                shift=self._current_flow_shift,
-            )
+            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
             scheduler = self.scheduler
             return self.diffuse(
                 latents=start_latents,
@@ -3544,7 +3314,6 @@ class Cosmos3OmniDiffusersPipeline(
                 guidance_interval=guidance_interval,
                 raw_action_dim=raw_action_dim,
                 scheduler=scheduler,
-                generator=generator,
             )
 
         if is_t2i and batch_size > 1:

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -396,6 +396,7 @@ class TimestepEmbedder(nn.Module):
         hidden_size: int,
         frequency_embedding_size: int = 256,
         max_period: int = 10000,
+        target_dtype: torch.dtype = torch.bfloat16,
     ) -> None:
         super().__init__()
         # Following diffusers naming pattern here for checkpoint compatibility.
@@ -406,7 +407,7 @@ class TimestepEmbedder(nn.Module):
         self.hidden_size = hidden_size
 
         half = frequency_embedding_size // 2
-        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half)
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=target_dtype) / half)
         self.register_buffer("freqs", freqs, persistent=False)
 
     def forward(self, t: torch.Tensor) -> torch.Tensor:
@@ -586,7 +587,6 @@ class Cosmos3CrossAttention(nn.Module):
         head_dim: int,
         rms_norm_eps: float,
         quant_config: QuantizationConfig | None = None,
-        qk_norm: bool = True,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -636,10 +636,8 @@ class Cosmos3CrossAttention(nn.Module):
             prefix=f"{prefix}.to_out",
         )
 
-        self.qk_norm = qk_norm
-        if self.qk_norm:
-            self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
-            self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
 
         self.attn = FrameworkAttention(
             num_heads=self.num_heads,
@@ -723,9 +721,8 @@ class Cosmos3CrossAttention(nn.Module):
         v = self.to_v(hidden_states).view(B, S_gen, self.num_kv_heads_local, self.head_dim)
 
         # Per-head QK norm
-        if self.qk_norm:
-            q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon)
-            k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon)
+        q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon)
+        k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon)
 
         # Qwen3-style RoPE
         q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
@@ -809,8 +806,6 @@ class Cosmos3GenDecoderLayer(nn.Module):
         head_dim: int,
         rms_norm_eps: float,
         quant_config: QuantizationConfig | None = None,
-        mlp_cls: type[nn.Module] = Cosmos3GatedMLP,
-        qk_norm: bool = True,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -822,12 +817,11 @@ class Cosmos3GenDecoderLayer(nn.Module):
             head_dim=head_dim,
             rms_norm_eps=rms_norm_eps,
             quant_config=quant_config,
-            qk_norm=qk_norm,
             prefix=f"{prefix}.cross_attention",
         )
         self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.mlp = mlp_cls(
+        self.mlp = Cosmos3GatedMLP(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
             quant_config=quant_config,
@@ -1002,9 +996,6 @@ class Cosmos3VFMTransformer(nn.Module):
 
     packed_modules_mapping = {}
 
-    _language_model_cls = Cosmos3LanguageModel
-    _gen_mlp_cls = Cosmos3GatedMLP
-
     @staticmethod
     def _is_transformer_block(name: str, module) -> bool:
         return ("gen_layers" in name or "language_model.layers" in name) and name.split(".")[-1].isdigit()
@@ -1041,25 +1032,6 @@ class Cosmos3VFMTransformer(nn.Module):
             if actual != expected:
                 raise ValueError(f"Unsupported Cosmos3 transformer config: {key}={actual!r}; expected {expected!r}.")
 
-    @classmethod
-    def _resolve_rms_norm_eps(cls, model_config: Any) -> float:
-        return float(_tf_config_get(model_config, "rms_norm_eps", 1e-6))
-
-    @classmethod
-    def _resolve_rope_theta(cls, model_config: Any) -> float:
-        return float(_tf_config_get(model_config, "rope_theta", 5_000_000))
-
-    @classmethod
-    def _resolve_mrope_section(cls, model_config: Any) -> list[int]:
-        rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {}
-        return list(rope_scaling.get("mrope_section", [24, 20, 20]))
-
-    def _language_model_kwargs(self) -> dict[str, Any]:
-        return {}
-
-    def validate_loaded_weights(self, loaded: set[str]) -> None:
-        del loaded
-
     def __init__(
         self,
         od_config: OmniDiffusionConfig,
@@ -1072,6 +1044,7 @@ class Cosmos3VFMTransformer(nn.Module):
         super().__init__()
         model_config = od_config.tf_model_config
         self._validate_supported_config(model_config)
+        rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {}
 
         self.hidden_size = int(_tf_config_get(model_config, "hidden_size", 4096))
         self.num_hidden_layers = int(_tf_config_get(model_config, "num_hidden_layers", 36))
@@ -1080,10 +1053,9 @@ class Cosmos3VFMTransformer(nn.Module):
         self.head_dim = int(_tf_config_get(model_config, "head_dim", 128))
         self.intermediate_size = int(_tf_config_get(model_config, "intermediate_size", 12288))
         self.vocab_size = int(_tf_config_get(model_config, "vocab_size", 151936))
-        self.rms_norm_eps = self._resolve_rms_norm_eps(model_config)
-        self.rope_theta = self._resolve_rope_theta(model_config)
-        self.mrope_section = self._resolve_mrope_section(model_config)
-        self.qk_norm_for_diffusion = bool(_tf_config_get(model_config, "qk_norm_for_diffusion", True))
+        self.rms_norm_eps = float(_tf_config_get(model_config, "rms_norm_eps", 1e-6))
+        self.rope_theta = float(_tf_config_get(model_config, "rope_theta", 5_000_000))
+        self.mrope_section = list(rope_scaling.get("mrope_section", [24, 20, 20]))
         self.latent_patch_size = int(_tf_config_get(model_config, "latent_patch_size", 2))
         self.latent_channel_size = int(_tf_config_get(model_config, "latent_channel", 48))
         self.timestep_scale = float(_tf_config_get(model_config, "timestep_scale", 0.001))
@@ -1123,12 +1095,10 @@ class Cosmos3VFMTransformer(nn.Module):
         )
         self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size
 
-        self.use_k_norm_und_for_gen = _tf_config_get(model_config, "use_k_norm_und_for_gen", None)
-
         dtype = od_config.dtype
         quant_config = getattr(od_config, "quantization_config", None) if od_config else None
 
-        self.language_model = self._language_model_cls(
+        self.language_model = Cosmos3LanguageModel(
             hidden_size=self.hidden_size,
             intermediate_size=self.intermediate_size,
             num_hidden_layers=self.num_hidden_layers,
@@ -1141,13 +1111,12 @@ class Cosmos3VFMTransformer(nn.Module):
             mrope_section=self.mrope_section,
             quant_config=quant_config,
             prefix="language_model",
-            **self._language_model_kwargs(),
         )
 
         # Video projection layers are small; not worth quantizing.
         self.proj_in = nn.Linear(self.patch_latent_dim, self.hidden_size)
         self.proj_out = nn.Linear(self.hidden_size, self.patch_latent_dim)
-        self.time_embedder = TimestepEmbedder(self.hidden_size)
+        self.time_embedder = TimestepEmbedder(self.hidden_size, target_dtype=dtype)
         if self.action_gen:
             self.action_proj_in = DomainAwareLinear(
                 self.action_dim,
@@ -1178,8 +1147,6 @@ class Cosmos3VFMTransformer(nn.Module):
                     head_dim=self.head_dim,
                     rms_norm_eps=self.rms_norm_eps,
                     quant_config=quant_config,
-                    mlp_cls=self._gen_mlp_cls,
-                    qk_norm=self.qk_norm_for_diffusion,
                     prefix=f"gen_layers.{i}",
                 )
                 for i in range(self.num_hidden_layers)

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py

@@ -1,396 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Cosmos3 Edge transformer variant with a Nemotron dense UND backbone."""
-
-from __future__ import annotations
-
-from typing import Any
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from vllm.distributed import get_tensor_model_parallel_world_size
-from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
-from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
-
-from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
-
-from .transformer_cosmos3 import (
-    Cosmos3VFMTransformer,
-    Qwen3VLTextRotaryEmbedding,
-    RMSNorm,
-    _apply_rotary_pos_emb,
-    _nested_get,
-    _tf_config_get,
-)
-
-COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense"
-
-
-def _tf_config_nested_get(config: Any, key: str, default: Any = None) -> Any:
-    if config is None:
-        return default
-    if isinstance(config, dict):
-        value = _nested_get(config, key)
-    elif hasattr(config, "to_dict"):
-        value = _nested_get(config.to_dict(), key)
-    elif hasattr(config, "params"):
-        value = _nested_get(config.params, key)
-    else:
-        value = getattr(config, key, None)
-    return default if value is None else value
-
-
-class Cosmos3Relu2MLP(nn.Module):
-    """Nemotron dense MLP: down_proj(relu(up_proj(x)) ** 2)."""
-
-    def __init__(
-        self,
-        hidden_size: int,
-        intermediate_size: int,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.up_proj = ColumnParallelLinear(
-            hidden_size,
-            intermediate_size,
-            bias=False,
-            gather_output=False,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.up_proj",
-        )
-        self.down_proj = RowParallelLinear(
-            intermediate_size,
-            hidden_size,
-            bias=False,
-            input_is_parallel=True,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.down_proj",
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        hidden = F.relu(self.up_proj(x))
-        return self.down_proj(hidden * hidden)
-
-
-class Cosmos3EdgeCausalAttention(nn.Module):
-    """Edge UND causal attention.
-
-    Reasoner self-attention uses raw UND Q/K with RoPE because
-    ``qk_norm_for_text=false``.  The GEN-facing UND key cache is a separate
-    normalized-then-RoPE'd view of the raw K tensor when GEN QK norm is enabled;
-    no-GEN-QK-norm checkpoints reuse the raw RoPE'd UND K.
-    """
-
-    def __init__(
-        self,
-        *,
-        hidden_size: int,
-        num_attention_heads: int,
-        num_key_value_heads: int,
-        head_dim: int,
-        rms_norm_eps: float,
-        use_k_norm_und_for_gen: bool,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.num_heads = num_attention_heads
-        self.num_kv_heads = num_key_value_heads
-        self.head_dim = head_dim
-
-        tp_size = get_tensor_model_parallel_world_size()
-        self.num_heads_local = self.num_heads // tp_size
-        self.num_kv_heads_local = self.num_kv_heads // tp_size
-
-        self.to_q = ColumnParallelLinear(
-            hidden_size,
-            self.num_heads * self.head_dim,
-            bias=False,
-            gather_output=False,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.to_q",
-        )
-        self.to_k = ColumnParallelLinear(
-            hidden_size,
-            self.num_kv_heads * self.head_dim,
-            bias=False,
-            gather_output=False,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.to_k",
-        )
-        self.to_v = ColumnParallelLinear(
-            hidden_size,
-            self.num_kv_heads * self.head_dim,
-            bias=False,
-            gather_output=False,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.to_v",
-        )
-        self.to_out = RowParallelLinear(
-            self.num_heads * self.head_dim,
-            hidden_size,
-            bias=False,
-            input_is_parallel=True,
-            return_bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.to_out",
-        )
-
-        self.k_norm_und_for_gen = RMSNorm(self.head_dim, eps=rms_norm_eps) if use_k_norm_und_for_gen else None
-        self.attn = FrameworkAttention(
-            num_heads=self.num_heads,
-            head_size=self.head_dim,
-            causal=True,
-            softmax_scale=1.0 / (self.head_dim**0.5),
-            num_kv_heads=self.num_kv_heads,
-            skip_sequence_parallel=True,
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        freqs_cos: torch.Tensor,
-        freqs_sin: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        B, S, _ = hidden_states.shape
-
-        q = self.to_q(hidden_states).view(B, S, self.num_heads_local, self.head_dim)
-        k = self.to_k(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim)
-        v = self.to_v(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim)
-
-        q_reasoner, k_reasoner = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
-
-        if self.k_norm_und_for_gen is not None:
-            k_gen = F.rms_norm(
-                k,
-                (self.head_dim,),
-                self.k_norm_und_for_gen.weight,
-                eps=self.k_norm_und_for_gen.variance_epsilon,
-            )
-            _, k_gen = _apply_rotary_pos_emb(q, k_gen, freqs_cos, freqs_sin)
-        else:
-            k_gen = k_reasoner
-
-        out = self.attn(q_reasoner, k_reasoner, v).reshape(B, S, -1)
-        return self.to_out(out), k_gen, v
-
-
-class Cosmos3EdgeUndDecoderLayer(nn.Module):
-    """Edge UND decoder layer: causal self-attention + ReLU2 MLP."""
-
-    def __init__(
-        self,
-        *,
-        hidden_size: int,
-        intermediate_size: int,
-        num_attention_heads: int,
-        num_key_value_heads: int,
-        head_dim: int,
-        rms_norm_eps: float,
-        use_k_norm_und_for_gen: bool,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.self_attn = Cosmos3EdgeCausalAttention(
-            hidden_size=hidden_size,
-            num_attention_heads=num_attention_heads,
-            num_key_value_heads=num_key_value_heads,
-            head_dim=head_dim,
-            rms_norm_eps=rms_norm_eps,
-            use_k_norm_und_for_gen=use_k_norm_und_for_gen,
-            quant_config=quant_config,
-            prefix=f"{prefix}.self_attn",
-        )
-        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.mlp = Cosmos3Relu2MLP(
-            hidden_size=hidden_size,
-            intermediate_size=intermediate_size,
-            quant_config=quant_config,
-            prefix=f"{prefix}.mlp",
-        )
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        freqs: tuple[torch.Tensor, torch.Tensor],
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
-
-        cos, sin = freqs
-        attn_out, k, v = self.self_attn(hidden_states, cos, sin)
-        hidden_states = residual + attn_out
-
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = residual + self.mlp(hidden_states)
-
-        return hidden_states, k, v
-
-
-class Cosmos3EdgeLanguageModel(nn.Module):
-    """Nemotron dense UND tower that returns GEN-facing K plus raw V per layer."""
-
-    _layerwise_offload_blocks_attrs = ["layers"]
-
-    def __init__(
-        self,
-        *,
-        hidden_size: int,
-        intermediate_size: int,
-        num_hidden_layers: int,
-        num_attention_heads: int,
-        num_key_value_heads: int,
-        head_dim: int,
-        vocab_size: int,
-        rms_norm_eps: float,
-        rope_theta: float,
-        mrope_section: list[int],
-        use_k_norm_und_for_gen: bool,
-        quant_config: QuantizationConfig | None = None,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
-        self.rotary_emb = Qwen3VLTextRotaryEmbedding(
-            head_dim=head_dim,
-            rope_theta=rope_theta,
-            mrope_section=mrope_section,
-        )
-        self.layers = nn.ModuleList(
-            [
-                Cosmos3EdgeUndDecoderLayer(
-                    hidden_size=hidden_size,
-                    intermediate_size=intermediate_size,
-                    num_attention_heads=num_attention_heads,
-                    num_key_value_heads=num_key_value_heads,
-                    head_dim=head_dim,
-                    rms_norm_eps=rms_norm_eps,
-                    use_k_norm_und_for_gen=use_k_norm_und_for_gen,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.layers.{i}",
-                )
-                for i in range(num_hidden_layers)
-            ]
-        )
-        self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
-
-    def forward(
-        self,
-        text_ids: torch.Tensor,
-        freqs: tuple[torch.Tensor, torch.Tensor],
-    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
-        hidden = self.embed_tokens(text_ids)
-
-        cached_kv: list[tuple[torch.Tensor, torch.Tensor]] = []
-        for layer in self.layers:
-            hidden, k, v = layer(hidden, freqs)
-            cached_kv.append((k, v))
-
-        return cached_kv
-
-
-class Cosmos3EdgeVFMTransformer(Cosmos3VFMTransformer):
-    """Cosmos3 Edge variant with Nemotron dense UND and shared GEN diffusion."""
-
-    _language_model_cls = Cosmos3EdgeLanguageModel
-    _gen_mlp_cls = Cosmos3Relu2MLP
-
-    @staticmethod
-    def _validate_supported_config(model_config: Any) -> None:
-        expected_values = {
-            "qk_norm_for_text": False,
-            "position_embedding_type": "unified_3d_mrope",
-            "unified_3d_mrope_reset_spatial_ids": True,
-            "joint_attn_implementation": "two_way",
-        }
-        for key, expected in expected_values.items():
-            actual = _tf_config_get(model_config, key, expected)
-            if actual != expected:
-                raise ValueError(
-                    f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}; expected {expected!r}."
-                )
-
-        backbone_type = _tf_config_get(model_config, "backbone_type", None)
-        if backbone_type is None:
-            raise ValueError(
-                f"Cosmos3 Edge transformer config must declare backbone_type={COSMOS3_EDGE_BACKBONE_TYPE!r}."
-            )
-        if backbone_type != COSMOS3_EDGE_BACKBONE_TYPE:
-            raise ValueError(
-                "Unsupported Cosmos3 Edge transformer config: "
-                f"backbone_type={backbone_type!r}; expected {COSMOS3_EDGE_BACKBONE_TYPE!r}."
-            )
-
-        required_values = {
-            "latent_channel": 48,
-            "latent_patch_size": 2,
-            "temporal_compression_factor": 4,
-        }
-        for key, expected in required_values.items():
-            actual = _tf_config_get(model_config, key, None)
-            if actual is None:
-                raise ValueError(f"Cosmos3 Edge transformer config must declare {key}={expected}.")
-            if int(actual) != expected:
-                raise ValueError(
-                    f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}; expected {expected!r}."
-                )
-
-    @classmethod
-    def _resolve_rms_norm_eps(cls, model_config: Any) -> float:
-        return float(
-            _tf_config_get(
-                model_config,
-                "rms_norm_eps",
-                _tf_config_get(model_config, "layer_norm_epsilon", 1e-5),
-            )
-        )
-
-    @classmethod
-    def _resolve_rope_theta(cls, model_config: Any) -> float:
-        return float(_tf_config_get(model_config, "rope_theta", 100_000_000))
-
-    @classmethod
-    def _resolve_mrope_section(cls, model_config: Any) -> list[int]:
-        mrope_section = _tf_config_get(model_config, "mrope_section", None)
-        if mrope_section is None:
-            mrope_section = _tf_config_nested_get(model_config, "mrope_section", [24, 20, 20])
-        return list(mrope_section)
-
-    def _language_model_kwargs(self) -> dict[str, Any]:
-        return {"use_k_norm_und_for_gen": self.use_k_norm_und_for_gen}
-
-    def validate_loaded_weights(self, loaded: set[str]) -> None:
-        missing: list[str] = []
-        for layer_idx in range(self.num_hidden_layers):
-            required_markers = (
-                f"language_model.layers.{layer_idx}.mlp.up_proj.",
-                f"language_model.layers.{layer_idx}.mlp.down_proj.",
-                f"gen_layers.{layer_idx}.mlp.up_proj.",
-                f"gen_layers.{layer_idx}.mlp.down_proj.",
-            )
-            if self.use_k_norm_und_for_gen:
-                required_markers = (
-                    *required_markers,
-                    f"language_model.layers.{layer_idx}.self_attn.k_norm_und_for_gen.",
-                )
-            missing.extend(
-                marker.rstrip(".") for marker in required_markers if not any(marker in name for name in loaded)
-            )
-        if missing:
-            preview = ", ".join(missing[:10])
-            suffix = " ..." if len(missing) > 10 else ""
-            raise ValueError(
-                "Cosmos3 Edge transformer checkpoint is missing required weights for "
-                f"{preview}{suffix}. Use a converted Edge checkpoint with ReLU2 UND/GEN MLP weights."
-            )

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -53,7 +53,6 @@ from vllm_omni.diffusion.models.dreamzero.utils import (
 )
 from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 MAX_DREAMZERO_SESSIONS = 64
@@ -611,12 +610,12 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         latents: tuple[torch.Tensor, torch.Tensor],
         do_true_cfg: bool,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Post-step sync: .contiguous() + platform synchronize."""
+        """Post-step sync: .contiguous() + cuda.synchronize()"""
         latents = tuple(t.contiguous() for t in latents)
         if do_true_cfg and get_classifier_free_guidance_world_size() > 1:
-            # Platform-dispatched sync (not torch.cuda directly) so DreamZero
-            # keeps working on non-CUDA hardware (XPU/NPU, see PR #4399).
-            current_omni_platform.synchronize()
+            device = next((t.device for t in latents if t.is_cuda), None)
+            if device is not None:
+                torch.cuda.current_stream(device).synchronize()
         return latents
 
     # -----------------------------------------------------------------------

vllm_omni/diffusion/models/schedulers/__init__.py

@@ -2,15 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from vllm_omni.diffusion.models.schedulers.scheduling_dmd2_euler import DMD2EulerScheduler
-from vllm_omni.diffusion.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler,
-)
 from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import (
     FlowUniPCMultistepScheduler,
 )
 
 __all__ = [
     "DMD2EulerScheduler",
-    "FlowMatchEulerDiscreteScheduler",
     "FlowUniPCMultistepScheduler",
 ]

vllm_omni/diffusion/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -1,660 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-# Originally from https://github.com/huggingface/diffusers/blob/v0.39.0/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py
-# Copyright 2025 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Flow Match Euler scheduler copied from Diffusers v0.39.0.
-
-Modified by vLLM-Omni contributors for local integration and RNG compatibility.
-"""
-
-import math
-from dataclasses import dataclass
-from typing import Literal
-
-import numpy as np
-import torch
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.schedulers.scheduling_utils import SchedulerMixin
-from diffusers.utils import BaseOutput, is_scipy_available, logging
-from diffusers.utils.torch_utils import randn_tensor
-
-if is_scipy_available():
-    import scipy.stats
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-@dataclass
-class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
-    """
-    Output class for the scheduler's `step` function output.
-
-    Args:
-        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
-            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
-            denoising loop.
-    """
-
-    prev_sample: torch.FloatTensor
-
-
-class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
-    """
-    Euler scheduler.
-
-    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
-    methods the library implements for all schedulers such as loading and saving.
-
-    Args:
-        num_train_timesteps (`int`, defaults to 1000):
-            The number of diffusion steps to train the model.
-        shift (`float`, defaults to 1.0):
-            The shift value for the timestep schedule.
-        use_dynamic_shifting (`bool`, defaults to False):
-            Whether to apply timestep shifting on-the-fly based on the image resolution.
-        base_shift (`float`, defaults to 0.5):
-            Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent
-            with desired output.
-        max_shift (`float`, defaults to 1.15):
-            Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be
-            more exaggerated or stylized.
-        base_image_seq_len (`int`, defaults to 256):
-            The base image sequence length.
-        max_image_seq_len (`int`, defaults to 4096):
-            The maximum image sequence length.
-        invert_sigmas (`bool`, defaults to False):
-            Whether to invert the sigmas.
-        shift_terminal (`float`, defaults to None):
-            The end value of the shifted timestep schedule.
-        use_karras_sigmas (`bool`, defaults to False):
-            Whether to use Karras sigmas for step sizes in the noise schedule during sampling.
-        use_exponential_sigmas (`bool`, defaults to False):
-            Whether to use exponential sigmas for step sizes in the noise schedule during sampling.
-        use_beta_sigmas (`bool`, defaults to False):
-            Whether to use beta sigmas for step sizes in the noise schedule during sampling.
-        time_shift_type (`str`, defaults to "exponential"):
-            The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".
-        stochastic_sampling (`bool`, defaults to False):
-            Whether to use stochastic sampling.
-    """
-
-    _compatibles = []
-    order = 1
-
-    @register_to_config
-    def __init__(
-        self,
-        num_train_timesteps: int = 1000,
-        shift: float = 1.0,
-        use_dynamic_shifting: bool = False,
-        base_shift: float | None = 0.5,
-        max_shift: float | None = 1.15,
-        base_image_seq_len: int = 256,
-        max_image_seq_len: int = 4096,
-        invert_sigmas: bool = False,
-        shift_terminal: float = None,
-        use_karras_sigmas: bool = False,
-        use_exponential_sigmas: bool = False,
-        use_beta_sigmas: bool = False,
-        time_shift_type: Literal["exponential", "linear"] = "exponential",
-        stochastic_sampling: bool = False,
-    ):
-        if self.config.use_beta_sigmas and not is_scipy_available():
-            raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
-        if (
-            sum(
-                [
-                    self.config.use_beta_sigmas,
-                    self.config.use_exponential_sigmas,
-                    self.config.use_karras_sigmas,
-                ]
-            )
-            > 1
-        ):
-            raise ValueError(
-                "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, "
-                "`config.use_karras_sigmas` can be used."
-            )
-        if time_shift_type not in {"exponential", "linear"}:
-            raise ValueError("`time_shift_type` must either be 'exponential' or 'linear'.")
-
-        timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
-        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
-
-        sigmas = timesteps / num_train_timesteps
-        if not use_dynamic_shifting:
-            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
-            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
-
-        self.timesteps = sigmas * num_train_timesteps
-
-        self._step_index = None
-        self._begin_index = None
-
-        self._shift = shift
-
-        self.sigmas = sigmas.to("cpu")  # to avoid too much CPU/GPU communication
-        self.sigma_min = self.sigmas[-1].item()
-        self.sigma_max = self.sigmas[0].item()
-
-    @property
-    def shift(self):
-        """
-        The value used for shifting.
-        """
-        return self._shift
-
-    @property
-    def step_index(self):
-        """
-        The index counter for current timestep. It will increase 1 after each scheduler step.
-        """
-        return self._step_index
-
-    @property
-    def begin_index(self):
-        """
-        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
-        """
-        return self._begin_index
-
-    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
-    def set_begin_index(self, begin_index: int = 0):
-        """
-        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
-
-        Args:
-            begin_index (`int`, defaults to `0`):
-                The begin index for the scheduler.
-        """
-        self._begin_index = begin_index
-
-    def set_shift(self, shift: float):
-        """
-        Sets the shift value for the scheduler.
-
-        Args:
-            shift (`float`):
-                The shift value to be set.
-        """
-        self._shift = shift
-
-    def scale_noise(
-        self,
-        sample: torch.FloatTensor,
-        timestep: float | torch.FloatTensor,
-        noise: torch.FloatTensor | None = None,
-    ) -> torch.FloatTensor:
-        """
-        Forward process in flow-matching
-
-        Args:
-            sample (`torch.FloatTensor`):
-                The input sample.
-            timestep (`torch.FloatTensor`):
-                The current timestep in the diffusion chain.
-            noise (`torch.FloatTensor`):
-                The noise tensor.
-
-        Returns:
-            `torch.FloatTensor`:
-                A scaled input sample.
-        """
-        # Make sure sigmas and timesteps have the same device and dtype as original_samples
-        sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
-
-        if sample.device.type == "mps" and torch.is_floating_point(timestep):
-            # mps does not support float64
-            schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
-            timestep = timestep.to(sample.device, dtype=torch.float32)
-        else:
-            schedule_timesteps = self.timesteps.to(sample.device)
-            timestep = timestep.to(sample.device)
-
-        # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
-        if self.begin_index is None:
-            step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
-        elif self.step_index is not None:
-            # add_noise is called after first denoising step (for inpainting)
-            step_indices = [self.step_index] * timestep.shape[0]
-        else:
-            # add noise is called before first denoising step to create initial latent(img2img)
-            step_indices = [self.begin_index] * timestep.shape[0]
-
-        sigma = sigmas[step_indices].flatten()
-        while len(sigma.shape) < len(sample.shape):
-            sigma = sigma.unsqueeze(-1)
-
-        sample = sigma * noise + (1.0 - sigma) * sample
-
-        return sample
-
-    def _sigma_to_t(self, sigma) -> float:
-        return sigma * self.config.num_train_timesteps
-
-    def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
-        """
-        Apply time shifting to the sigmas.
-
-        Args:
-            mu (`float`):
-                The mu parameter for the time shift.
-            sigma (`float`):
-                The sigma parameter for the time shift.
-            t (`torch.Tensor`):
-                The input timesteps.
-
-        Returns:
-            `torch.Tensor`:
-                The time-shifted timesteps.
-        """
-        if self.config.time_shift_type == "exponential":
-            return self._time_shift_exponential(mu, sigma, t)
-        elif self.config.time_shift_type == "linear":
-            return self._time_shift_linear(mu, sigma, t)
-
-    def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor:
-        r"""
-        Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config
-        value.
-
-        Reference:
-        https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51
-
-        Args:
-            t (`torch.Tensor`):
-                A tensor of timesteps to be stretched and shifted.
-
-        Returns:
-            `torch.Tensor`:
-                A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`.
-        """
-        one_minus_z = 1 - t
-        scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal)
-        stretched_t = 1 - (one_minus_z / scale_factor)
-        return stretched_t
-
-    def set_timesteps(
-        self,
-        num_inference_steps: int | None = None,
-        device: str | torch.device = None,
-        sigmas: list[float] | None = None,
-        mu: float | None = None,
-        timesteps: list[float] | None = None,
-    ):
-        """
-        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
-
-        Args:
-            num_inference_steps (`int`, *optional*):
-                The number of diffusion steps used when generating samples with a pre-trained model.
-            device (`str` or `torch.device`, *optional*):
-                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-            sigmas (`list[float]`, *optional*):
-                Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed
-                automatically.
-            mu (`float`, *optional*):
-                Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep
-                shifting.
-            timesteps (`list[float]`, *optional*):
-                Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed
-                automatically.
-        """
-        if self.config.use_dynamic_shifting and mu is None:
-            raise ValueError("`mu` must be passed when `use_dynamic_shifting` is set to be `True`")
-
-        if sigmas is not None and timesteps is not None:
-            if len(sigmas) != len(timesteps):
-                raise ValueError("`sigmas` and `timesteps` should have the same length")
-
-        if num_inference_steps is not None:
-            if (sigmas is not None and len(sigmas) != num_inference_steps) or (
-                timesteps is not None and len(timesteps) != num_inference_steps
-            ):
-                raise ValueError(
-                    "`sigmas` and `timesteps` should have the same length as num_inference_steps, "
-                    "if `num_inference_steps` is provided"
-                )
-        else:
-            num_inference_steps = len(sigmas) if sigmas is not None else len(timesteps)
-
-        self.num_inference_steps = num_inference_steps
-
-        # 1. Prepare default sigmas
-        is_timesteps_provided = timesteps is not None
-
-        if is_timesteps_provided:
-            timesteps = np.array(timesteps).astype(np.float32)
-
-        if sigmas is None:
-            if timesteps is None:
-                timesteps = np.linspace(
-                    self._sigma_to_t(self.sigma_max),
-                    self._sigma_to_t(self.sigma_min),
-                    num_inference_steps,
-                )
-            sigmas = timesteps / self.config.num_train_timesteps
-        else:
-            sigmas = np.array(sigmas).astype(np.float32)
-            num_inference_steps = len(sigmas)
-
-        # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of
-        #    "exponential" or "linear" type is applied
-        if self.config.use_dynamic_shifting:
-            sigmas = self.time_shift(mu, 1.0, sigmas)
-        else:
-            sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)
-
-        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value
-        if self.config.shift_terminal:
-            sigmas = self.stretch_shift_to_terminal(sigmas)
-
-        # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules
-        if self.config.use_karras_sigmas:
-            sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
-        elif self.config.use_exponential_sigmas:
-            sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
-        elif self.config.use_beta_sigmas:
-            sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
-
-        # 5. Convert sigmas and timesteps to tensors and move to specified device
-        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
-        if not is_timesteps_provided:
-            timesteps = sigmas * self.config.num_train_timesteps
-        else:
-            timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device)
-
-        # 6. Append the terminal sigma value.
-        #    If a model requires inverted sigma schedule for denoising but timesteps without inversion, the
-        #    `invert_sigmas` flag can be set to `True`. This case is only required in Mochi
-        if self.config.invert_sigmas:
-            sigmas = 1.0 - sigmas
-            timesteps = sigmas * self.config.num_train_timesteps
-            sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)])
-        else:
-            sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
-
-        self.timesteps = timesteps
-        self.sigmas = sigmas
-        self._step_index = None
-        self._begin_index = None
-
-    def index_for_timestep(
-        self,
-        timestep: float | torch.FloatTensor,
-        schedule_timesteps: torch.FloatTensor | None = None,
-    ) -> int:
-        """
-        Get the index for the given timestep.
-
-        Args:
-            timestep (`float` or `torch.FloatTensor`):
-                The timestep to find the index for.
-            schedule_timesteps (`torch.FloatTensor`, *optional*):
-                The schedule timesteps to validate against. If `None`, the scheduler's timesteps are used.
-
-        Returns:
-            `int`:
-                The index of the timestep.
-        """
-        if schedule_timesteps is None:
-            schedule_timesteps = self.timesteps
-
-        indices = (schedule_timesteps == timestep).nonzero()
-
-        # The sigma index that is taken for the **very** first `step`
-        # is always the second index (or the last index if there is only 1)
-        # This way we can ensure we don't accidentally skip a sigma in
-        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
-        pos = 1 if len(indices) > 1 else 0
-
-        return indices[pos].item()
-
-    def _init_step_index(self, timestep: float | torch.FloatTensor) -> None:
-        if self.begin_index is None:
-            if isinstance(timestep, torch.Tensor):
-                timestep = timestep.to(self.timesteps.device)
-            self._step_index = self.index_for_timestep(timestep)
-        else:
-            self._step_index = self._begin_index
-
-    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,
-    ) -> FlowMatchEulerDiscreteSchedulerOutput | 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).
-
-        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`, defaults to `True`):
-                Whether or not to return a
-                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple.
-
-        Returns:
-            [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`:
-                If return_dict is `True`,
-                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned,
-                otherwise a tuple is returned where the first element is the sample tensor.
-        """
-
-        if isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) or isinstance(timestep, 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 per_token_timesteps is not None:
-            per_token_sigmas = per_token_timesteps / self.config.num_train_timesteps
-
-            sigmas = self.sigmas[:, None, None]
-            lower_mask = sigmas < per_token_sigmas[None] - 1e-6
-            lower_sigmas = lower_mask * sigmas
-            lower_sigmas, _ = lower_sigmas.max(dim=0)
-
-            current_sigma = per_token_sigmas[..., None]
-            next_sigma = lower_sigmas[..., None]
-            dt = current_sigma - next_sigma
-        else:
-            sigma_idx = self.step_index
-            sigma = self.sigmas[sigma_idx]
-            sigma_next = self.sigmas[sigma_idx + 1]
-
-            current_sigma = sigma
-            next_sigma = sigma_next
-            dt = sigma_next - sigma
-
-        if self.config.stochastic_sampling:
-            x0 = sample - current_sigma * model_output
-            noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
-            prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise
-        else:
-            prev_sample = sample + dt * model_output
-
-        # 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,)
-
-        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
-
-    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
-    def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
-        """
-        Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative
-        Models](https://huggingface.co/papers/2206.00364).
-
-        Args:
-            in_sigmas (`torch.Tensor`):
-                The input sigma values to be converted.
-            num_inference_steps (`int`):
-                The number of inference steps to generate the noise schedule for.
-
-        Returns:
-            `torch.Tensor`:
-                The converted sigma values following the Karras noise schedule.
-        """
-
-        # Hack to make sure that other schedulers which copy this function don't break
-        # TODO: Add this logic to the other schedulers
-        if hasattr(self.config, "sigma_min"):
-            sigma_min = self.config.sigma_min
-        else:
-            sigma_min = None
-
-        if hasattr(self.config, "sigma_max"):
-            sigma_max = self.config.sigma_max
-        else:
-            sigma_max = None
-
-        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
-        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
-
-        rho = 7.0  # 7.0 is the value used in the paper
-        ramp = np.linspace(0, 1, num_inference_steps)
-        min_inv_rho = sigma_min ** (1 / rho)
-        max_inv_rho = sigma_max ** (1 / rho)
-        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
-        return sigmas
-
-    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential
-    def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
-        """
-        Construct an exponential noise schedule.
-
-        Args:
-            in_sigmas (`torch.Tensor`):
-                The input sigma values to be converted.
-            num_inference_steps (`int`):
-                The number of inference steps to generate the noise schedule for.
-
-        Returns:
-            `torch.Tensor`:
-                The converted sigma values following an exponential schedule.
-        """
-
-        # Hack to make sure that other schedulers which copy this function don't break
-        # TODO: Add this logic to the other schedulers
-        if hasattr(self.config, "sigma_min"):
-            sigma_min = self.config.sigma_min
-        else:
-            sigma_min = None
-
-        if hasattr(self.config, "sigma_max"):
-            sigma_max = self.config.sigma_max
-        else:
-            sigma_max = None
-
-        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
-        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
-
-        sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps))
-        return sigmas
-
-    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta
-    def _convert_to_beta(
-        self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6
-    ) -> torch.Tensor:
-        """
-        Construct a beta noise schedule as proposed in [Beta Sampling is All You
-        Need](https://huggingface.co/papers/2407.12173).
-
-        Args:
-            in_sigmas (`torch.Tensor`):
-                The input sigma values to be converted.
-            num_inference_steps (`int`):
-                The number of inference steps to generate the noise schedule for.
-            alpha (`float`, *optional*, defaults to `0.6`):
-                The alpha parameter for the beta distribution.
-            beta (`float`, *optional*, defaults to `0.6`):
-                The beta parameter for the beta distribution.
-
-        Returns:
-            `torch.Tensor`:
-                The converted sigma values following a beta distribution schedule.
-        """
-
-        # Hack to make sure that other schedulers which copy this function don't break
-        # TODO: Add this logic to the other schedulers
-        if hasattr(self.config, "sigma_min"):
-            sigma_min = self.config.sigma_min
-        else:
-            sigma_min = None
-
-        if hasattr(self.config, "sigma_max"):
-            sigma_max = self.config.sigma_max
-        else:
-            sigma_max = None
-
-        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
-        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
-
-        sigmas = np.array(
-            [
-                sigma_min + (ppf * (sigma_max - sigma_min))
-                for ppf in [
-                    scipy.stats.beta.ppf(timestep, alpha, beta)
-                    for timestep in 1 - np.linspace(0, 1, num_inference_steps)
-                ]
-            ]
-        )
-        return sigmas
-
-    def _time_shift_exponential(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
-        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
-
-    def _time_shift_linear(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
-        return mu / (mu + (1 / t - 1) ** sigma)
-
-    def __len__(self) -> int:
-        return self.config.num_train_timesteps

vllm_omni/diffusion/registry.py

@@ -276,11 +276,6 @@ _DIFFUSION_MODELS = {
         "pipeline_cosmos3",
         "Cosmos3OmniDiffusersPipeline",
     ),
-    "Cosmos3OmniPipeline": (
-        "cosmos3",
-        "pipeline_cosmos3",
-        "Cosmos3OmniDiffusersPipeline",
-    ),
     "SoulXSingerPipeline": (
         "soulx_singer",
         "pipeline_soulx_singer_svs",
@@ -541,7 +536,6 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func",
     "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func",
-    "Cosmos3OmniPipeline": "get_cosmos3_post_process_func",
     "HiDreamImagePipeline": "get_hidream_image_post_process_func",
     "StableDiffusionXLPipeline": "get_sdxl_image_post_process_func",
     "Krea2Pipeline": "get_krea2_post_process_func",
@@ -552,7 +546,6 @@ _DIFFUSION_IR_OP_PRIORITY_FUNCS = {
     # `ir_op_priority_func` function must be placed in {mod_folder}/{mod_relname}.py,
     # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key.
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_ir_op_priority_func",
-    "Cosmos3OmniPipeline": "get_cosmos3_ir_op_priority_func",
 }
 
 _DIFFUSION_PRE_PROCESS_FUNCS = {
@@ -577,7 +570,6 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func",
     "MagiHumanPipeline": "get_magi_human_pre_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func",
-    "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func",
     "SoulXSingerPipeline": "get_soulxsinger_pre_process_func",
     "SoulXSingerSVCPipeline": "get_soulxsinger_svc_pre_process_func",
 }

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -83,13 +83,7 @@ class _DiffusionVllmModelConfig:
 
     @property
     def is_diffusion(self) -> bool:
-        # Duck-types ``vllm.config.ModelConfig.is_diffusion``, which detects
-        # *discrete diffusion (dLLM) text models* via ``hf_config.canvas_length``
-        # and gates dLLM-only scheduler/spec-decode paths in vLLM.  Omni's DiT
-        # pipelines run in their own engine and must never enable those paths,
-        # so mirror the upstream check (False for every DiT model) instead of
-        # leaving the attribute missing.
-        return getattr(self.hf_config, "canvas_length", None) is not None
+        return False
 
 
 def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig:
@@ -294,12 +288,6 @@ class DiffusionWorker:
         # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping.
         # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion.
         vllm_config.kernel_config.ir_op_priority = _resolve_ir_op_priority(self.od_config, vllm_config)
-        if self.od_config.moe_backend != "auto":
-            logger.warning(
-                "Overriding MoE backend from default 'auto' to '%s' per deploy config.",
-                self.od_config.moe_backend,
-            )
-            vllm_config.kernel_config.moe_backend = self.od_config.moe_backend
         logger.info(
             "Final IR op priority after setting vLLM-Omni overrides: %s", vllm_config.kernel_config.ir_op_priority
         )

vllm_omni/distributed/omni_connectors/utils/serialization.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from collections.abc import Mapping
-from dataclasses import asdict, is_dataclass
+from dataclasses import asdict, fields, is_dataclass
 from typing import Any
 
 import msgspec
@@ -59,7 +59,10 @@ class OmniMsgpackEncoder:
         if isinstance(obj, Image.Image):
             return self._encode_pil_image(obj)
 
-        # RequestOutput (not a dataclass, needs special handling)
+        # RequestOutput (not a dataclass, needs special handling).
+        # Note: OmniRequestOutput subclasses RequestOutput but is a dataclass,
+        # so msgspec encodes it natively (field map) and it never reaches
+        # this hook.
         if isinstance(obj, RequestOutput):
             return self._encode_request_output(obj)
 
@@ -252,23 +255,30 @@ class OmniMsgpackDecoder:
     def _decode_omni_request_output(self, obj: dict[str, Any]) -> Any:
         """Decode dict to OmniRequestOutput.
 
-        OmniRequestOutput is a dataclass, so we can use msgspec.convert
-        or construct it directly.
+        By the time this runs, ``_post_process`` has already reconstructed
+        nested values bottom-up (CompletionOutput, tensors, images, ...), so
+        the object is built directly from the known dataclass fields.
         """
         from vllm_omni.outputs import OmniRequestOutput
 
+        data = dict(obj)
+        # Legacy wire format (pre-inheritance): the stage output was nested
+        # under "request_output". Merge its content via the compat setter.
+        legacy_inner = data.pop("request_output", None)
+        multi_modal_placeholders = data.pop("multi_modal_placeholders", None)
+        field_names = {f.name for f in fields(OmniRequestOutput)}
+        known = {k: v for k, v in data.items() if k in field_names}
         try:
-            # Use msgspec.convert for dataclass reconstruction
-            return msgspec.convert(obj, OmniRequestOutput)
+            out = OmniRequestOutput(**known)
         except Exception:
-            try:
-                # Fallback: construct directly if msgspec.convert fails
-                # (e.g., if some fields are missing or have wrong types)
-                return OmniRequestOutput(**obj)
-            except Exception:
-                # If both attempts fail, return dict as-is (defensive fallback)
-                # This should rarely happen if _is_omni_request_output is correct
-                return obj
+            # Defensive fallback; should rarely happen if
+            # _is_omni_request_output is correct.
+            return obj
+        if legacy_inner is not None and not isinstance(legacy_inner, dict):
+            out.request_output = legacy_inner
+        if multi_modal_placeholders is not None:
+            setattr(out, "multi_modal_placeholders", multi_modal_placeholders)
+        return out
 
     def _decode_tensor(self, obj: dict[str, Any]) -> torch.Tensor:
         """Decode dict to torch.Tensor."""

vllm_omni/entrypoints/cli/serve.py

@@ -148,8 +148,7 @@ class OmniServeCommand(CLISubcommand):
                     "The following CLI args are not supported under --omni: "
                     f"{', '.join(offenders)}. Configure parallelism through the "
                     "per-stage YAML (`--deploy-config` / `--stage-configs-path`) "
-                    "and replica count via the per-stage `num_replicas` config field "
-                    "(single-runtime) or `--omni-dp-size-local` (headless / multi-runtime)."
+                    "and replica count via `--omni-dp-size-local`."
                 )
 
         # --omni-lb-policy is validated against the LoadBalancingPolicy enum.

vllm_omni/entrypoints/omni_base.py

@@ -577,7 +577,6 @@ class OmniBase(PDDisaggregationMixin):
         self.prom_metrics.set_running(running)
         self.prom_metrics.set_waiting(max(0, total - running))
 
-        images = getattr(engine_outputs, "images", []) if output_type == "image" else []
         response_metrics: dict[str, Any] = {}
         stage_metrics: dict[str, dict[str, Any]] = {}
         rid_key = str(req_id)
@@ -602,17 +601,14 @@ class OmniBase(PDDisaggregationMixin):
                 response_metrics["final_output_type"] = current_stage_metrics["final_output_type"]
                 response_metrics["num_tokens_in"] = current_stage_metrics["num_tokens_in"]
                 response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"]
+        # Generation content (outputs, prompt, images, trajectory_*, ...) is
+        # copied from engine_outputs onto the returned object by
+        # OmniRequestOutput.__post_init__.
         return OmniRequestOutput(
             request_id=req_id or "",
             stage_id=stage_id,
             final_output_type=output_type,
             request_output=engine_outputs,
-            images=images,
-            trajectory_latents=getattr(engine_outputs, "trajectory_latents", None),
-            trajectory_timesteps=getattr(engine_outputs, "trajectory_timesteps", None),
-            trajectory_log_probs=getattr(engine_outputs, "trajectory_log_probs", None),
-            trajectory_decoded=getattr(engine_outputs, "trajectory_decoded", None),
-            _custom_output=getattr(engine_outputs, "_custom_output", {}),
             metrics=response_metrics,
             stage_durations=stage_durations,
             peak_memory_mb=peak_memory_mb,

vllm_omni/entrypoints/openai/serving_chat.py

@@ -3113,7 +3113,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 request_id=request_id,
             )
 
-        images = getattr(result.request_output, "images", [])
+        images = getattr(result, "images", [])
         stage_durations = result.stage_durations
         peak_memory_mb = result.peak_memory_mb
         cot_output = None
@@ -3184,7 +3184,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         )
                         yield f"data: {chunk.model_dump_json()}\n\n"
                 elif final_output_type == "image":
-                    images = self._flatten_diffusion_images(getattr(output.request_output, "images", []))
+                    images = self._flatten_diffusion_images(getattr(output, "images", []))
                     if not images:
                         raise RuntimeError("Streaming image edit produced an empty final image output.")
                     image_data = [
@@ -3488,8 +3488,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
             # Image output path (text2img / img2img)
             final_output_type = getattr(result, "final_output_type", "image")
-            # Handle nested OmniRequestOutput structure where images might be in request_output
-            images = getattr(result.request_output, "images", [])
+            images = getattr(result, "images", [])
             multimodal_output = getattr(result, "multimodal_output", {}) or {}
             stage_durations = result.stage_durations
             peak_memory_mb = result.peak_memory_mb

vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py

@@ -302,11 +302,6 @@ class Qwen3TTSPromptEmbedsBuilder:
             artifacts. ``0`` disables the cache.
     """
 
-    # The embedding path always runs in bf16 regardless of the configured
-    # model dtype (mirrors ``Qwen3TTSTalkerForConditionalGeneration``).
-    # Class attribute so test stubs built via ``__new__`` inherit it.
-    _embedding_dtype = torch.bfloat16
-
     def __init__(
         self,
         *,
@@ -334,6 +329,7 @@ class Qwen3TTSPromptEmbedsBuilder:
         self._tts_pad_embed_buffer = tts_pad_embed
         self._encode_ref_audio_batch_fn = encode_ref_audio_batch
         self._speaker_cache = speaker_cache
+        self._embedding_dtype = torch.bfloat16
 
         self._text_tokenizer: Any | None = None
 

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -266,11 +266,6 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
     """vLLM-AR talker: step-wise layer-0 codec decoding.
     Predicts residual codebooks (1..Q-1) into `audio_codes` and streams text via `tailing_text_hidden`."""
 
-    # The embedding path always runs in bf16 regardless of the configured
-    # model dtype (encoder weights are cast to bf16 in ``__init__``).  Class
-    # attribute so test stubs built via ``__new__`` inherit it.
-    _embedding_dtype = torch.bfloat16
-
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             # Talker backbone (Qwen3 decoder-only).
@@ -448,6 +443,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             torch.zeros(1, int(self.talker_config.hidden_size), dtype=model_dtype),
             persistent=False,
         )
+        self._embedding_dtype = torch.bfloat16
 
         tokenizer_config = Qwen3TTSTokenizerV2Config.from_pretrained(
             self.model_path,

vllm_omni/model_extras/registry.py

@@ -153,11 +153,6 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
         "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt,
     },
-    "Cosmos3OmniPipeline": {
-        "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS,
-        "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS,
-        "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt,
-    },
     "MagiHumanPipeline": {
         "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS,
         "extra_output_params": MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,

vllm_omni/outputs/__init__.py

@@ -3,7 +3,7 @@ from typing import Any
 
 import torch
 from PIL import Image
-from vllm.outputs import RequestOutput
+from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.v1.outputs import ModelRunnerOutput
 
 from vllm_omni.inputs.data import OmniPromptType
@@ -59,12 +59,58 @@ class OmniModelRunnerOutput(ModelRunnerOutput):
     omni_connector_output: OmniConnectorOutput | None = None
 
 
+# RequestOutput attributes that hold generation content.  When wrapping a
+# stage output, these are copied onto the OmniRequestOutput itself so that
+# the object IS the RequestOutput instead of holding one.
+_REQUEST_OUTPUT_CONTENT_ATTRS = (
+    "prompt",
+    "prompt_token_ids",
+    "prompt_logprobs",
+    "outputs",
+    "finished",
+    "lora_request",
+    "encoder_prompt",
+    "encoder_prompt_token_ids",
+    "num_cached_tokens",
+    "kv_transfer_params",
+)
+
+# Omni-specific content copied when the wrapped stage output is itself an
+# OmniRequestOutput (e.g. a diffusion stage inside a multi-stage pipeline).
+_OMNI_CONTENT_ATTRS = (
+    "images",
+    "latents",
+    "trajectory_latents",
+    "trajectory_timesteps",
+    "trajectory_log_probs",
+    "trajectory_decoded",
+    "_multimodal_output",
+    "_custom_output",
+)
+
+# Dynamic attributes that must not be copied verbatim onto the wrapper:
+# they collide with properties defined on OmniRequestOutput.
+_COPY_SKIP_ATTRS = frozenset({"request_output", "multimodal_output", "custom_output"})
+
+
 @dataclass
-class OmniRequestOutput:
+class OmniRequestOutput(RequestOutput):
     """Unified request output for both pipeline stages and diffusion models.
 
+    Extends vLLM's ``RequestOutput`` so that omni outputs can flow directly
+    through vLLM serving codepaths (which expect ``prompt_token_ids``,
+    ``outputs``, etc. as real attributes).  The inherited fields store the
+    LLM generation content; omni-specific fields store pipeline/diffusion
+    extras.
+
+    Note: ``RequestOutput`` is a plain class (not a dataclass), so all of its
+    attributes are redeclared below as dataclass fields with defaults — the
+    dataclass-generated ``__init__`` replaces ``RequestOutput.__init__`` and
+    must set them itself.
+
     This class handles outputs from:
-    1. Multi-stage LLM pipelines (with stage_id, final_output_type, request_output)
+    1. Multi-stage LLM pipelines (with stage_id, final_output_type, and the
+       inherited RequestOutput fields carrying the stage's generation content)
     2. Diffusion models (with images, prompt, metrics)
 
     Attributes:
@@ -72,30 +118,38 @@ class OmniRequestOutput:
         finished: Whether generation is complete
         stage_id: Identifier of the stage that produced this output (pipeline mode)
         final_output_type: Type of output ("text", "image", "audio", "latents")
-        request_output: The underlying RequestOutput from the stage (pipeline mode)
         images: List of generated PIL images (diffusion mode)
-        prompt: The prompt used for generation (diffusion mode)
+        prompt: The prompt used for generation
         latents: Optional tensor of latent representations (diffusion mode)
-        metrics: Optional dictionary of generation metrics
+        metrics: Generation metrics. A plain dict for omni outputs; may carry
+            vLLM's request stats object when copied from a raw RequestOutput.
     """
 
+    # --- Inherited RequestOutput attributes (redeclared as fields) ---
     request_id: str = ""
+    prompt: OmniPromptType | None = None
+    prompt_token_ids: list[int] | None = None
+    prompt_logprobs: Any = None
+    outputs: list[CompletionOutput] = field(default_factory=list)
     finished: bool = True
+    metrics: Any = field(default_factory=dict)
+    lora_request: Any = None
+    encoder_prompt: str | None = None
+    encoder_prompt_token_ids: list[int] | None = None
+    num_cached_tokens: int | None = None
+    kv_transfer_params: dict[str, Any] | None = None
 
-    # Pipeline stage fields
+    # --- Pipeline stage fields ---
     stage_id: int | None = None
     final_output_type: str = "text"
-    request_output: RequestOutput | None = None
 
-    # Diffusion model fields
+    # --- Diffusion model fields ---
     images: list[Image.Image] = field(default_factory=list)
-    prompt: OmniPromptType | None = None
     latents: torch.Tensor | None = None
     trajectory_latents: torch.Tensor | None = None
     trajectory_timesteps: torch.Tensor | None = None
     trajectory_log_probs: torch.Tensor | None = None
     trajectory_decoded: list | None = None
-    metrics: dict[str, Any] = field(default_factory=dict)
     _multimodal_output: dict[str, Any] = field(default_factory=dict)
     _custom_output: dict[str, Any] = field(default_factory=dict)
 
@@ -110,6 +164,44 @@ class OmniRequestOutput:
     error_status_code: int | None = None
     error_type: str | None = None
 
+    def _copy_content_from(self, source: Any) -> None:
+        """Copy generation content from a stage output into this object.
+
+        ``source`` may be a vLLM ``RequestOutput``, another
+        ``OmniRequestOutput``, or a duck-typed stand-in (tests use mocks and
+        SimpleNamespace). Attribute presence is checked against the source's
+        instance ``__dict__`` so that auto-created Mock attributes are not
+        copied. Control fields set by the wrapper (stage_id,
+        final_output_type, metrics, stage_durations, ...) are intentionally
+        not copied.
+        """
+        source_vars = getattr(source, "__dict__", None)
+
+        def _has(name: str) -> bool:
+            if source_vars is not None:
+                return name in source_vars
+            return hasattr(source, name)
+
+        for name in _REQUEST_OUTPUT_CONTENT_ATTRS + _OMNI_CONTENT_ATTRS:
+            if _has(name):
+                setattr(self, name, getattr(source, name))
+        if not self.request_id:
+            self.request_id = getattr(source, "request_id", "") or ""
+        # Dynamic attribute that vLLM code attaches to RequestOutput; it is a
+        # read-only property on this class, so store it in _multimodal_output.
+        if not isinstance(source, OmniRequestOutput) and _has("multimodal_output"):
+            mm = source.multimodal_output
+            if mm is not None:
+                self._multimodal_output = mm
+        # Preserve remaining dynamic attributes (e.g. dynamically annotated
+        # stage outputs) so they stay reachable through `request_output`.
+        if source_vars:
+            known = type(self).__dataclass_fields__.keys()
+            for name, value in source_vars.items():
+                if name.startswith("_") or name in known or name in _COPY_SKIP_ATTRS:
+                    continue
+                setattr(self, name, value)
+
     @classmethod
     def from_error(
         cls,
@@ -143,7 +235,10 @@ class OmniRequestOutput:
         final_output_type: str,
         request_output: RequestOutput,
     ) -> "OmniRequestOutput":
-        """Create output from pipeline stage.
+        """Create output from a pipeline stage.
+
+        The stage output's content is copied into the new object's own
+        (inherited) fields.
 
         Args:
             stage_id: Stage identifier
@@ -153,13 +248,13 @@ class OmniRequestOutput:
         Returns:
             OmniRequestOutput configured for pipeline mode
         """
-        return cls(
-            request_id=getattr(request_output, "request_id", ""),
+        out = cls(
             stage_id=stage_id,
             final_output_type=final_output_type,
             request_output=request_output,
-            finished=True,
         )
+        out.finished = True
+        return out
 
     @classmethod
     def from_diffusion(
@@ -220,35 +315,21 @@ class OmniRequestOutput:
 
     @property
     def multimodal_output(self) -> Any:
-        """Return multimodal output from the underlying request output or local field.
+        """Return the multimodal output payload.
 
-        For pipeline outputs, this checks completion outputs first, then request_output.
-        For diffusion outputs, this returns the local _multimodal_output field.
+        Checks completion outputs first (where multimodal_output is attached
+        by AR stages), then the local _multimodal_output field.
 
         Returns either a MultimodalPayload (Phase 3+) or a plain dict (legacy).
         """
-        if self.request_output is None:
-            return self._multimodal_output
-
-        # Check completion outputs first (where multimodal_output is attached)
-        for output in getattr(self.request_output, "outputs", []):
+        for output in self.outputs:
             if mm := getattr(output, "multimodal_output", None):
                 return mm
-        if mm := getattr(self.request_output, "multimodal_output", None):
-            return mm
         return self._multimodal_output
 
     @property
     def custom_output(self) -> dict[str, Any]:
-        """Return custom output data from diffusion pipelines.
-
-        For diffusion outputs, returns the local _custom_output field.
-        For pipeline outputs with an inner OmniRequestOutput, forwards
-        the custom_output from the inner request output.
-        """
-        if self.request_output is not None:
-            if isinstance(self.request_output, OmniRequestOutput):
-                return self.request_output._custom_output
+        """Return custom output data from diffusion pipelines."""
         return self._custom_output
 
     @custom_output.setter
@@ -260,58 +341,6 @@ class OmniRequestOutput:
         """Return the number of generated images."""
         return len(self.images)
 
-    # Pass-through properties keep vLLM serving codepaths compatible with
-    # OmniRequestOutput for pipeline outputs (Issue #345).
-    @property
-    def prompt_token_ids(self) -> list[int] | None:
-        """Return prompt token IDs from the underlying request output.
-
-        This property is required for compatibility with vLLM's streaming
-        chat completion generator which checks res.prompt_token_ids.
-        """
-        if self.request_output is not None:
-            return getattr(self.request_output, "prompt_token_ids", None)
-        return None
-
-    @property
-    def outputs(self) -> list[Any]:
-        """Return outputs from the underlying request output.
-
-        This property is required for compatibility with vLLM's streaming
-        and non-streaming chat completion generators.
-        """
-        if self.request_output is not None:
-            return getattr(self.request_output, "outputs", [])
-        return []
-
-    @property
-    def encoder_prompt_token_ids(self) -> list[int] | None:
-        """Return encoder prompt token IDs from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "encoder_prompt_token_ids", None)
-        return None
-
-    @property
-    def prompt_logprobs(self) -> Any:
-        """Return prompt logprobs from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "prompt_logprobs", None)
-        return None
-
-    @property
-    def num_cached_tokens(self) -> int | None:
-        """Return number of cached tokens from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "num_cached_tokens", None)
-        return None
-
-    @property
-    def kv_transfer_params(self) -> Any:
-        """Return KV transfer params from the underlying request output."""
-        if self.request_output is not None:
-            return getattr(self.request_output, "kv_transfer_params", None)
-        return None
-
     @property
     def is_diffusion_output(self) -> bool:
         """Check if this is a diffusion model output."""
@@ -320,60 +349,33 @@ class OmniRequestOutput:
     @property
     def is_pipeline_output(self) -> bool:
         """Check if this is a pipeline stage output."""
-        return self.stage_id is not None and self.request_output is not None
+        return self.stage_id is not None
 
     def unwrap(self) -> "OmniRequestOutput":
-        """Unwrap nested OmniRequestOutput to get the innermost result.
+        """Return the output itself.
 
-        This helper handles the common pattern where pipeline outputs may wrap
-        other OmniRequestOutput instances. It recursively unwraps until it reaches
-        the final output with actual content (images, text, etc.).
-
-        Returns:
-            The innermost OmniRequestOutput containing the actual generation results.
-
-        Example:
-            ```python
-            result = omni.generate(...)
-            output = OmniRequestOutput.unwrap_result(result)
-            if output.images:
-                # Access images directly
-                video_frames = output.images
-            ```
+        Kept for backward compatibility: outputs used to nest a stage's
+        OmniRequestOutput inside ``request_output``; content is now flattened
+        onto the object itself, so there is nothing to unwrap.
         """
-        current = self
-        # Unwrap nested pipeline outputs
-        while current.is_pipeline_output and current.request_output is not None:
-            if isinstance(current.request_output, OmniRequestOutput):
-                current = current.request_output
-            else:
-                break
-        return current
+        return self
 
     @staticmethod
     def unwrap_result(result: Any) -> "OmniRequestOutput":
-        """Unwrap result from omni.generate() to get the final OmniRequestOutput.
+        """Extract the final OmniRequestOutput from omni.generate() results.
 
-        This static helper handles the full unwrapping pattern including:
+        Handles:
         1. Extracting from list if needed
         2. Type validation
-        3. Recursive unwrapping of nested pipeline outputs
 
         Args:
             result: The result from omni.generate() - may be a list or OmniRequestOutput
 
         Returns:
-            The innermost OmniRequestOutput with actual content
+            The OmniRequestOutput with actual content
 
         Raises:
             ValueError: If result is not an OmniRequestOutput or list containing one
-
-        Example:
-            ```python
-            result = omni.generate(...)
-            output = OmniRequestOutput.unwrap_result(result)
-            # output is guaranteed to be the final OmniRequestOutput
-            ```
         """
         # Handle list wrapper
         if isinstance(result, list):
@@ -385,7 +387,6 @@ class OmniRequestOutput:
         if not isinstance(result, OmniRequestOutput):
             raise ValueError(f"Expected OmniRequestOutput, got {type(result)}")
 
-        # Unwrap nested outputs
         return result.unwrap()
 
     def to_dict(self) -> dict[str, Any]:
@@ -425,7 +426,8 @@ class OmniRequestOutput:
             f"finished={self.finished}",
             f"stage_id={self.stage_id}",
             f"final_output_type={self.final_output_type!r}",
-            f"request_output={self.request_output}",
+            f"prompt_token_ids={self.prompt_token_ids}",
+            f"outputs={self.outputs}",
             f"images={images_repr}",
             f"prompt={self.prompt!r}",
             f"latents={self.latents}",
@@ -437,3 +439,45 @@ class OmniRequestOutput:
         ]
 
         return f"OmniRequestOutput({', '.join(parts)})"
+
+
+def _request_output_get(self: OmniRequestOutput) -> RequestOutput | None:
+    """Backward-compat view of the wrapped stage output.
+
+    The omni output now IS the RequestOutput, so pipeline/LLM outputs return
+    ``self``; pure diffusion or error outputs return None (matching the old
+    composition semantics where ``request_output`` was unset for them).
+    """
+    if self.stage_id is not None or self.outputs:
+        return self
+    return None
+
+
+def _request_output_set(self: OmniRequestOutput, value: Any) -> None:
+    if value is not None:
+        self._copy_content_from(value)
+
+
+# `request_output` is attached after class creation, NOT declared as a field
+# or InitVar: msgspec serializes dataclasses natively via __dataclass_fields__
+# (which includes InitVar pseudo-fields), so a request_output entry would be
+# read through the property during encoding — returning `self` and recursing
+# forever. As a plain property it is invisible to both dataclasses and
+# msgspec.
+OmniRequestOutput.request_output = property(  # type: ignore[assignment]
+    _request_output_get, _request_output_set
+)
+
+# Keep `request_output=` working as a constructor keyword for
+# backward compatibility: the wrapped stage output's content is copied onto
+# the new object (see _copy_content_from).
+_dataclass_init = OmniRequestOutput.__init__
+
+
+def _omni_request_output_init(self: OmniRequestOutput, *args: Any, request_output: Any = None, **kwargs: Any) -> None:
+    _dataclass_init(self, *args, **kwargs)
+    if request_output is not None:
+        self._copy_content_from(request_output)
+
+
+OmniRequestOutput.__init__ = _omni_request_output_init  # type: ignore[method-assign]

vllm_omni/quantization/quack_fp8.py

@@ -94,7 +94,7 @@ def quack_scaled_fp8_mm(
         return None
     out = torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=out_dtype)
     alpha = scale_a.reshape(1).float() * scale_b.reshape(1).float()
-    gemm.gemm(a, b, out=out, bias=bias, alpha=alpha, tuned=True)
+    gemm.gemm_tuned(a, b, out, bias=bias, alpha=alpha)
     return out