[WIP] Optimize Helios accuracy-maintained path

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

.buildkite/test-ready.yml

@@ -457,6 +457,51 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: "Diffusion · Helios Distilled Acc/Perf Compare"
+        commands:
+          - |
+            timeout 90m bash -c "
+              set -euxo pipefail
+              export VLLM_IMAGE_FETCH_TIMEOUT=60
+              pytest -s -v tests/e2e/accuracy/test_helios_distilled.py::test_helios_distilled_pr_matches_main_acc_perf \
+                -m 'full_model and benchmark' \
+                --run-level 'full_model'
+            "
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/helios_distilled_acc_perf/result.json"
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/helios_distilled_acc_perf/artifacts/**/*.mp4"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
 
       - label: "Diffusion · Qwen Image Test"
         commands:

benchmarks/diffusion/backends.py

@@ -278,7 +278,7 @@ async def async_request_v1_videos(
                 return output
 
         # invoke a poll request (GET /v1/videos/{video_id})
-        poll_interval = 2.0  # Unit(s)
+        poll_interval = 0.2  # Unit(s)
         timeout_seconds = 600.0
         deadline = time.perf_counter() + timeout_seconds
         job_url = f"{input.api_url}/{job_id}"

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -68,6 +68,7 @@ import glob
 import json
 import logging
 import os
+import pathlib
 import random
 import tempfile
 import time
@@ -598,6 +599,13 @@ class RandomDataset(BaseDataset):
         if self.enable_negative_prompt:
             extra_body["negative_prompt"] = f"Negative prompt {idx} for benchmarking diffusion models"
 
+        request_fields = {
+            "width",
+            "height",
+            "num_frames",
+            "num_inference_steps",
+            "fps",
+        }
         params = {
             "width": self.args.width,
             "height": self.args.height,
@@ -607,9 +615,11 @@ class RandomDataset(BaseDataset):
         }
         if self._sampled_requests:
             profile = self._sampled_requests[idx]
-            params.update(profile)
+            params.update({key: value for key, value in profile.items() if key in request_fields})
+            extra_body.update({key: value for key, value in profile.items() if key not in request_fields})
+        prompt = self.args.random_prompt or f"Random prompt {idx} for benchmarking diffusion models"
         return RequestFuncInput(
-            prompt=f"Random prompt {idx} for benchmarking diffusion models",
+            prompt=prompt,
             api_url=self.api_url,
             model=self.model,
             seed=self.args.seed,
@@ -869,6 +879,35 @@ def calculate_metrics(
     return metrics
 
 
+def _response_suffix(endpoint: str) -> str:
+    if endpoint == "/v1/videos":
+        return ".mp4"
+    return ".json"
+
+
+def _save_success_responses(
+    outputs: list[RequestFuncOutput],
+    requests_list: list[RequestFuncInput],
+    output_dir: str,
+    endpoint: str,
+) -> list[str]:
+    path = pathlib.Path(output_dir).expanduser().resolve()
+    path.mkdir(parents=True, exist_ok=True)
+    saved_paths: list[str] = []
+    suffix = _response_suffix(endpoint)
+
+    for index, (request, output) in enumerate(zip(requests_list, outputs, strict=False)):
+        if not output.success:
+            continue
+        item_path = path / f"{index:05d}_{request.request_id}{suffix}"
+        if isinstance(output.response_body, bytes):
+            item_path.write_bytes(output.response_body)
+        else:
+            item_path.write_text(json.dumps(output.response_body, indent=2, sort_keys=True), encoding="utf-8")
+        saved_paths.append(str(item_path))
+    return saved_paths
+
+
 def wait_for_service(base_url: str, timeout: int = 120) -> None:
     print(f"Waiting for service at {base_url}...")
     start_time = time.time()
@@ -998,6 +1037,13 @@ async def benchmark(args):
     metrics["model"] = args.model
     metrics["dataset"] = args.dataset
     metrics["task"] = args.task
+    if args.save_response_dir:
+        metrics["saved_responses"] = _save_success_responses(
+            outputs,
+            requests_list,
+            args.save_response_dir,
+            args.endpoint,
+        )
 
     print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
 
@@ -1154,6 +1200,12 @@ if __name__ == "__main__":
     )
     parser.add_argument("--fps", type=int, default=None, help="FPS (for video).")
     parser.add_argument("--output-file", type=str, default=None, help="Output JSON file for metrics.")
+    parser.add_argument(
+        "--save-response-dir",
+        type=str,
+        default=None,
+        help="Directory to save successful raw responses. Video responses are saved as MP4 files.",
+    )
     parser.add_argument(
         "--slo",
         action="store_true",
@@ -1189,6 +1241,12 @@ if __name__ == "__main__":
             '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
         ),
     )
+    parser.add_argument(
+        "--random-prompt",
+        type=str,
+        default=None,
+        help="Prompt to use for every synthetic random request. Defaults to generated placeholder prompts.",
+    )
     parser.add_argument(
         "--num-input-images",
         type=int,

benchmarks/diffusion/performance_dashboard/helios_distilled_serving_performance.md

@@ -0,0 +1,222 @@
+# Helios-Distilled Serving Performance Dashboard
+
+This document describes how to deploy and benchmark **BestWishYsh/Helios-Distilled** using vLLM-Omni. It follows the same serving benchmark structure as the Qwen-Image and Wan2.2 dashboards.
+
+---
+
+# 1. Overview
+
+Helios-Distilled is a text-to-video diffusion model. Helios uses an autoregressive chunking pattern that generates 33 frames per chunk. For stable comparisons, choose `num_frames` as a multiple of 33, such as 99.
+
+This document covers:
+
+* Service launch configuration
+* Benchmark scripts and usage
+* Dataset and workload settings
+* Performance and accuracy comparison methodology
+* Reproducibility guidelines
+
+---
+
+# 2. Test Environment
+
+| Component | Specification |
+|------------|----------------|
+| GPU | NVIDIA H100 for CI accuracy/perf gate; local bring-up also used L20X |
+| Diffusion Attention Backend | FlashAttention |
+| Model | `BestWishYsh/Helios-Distilled` |
+
+---
+
+# 3. Service Launch Configuration
+
+## 3.1 Basic Serving Command
+
+```bash
+vllm serve BestWishYsh/Helios-Distilled --omni \
+    --port 8091
+```
+
+## 3.2 Key Parameters
+
+| Parameter | Description |
+| --- | --- |
+| `--port` | Serving port |
+| `--max-num-seqs` | Maximum number of active sequences |
+| `--enable-diffusion-pipeline-profiler` | Optional stage-level profiling |
+| `--profiler-config` | Optional torch profiler configuration |
+
+Record these parameters when reporting performance results.
+
+---
+
+# 4. Benchmark Scripts
+
+## 4.1 Online Serving Benchmark
+
+Use the shared diffusion serving benchmark, as with Qwen-Image and Wan2.2:
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --base-url http://localhost:8091 \
+    --endpoint /v1/videos \
+    --model BestWishYsh/Helios-Distilled \
+    --dataset random \
+    --task t2v \
+    --num-prompts 1 \
+    --max-concurrency 1 \
+    --disable-tqdm \
+    --random-request-config '[
+        {
+            "width": 640,
+            "height": 384,
+            "num_inference_steps": 50,
+            "num_frames": 99,
+            "fps": 16,
+            "guidance_scale": 1.0,
+            "is_enable_stage2": true,
+            "pyramid_num_stages": 3,
+            "pyramid_num_inference_steps_list": [1, 1, 1],
+            "is_amplify_first_chunk": false,
+            "weight": 1
+        }
+    ]'
+```
+
+`--random-request-config` forwards standard fields such as `width`, `height`, `num_frames`, `fps`, and `num_inference_steps` as first-class request fields. Model-specific fields are forwarded in the request body, which is required for Helios stage-2 settings.
+
+## 4.2 Warmed Serving Benchmark
+
+For local bring-up, use the same shared serving benchmark with one warmup request and one measured request:
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --base-url http://localhost:8091 \
+    --endpoint /v1/videos \
+    --model BestWishYsh/Helios-Distilled \
+    --dataset random \
+    --task t2v \
+    --num-prompts 1 \
+    --max-concurrency 1 \
+    --warmup-requests 1 \
+    --warmup-num-inference-steps 1 \
+    --random-prompt "A cat wearing sunglasses dances on a beach at sunset, cinematic lighting." \
+    --seed 42 \
+    --disable-tqdm \
+    --output-file metrics.json \
+    --save-response-dir responses \
+    --random-request-config '[
+        {
+            "width": 640,
+            "height": 384,
+            "num_inference_steps": 50,
+            "num_frames": 99,
+            "fps": 16,
+            "guidance_scale": 1.0,
+            "is_enable_stage2": true,
+            "pyramid_num_stages": 3,
+            "pyramid_num_inference_steps_list": [1, 1, 1],
+            "is_amplify_first_chunk": false,
+            "weight": 1
+        }
+    ]'
+```
+
+## 4.3 PR-vs-Main Accuracy and Performance Comparison
+
+The PR-vs-main accuracy/performance gate lives in the Helios accuracy pytest. It checks out `origin/main`, starts a server for each checkout, and drives both with `diffusion_benchmark_serving.py`:
+
+```bash
+pytest -s -v tests/e2e/accuracy/test_helios_distilled.py::test_helios_distilled_pr_matches_main_acc_perf \
+    -m 'full_model and benchmark' \
+    --run-level 'full_model'
+```
+
+---
+
+# 5. Dataset & Workload Settings
+
+## 5.1 Recommended Evaluation Configurations
+
+### Dataset A (99 frames, 384x640)
+
+* Dataset: `random`
+* Task: t2v
+* Concurrency: 1
+* Frame count: 99 (`33 x 3`)
+* FPS: 16
+* Stage-2 steps: `[1, 1, 1]`
+
+```json
+[
+  {
+    "width": 640,
+    "height": 384,
+    "num_inference_steps": 50,
+    "num_frames": 99,
+    "fps": 16,
+    "guidance_scale": 1.0,
+    "is_enable_stage2": true,
+    "pyramid_num_stages": 3,
+    "pyramid_num_inference_steps_list": [1, 1, 1],
+    "is_amplify_first_chunk": false,
+    "weight": 1
+  }
+]
+```
+
+---
+
+# 6. Performance Metrics
+
+| Metric | Description | Unit |
+| --- | --- | --- |
+| Mean Latency | Mean request latency | seconds |
+| P99 Latency | P99 request latency | seconds |
+| Generated FPS | `num_frames / measured_time_s` | frames/second |
+| Candidate/reference latency ratio | PR latency divided by reference latency | ratio |
+
+---
+
+# 7. Accuracy Metrics
+
+The PR-vs-main gate compares encoded video frames extracted from MP4 outputs:
+
+| Metric | Threshold | Direction |
+| --- | --- | --- |
+| PSNR | `>= 35 dB` | Higher is better |
+| MAE | `<= 3.0 / 255` | Lower is better |
+| Candidate/reference latency ratio | `<= 1.03` | Lower is better |
+
+Local L20X bring-up for this PR produced:
+
+| Variant | Measured Latency (ms) | Generated FPS |
+| --- | --- | --- |
+| `origin/main` | 5317.95 | 18.62 |
+| Candidate PR | 5283.27 | 18.74 |
+
+Encoded-video similarity:
+
+| Metric | Value |
+| --- | --- |
+| PSNR | 38.66 dB |
+| MAE | 1.89 / 255 |
+| Cosine similarity | 0.99978 |
+
+---
+
+# 8. Reproducibility Checklist
+
+To ensure consistent and comparable benchmark results:
+
+* Record GPU type
+* Record vLLM and vLLM-Omni versions
+* Record benchmark parameters: resolution, frame count, FPS, seed, concurrency, and stage-2 steps
+* Use frame counts that are multiples of 33
+* Keep prompt, seed, resolution, and sampling parameters identical when comparing PR vs main
+* Ensure no background workload on GPUs during testing
+* Use the same MP4 export and frame extraction path for both variants
+
+---
+
+This document serves as the Helios-Distilled serving performance reference under vLLM-Omni.

docs/community/governance.md

@@ -39,34 +39,11 @@ Committers have write access and merge rights. They typically have deep expertis
 - [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
 - [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
 - [@princepride](https://github.com/princepride): Diffusion and Omni Support
-- [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs
+- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
-- [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache
 - [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
-- [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs
+- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
 - [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
-- [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp
+- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
-- [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support
-
-## Diffusion Workload Division
-
-This section breaks down diffusion-related responsibilities. If you have PRs touching these areas, please ping the listed owners for review.
-
-### Runtime-related
-
-- Diffusion models: @RuixiangMa, @wtomin
-
-### Optimization
-
-- Parallel: @RuixiangMa, @wtomin
-- Attention backend: @ZJY0516
-- Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang
-- Offload: @yuanheng-zhao
-- Quantization: @lishunyang12, @david6666666
-- Kernel fusion / communication-computation: @ZJY0516
-
-### Docs & Test
-
-- Diffusion docs: @RuixiangMa, @wtomin
 
 ## Meetings
 

tests/dfx/conftest.py

@@ -279,10 +279,6 @@ def create_benchmark_indices(
     return indices
 
 
-# Omni perf/stability: ``vllm bench serve --omni`` result template (used if bench exits without writing JSON).
-OMNI_RESULT_TEMPLATE_PATH = Path(__file__).resolve().parent / "perf" / "scripts" / "result_omni_template.json"
-
-
 def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
     """Make a single path segment safe for result filenames on common filesystems."""
     if value is None:
@@ -293,7 +289,7 @@ def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
     return s if s else default
 
 
-def resolve_baseline_value(
+def _resolve_baseline_value(
     baseline_raw: Any,
     *,
     sweep_index: int | None,
@@ -334,7 +330,7 @@ def _baseline_thresholds_for_step(
 ) -> dict[str, Any]:
     """Resolve baseline config to one threshold per metric for this iteration."""
     return {
-        metric_name: resolve_baseline_value(
+        metric_name: _resolve_baseline_value(
             baseline_raw,
             sweep_index=sweep_index,
             max_concurrency=max_concurrency,
@@ -358,12 +354,7 @@ def run_benchmark(
     random_input_len: Any | None = None,
     random_output_len: Any | None = None,
 ) -> dict[str, Any]:
-    """Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
+    """Run one ``vllm bench serve --omni`` iteration and return parsed metrics."""
-
-    After ``vllm bench`` writes the JSON, ``result["baseline"]`` holds the resolved per-metric thresholds
-    (when ``baseline_config`` is provided). If the benchmark exits without writing a result file,
-    ``result_omni_template.json`` is used as a fallback.
-    """
     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)
@@ -412,17 +403,8 @@ def run_benchmark(
         result_dir = "./"
 
     result_path = os.path.join(result_dir, result_filename)
-    if not os.path.exists(result_path):
+    with open(result_path, encoding="utf-8") as f:
-        with open(OMNI_RESULT_TEMPLATE_PATH, encoding="utf-8") as f:
+        result = json.load(f)
-            template_result: dict[str, Any] = json.load(f)
-        Path(result_path).parent.mkdir(parents=True, exist_ok=True)
-        with open(result_path, "w", encoding="utf-8") as f:
-            json.dump(template_result, f, ensure_ascii=False, indent=2)
-        print(f"Benchmark result file not generated, fallback to template: {result_path}")
-        result = template_result
-    else:
-        with open(result_path, encoding="utf-8") as f:
-            result = json.load(f)
 
     if baseline_config:
         result["baseline"] = _baseline_thresholds_for_step(
@@ -450,12 +432,3 @@ def pytest_addoption(parser: pytest.Parser) -> None:
         default=None,
         help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"),
     )
-    parser.addoption(
-        "--assert-baseline",
-        action="store_true",
-        default=False,
-        help=(
-            "When set, omni/diffusion perf runners compare metrics against the baseline block in the JSON config "
-            "(default: off)."
-        ),
-    )

tests/dfx/perf/scripts/run_benchmark.py

@@ -12,17 +12,12 @@ from tests.dfx.conftest import (
     create_unique_server_params,
     get_benchmark_params_for_server,
     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.
-
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
@@ -51,6 +46,7 @@ if CONFIG_FILE_PATH is None:
     CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE
 
 BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
+OMNI_RESULT_TEMPLATE_PATH = Path(__file__).parent / "result_omni_template.json"
 
 
 DEPLOY_CONFIGS_DIR = Path(__file__).parent.parent / "deploy"
@@ -91,6 +87,16 @@ def omni_server(request):
         print("OmniServer stopped")
 
 
+def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
+    """Make a single path segment safe for result filenames on common filesystems."""
+    if value is None:
+        return default
+    s = str(value).strip()
+    for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
+        s = s.replace(bad, "_")
+    return s if s else default
+
+
 benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
 
 
@@ -120,23 +126,76 @@ def benchmark_params(request, omni_server):
     }
 
 
+def _resolve_baseline_value(
+    baseline_raw: Any,
+    *,
+    sweep_index: int | None,
+    max_concurrency: Any = None,
+    request_rate: Any = None,
+) -> Any:
+    """Pick the baseline threshold for this sweep step.
+
+    Supported shapes per metric:
+    - **Scalar** — same threshold for every concurrency / QPS.
+    - **List** — aligned with ``max_concurrency`` / ``request_rate`` sweep order; use ``sweep_index``.
+    - **Dict** — keyed by concurrency or rate, e.g. ``{"1": 500, "4": 800}`` (keys are strings in JSON).
+
+    For dict lookup, ``max_concurrency`` is preferred when both are set (concurrency sweep).
+    """
+    if baseline_raw is None:
+        # If no baseline is set, the maximum value will be used.
+        return 100000
+    if isinstance(baseline_raw, dict):
+        if max_concurrency is not None:
+            for key in (max_concurrency, str(max_concurrency)):
+                if key in baseline_raw:
+                    return baseline_raw[key]
+        if request_rate is not None:
+            for key in (request_rate, str(request_rate)):
+                if key in baseline_raw:
+                    return baseline_raw[key]
+        raise KeyError(
+            f"baseline dict has no key for max_concurrency={max_concurrency!r} "
+            f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
+        )
+    if isinstance(baseline_raw, (list, tuple)):
+        return baseline_raw[sweep_index]
+    return baseline_raw
+
+
+def _baseline_thresholds_for_step(
+    baseline_data: dict[str, Any],
+    *,
+    sweep_index: int | None = None,
+    max_concurrency: Any = None,
+    request_rate: Any = None,
+) -> dict[str, Any]:
+    """Resolve ``test.json`` ``baseline`` block to one threshold per metric (same as ``assert_result``)."""
+    return {
+        metric_name: _resolve_baseline_value(
+            baseline_raw,
+            sweep_index=sweep_index,
+            max_concurrency=max_concurrency,
+            request_rate=request_rate,
+        )
+        for metric_name, baseline_raw in baseline_data.items()
+    }
+
+
 def assert_result(
     result,
     params,
     num_prompt,
     *,
-    assert_baseline: bool,
     sweep_index: int | None = None,
-    max_concurrency: Any | None = None,
+    max_concurrency: Any = None,
-    request_rate: Any | None = None,
+    request_rate: Any = None,
 ) -> None:
     assert result["completed"] == num_prompt, "Request failures exist"
-    if not assert_baseline:
-        return
     baseline_data = params.get("baseline", {})
     for metric_name, baseline_raw in baseline_data.items():
         current_value = result[metric_name]
-        baseline_value = resolve_baseline_value(
+        baseline_value = _resolve_baseline_value(
             baseline_raw,
             sweep_index=sweep_index,
             max_concurrency=max_concurrency,
@@ -155,7 +214,7 @@ def assert_result(
 @pytest.mark.benchmark
 @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):
+def test_performance_benchmark(omni_server, benchmark_params):
     test_name = benchmark_params["test_name"]
     params = benchmark_params["params"]
     dataset_name = params.get("dataset_name", "")
@@ -167,8 +226,6 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
     print(f"Running benchmark for model: {model}")
     print(f"Benchmark parameters: {benchmark_params}")
 
-    assert_baseline = request.config.getoption("--assert-baseline", default=False)
-
     def to_list(value, default=None):
         if value is None:
             return [] if default is None else [default]
@@ -207,7 +264,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
         elif not isinstance(value, bool):
             args.extend([arg_name, str(value)])
 
-    # QPS / request-rate sweep
+    # QPS test (sweep_index aligns with qps_list / num_prompt_list for this loop)
     for i, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
         args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
         result = run_benchmark(
@@ -226,13 +283,12 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
         assert_result(
             result,
             params,
-            num_prompt,
+            num_prompt=num_prompt,
-            assert_baseline=assert_baseline,
             sweep_index=i,
             request_rate=qps,
         )
 
-    # concurrency test
+    # concurrency test (sweep_index aligns with max_concurrency_list for separate thresholds per concurrency)
     for i, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
         args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
         result = run_benchmark(
@@ -251,8 +307,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request):
         assert_result(
             result,
             params,
-            num_prompt,
+            num_prompt=num_prompt,
-            assert_baseline=assert_baseline,
             sweep_index=i,
             max_concurrency=concurrency,
         )

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -11,8 +11,6 @@ This runner separates two concepts:
 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).
-
 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,
@@ -654,14 +652,12 @@ def run_benchmark(
 # ---------------------------------------------------------------------------
 
 
-def assert_result(result: dict[str, Any], params: dict[str, Any], *, assert_baseline: bool) -> None:
+def assert_result(result: dict[str, Any], params: dict[str, Any]) -> None:
-    """Assert request completion; optionally assert metrics against ``baseline`` in *params*."""
+    """Assert that benchmark metrics satisfy the configured baselines."""
     num_prompts = params.get("num-prompts", 10)
     completed = result.get("completed_requests", result.get("completed", 0))
     assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}"
 
-    if not assert_baseline:
-        return
     if params.get("skip-performance-assertion", False):
         print("Skipping performance assertions.")
         return
@@ -703,14 +699,16 @@ def _resolve_benchmark_endpoint(server_cfg: dict[str, Any], params: dict[str, An
     indirect=True,
 )
 @pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
-def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request):
+def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
-    """Run the diffusion performance benchmark and verify request completion.
+    """Run the diffusion performance benchmark and assert against baselines.
 
     One server is started per unique parallel configuration (module scope).
     For each server, all benchmark parameter sets defined in the config JSON
-    are executed sequentially; metrics are recorded to the aggregated result file.
+    are executed sequentially; results are asserted against the baselines.
 
-    Pass ``--assert-baseline`` to compare ``throughput_qps``, ``latency_mean``, etc. to the JSON ``baseline`` block.
+    Tracked metrics:
+        - throughput_qps          (higher is better)
+        - latency_p50, latency_p99 (lower is better)
     """
     test_name = benchmark_params["test_name"]
     params = benchmark_params["params"]
@@ -746,5 +744,4 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
     print(f"\n  Aggregated results: {AGGREGATED_RESULT_FILE}")
     print("=" * 60)
 
-    assert_baseline = request.config.getoption("--assert-baseline", default=False)
+    assert_result(result, params)
-    assert_result(result, params, assert_baseline=assert_baseline)

tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json

@@ -3,67 +3,31 @@
         "test_name": "test_qwen3_omni_chunk",
         "server_params": {
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "extra_cli_args": [
+            "extra_cli_args": ["--async-chunk"]
-                "--async-chunk"
-            ]
         },
         "benchmark_params": [
             {
                 "dataset_name": "random",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [4, 16, 32, 64, 128],
-                    4,
+                "max_concurrency": [1, 4, 8, 16, 32],
-                    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": [
+                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000],
-                        400,
+                    "mean_audio_ttfp_ms": [1000, 3000, 5000, 7000, 9000],
-                        300,
+                    "mean_audio_rtf": [0.2, 0.35, 0.6, 0.85, 0.9]
-                        400,
-                        600,
-                        900
-                    ],
-                    "mean_e2el_ms": [
-                        26000,
-                        43000,
-                        52000,
-                        72000,
-                        109000
-                    ],
-                    "mean_audio_rtf": [
-                        0.15,
-                        0.30,
-                        0.35,
-                        0.50,
-                        0.75
-                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [10],
-                    10
+                "request_rate": [0.1],
-                ],
-                "request_rate": [
-                    0.1
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -78,27 +42,17 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [2000],
-                        600
+                    "mean_audio_ttfp_ms": [2000],
-                    ],
+                    "mean_audio_rtf": [0.25]
-                    "mean_e2el_ms": [
-                        6000
-                    ],
-                    "mean_audio_rtf": [
-                        0.20
-                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [40],
-                    40
+                "request_rate": [0.5],
-                ],
-                "request_rate": [
-                    0.5
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -115,27 +69,17 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [6000],
-                        300
+                    "mean_audio_ttfp_ms": [6000],
-                    ],
+                    "mean_audio_rtf": [0.7]
-                    "mean_e2el_ms": [
-                        7000
-                    ],
-                    "mean_audio_rtf": [
-                        0.25
-                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [100],
-                    100
+                "request_rate": [1.0],
-                ],
-                "request_rate": [
-                    1.0
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -154,79 +98,40 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [12000],
-                        700
+                    "mean_audio_ttfp_ms": [12000],
-                    ],
+                    "mean_audio_rtf": [1.0]
-                    "mean_e2el_ms": [
-                        11000
-                    ],
-                    "mean_audio_rtf": [
-                        0.35
-                    ]
                 }
             },
             {
                 "dataset_name": "random",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [4, 16, 32, 64, 128],
-                    4,
+                "max_concurrency": [1, 4, 8, 16, 32],
-                    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": [
+                    "modalities": ["text"]
-                        "text"
-                    ]
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000]
-                        100,
-                        300,
-                        500,
-                        800,
-                        1500
-                    ],
-                    "mean_e2el_ms": [
-                        5000,
-                        8000,
-                        10000,
-                        14000,
-                        20000
-                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [10],
-                    10
+                "request_rate": [0.1],
-                ],
-                "request_rate": [
-                    0.1
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": [
+                    "modalities": ["text"]
-                        "text"
-                    ]
                 },
                 "random_mm_base_items_per_request": 1,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -238,32 +143,21 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [2000]
-                        500
-                    ],
-                    "mean_e2el_ms": [
-                        1000
-                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [40],
-                    40
+                "request_rate": [0.5],
-                ],
-                "request_rate": [
-                    0.5
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": [
+                    "modalities": ["text"]
-                        "text"
-                    ]
                 },
                 "random_mm_base_items_per_request": 2,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -277,32 +171,21 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [6000]
-                        200
-                    ],
-                    "mean_e2el_ms": [
-                        800
-                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [100],
-                    100
+                "request_rate": [1.0],
-                ],
-                "request_rate": [
-                    1.0
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": [
+                    "modalities": ["text"]
-                        "text"
-                    ]
                 },
                 "random_mm_base_items_per_request": 3,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -318,12 +201,7 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [6000]
-                        600
-                    ],
-                    "mean_e2el_ms": [
-                        1300
-                    ]
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json

@@ -3,67 +3,31 @@
         "test_name": "test_qwen3_omni",
         "server_params": {
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
-            "extra_cli_args": [
+            "extra_cli_args": ["--no-async-chunk"]
-                "--no-async-chunk"
-            ]
         },
         "benchmark_params": [
             {
                 "dataset_name": "random",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [4, 16, 32, 64, 128],
-                    4,
+                "max_concurrency": [1, 4, 8, 16, 32],
-                    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": [
+                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000],
-                        400,
+                    "mean_audio_ttfp_ms": [30000, 60000, 90000, 120000, 150000],
-                        300,
+                    "mean_audio_rtf": [0.35, 0.45, 0.55, 0.65, 0.75]
-                        400,
-                        600,
-                        900
-                    ],
-                    "mean_e2el_ms": [
-                        32000,
-                        34000,
-                        45000,
-                        61000,
-                        90000
-                    ],
-                    "mean_audio_rtf": [
-                        0.20,
-                        0.25,
-                        0.30,
-                        0.45,
-                        0.60
-                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [10],
-                    10
+                "request_rate": [0.1],
-                ],
-                "request_rate": [
-                    0.1
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -78,27 +42,17 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [2000],
-                        500
+                    "mean_audio_ttfp_ms": [10000],
-                    ],
+                    "mean_audio_rtf": [0.25]
-                    "mean_e2el_ms": [
-                        6000
-                    ],
-                    "mean_audio_rtf": [
-                        0.20
-                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [40],
-                    40
+                "request_rate": [0.5],
-                ],
-                "request_rate": [
-                    0.5
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -115,27 +69,17 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [6000],
-                        300
+                    "mean_audio_ttfp_ms": [15000],
-                    ],
+                    "mean_audio_rtf": [0.45]
-                    "mean_e2el_ms": [
-                        7000
-                    ],
-                    "mean_audio_rtf": [
-                        0.20
-                    ]
                 }
             },
             {
                 "dataset_name": "random-mm",
                 "backend": "openai-chat-omni",
                 "endpoint": "/v1/chat/completions",
-                "num_prompts": [
+                "num_prompts": [100],
-                    100
+                "request_rate": [1.0],
-                ],
-                "request_rate": [
-                    1.0
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -154,15 +98,9 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [12000],
-                        700
+                    "mean_audio_ttfp_ms": [18000],
-                    ],
+                    "mean_audio_rtf": [0.9]
-                    "mean_e2el_ms": [
-                        9000
-                    ],
-                    "mean_audio_rtf": [
-                        0.30
-                    ]
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

@@ -4,9 +4,7 @@
         "server_params": {
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
             "use_omni": false,
-            "extra_cli_args": [
+            "extra_cli_args": ["--no-enable-prefix-caching"]
-                "--no-enable-prefix-caching"
-            ]
         },
         "benchmark_params": [
             {
@@ -14,20 +12,8 @@
                 "backend": "vllm",
                 "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
                 "endpoint": "/v1/completions",
-                "num_prompts": [
+                "num_prompts": [4, 16, 32, 64, 128],
-                    4,
+                "max_concurrency": [1, 4, 8, 16, 32],
-                    16,
-                    32,
-                    64,
-                    128
-                ],
-                "max_concurrency": [
-                    1,
-                    4,
-                    8,
-                    16,
-                    32
-                ],
                 "random_input_len": 2500,
                 "random_output_len": 900,
                 "temperature": 0.4,
@@ -37,26 +23,11 @@
                 "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": [
+                    "modalities": ["text"]
-                        "text"
-                    ]
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000]
-                        300,
-                        300,
-                        1000,
-                        600,
-                        800
-                    ],
-                    "mean_e2el_ms": [
-                        6000,
-                        8000,
-                        11000,
-                        13000,
-                        17000
-                    ]
                 }
             },
             {
@@ -64,12 +35,8 @@
                 "backend": "vllm",
                 "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
                 "endpoint": "/v1/completions",
-                "num_prompts": [
+                "num_prompts": [10],
-                    10
+                "request_rate": [0.1],
-                ],
-                "request_rate": [
-                    0.1
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -80,9 +47,7 @@
                 "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": [
+                    "modalities": ["text"]
-                        "text"
-                    ]
                 },
                 "random_mm_base_items_per_request": 1,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -94,12 +59,7 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [2000]
-                        30
-                    ],
-                    "mean_e2el_ms": [
-                        540
-                    ]
                 }
             },
             {
@@ -107,12 +67,8 @@
                 "backend": "vllm",
                 "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
                 "endpoint": "/v1/completions",
-                "num_prompts": [
+                "num_prompts": [40],
-                    40
+                "request_rate": [0.5],
-                ],
-                "request_rate": [
-                    0.5
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -123,9 +79,7 @@
                 "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": [
+                    "modalities": ["text"]
-                        "text"
-                    ]
                 },
                 "random_mm_base_items_per_request": 2,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -139,12 +93,7 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [6000]
-                        30
-                    ],
-                    "mean_e2el_ms": [
-                        600
-                    ]
                 }
             },
             {
@@ -152,12 +101,8 @@
                 "backend": "vllm",
                 "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
                 "endpoint": "/v1/completions",
-                "num_prompts": [
+                "num_prompts": [100],
-                    100
+                "request_rate": [1.0],
-                ],
-                "request_rate": [
-                    1.0
-                ],
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
@@ -168,9 +113,7 @@
                 "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
-                    "modalities": [
+                    "modalities": ["text"]
-                        "text"
-                    ]
                 },
                 "random_mm_base_items_per_request": 3,
                 "random_mm_num_mm_items_range_ratio": 0.5,
@@ -186,12 +129,7 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el",
                 "baseline": {
-                    "mean_ttft_ms": [
+                    "mean_ttft_ms": [6000]
-                        40
-                    ],
-                    "mean_e2el_ms": [
-                        700
-                    ]
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json

@@ -22,10 +22,10 @@
                 "num-input-images": 2,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.06,
+                    "throughput_qps": 0.05,
-                    "latency_mean": 20,
+                    "latency_mean": 18,
-                    "peak_memory_mb_max": 57000,
+                    "peak_memory_mb_max": 78500,
-                    "peak_memory_mb_mean": 57000
+                    "peak_memory_mb_mean": 78500
                 }
             },
             {
@@ -41,9 +41,9 @@
                 "enable-negative-prompt": true,
                 "baseline": {
                     "throughput_qps": 0.01,
-                    "latency_mean": 60,
+                    "latency_mean": 70,
-                    "peak_memory_mb_max": 67000,
+                    "peak_memory_mb_max": 81000,
-                    "peak_memory_mb_mean": 67000
+                    "peak_memory_mb_mean": 81000
                 }
             }
         ]
@@ -75,10 +75,10 @@
                 "num-input-images": 2,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.15,
+                    "throughput_qps": 0.1,
-                    "latency_mean": 10,
+                    "latency_mean": 12,
-                    "peak_memory_mb_max": 55000,
+                    "peak_memory_mb_max": 69000,
-                    "peak_memory_mb_mean": 55000
+                    "peak_memory_mb_mean": 69000
                 }
             },
             {
@@ -93,10 +93,10 @@
                 "num-input-images": 2,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.05,
+                    "throughput_qps": 0.03,
-                    "latency_mean": 20,
+                    "latency_mean": 28,
-                    "peak_memory_mb_max": 55000,
+                    "peak_memory_mb_max": 69000,
-                    "peak_memory_mb_mean": 55000
+                    "peak_memory_mb_mean": 69000
                 }
             }
         ]
@@ -138,10 +138,10 @@
                 "num-input-images": 2,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.21,
+                    "throughput_qps": 0.10,
-                    "latency_mean": 10,
+                    "latency_mean": 12,
-                    "peak_memory_mb_max": 57000,
+                    "peak_memory_mb_max": 73000,
-                    "peak_memory_mb_mean": 57000
+                    "peak_memory_mb_mean": 73000
                 }
             },
             {
@@ -156,10 +156,10 @@
                 "num-input-images": 2,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.08,
+                    "throughput_qps": 0.05,
                     "latency_mean": 20,
-                    "peak_memory_mb_max": 66000,
+                    "peak_memory_mb_max": 81000,
-                    "peak_memory_mb_mean": 66000
+                    "peak_memory_mb_mean": 81000
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json

@@ -21,10 +21,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.1,
+                    "throughput_qps": 0.05,
-                    "latency_mean": 10,
+                    "latency_mean": 15.0,
-                    "peak_memory_mb_max": 57000,
+                    "peak_memory_mb_max": 72500,
-                    "peak_memory_mb_mean": 57000
+                    "peak_memory_mb_mean": 72500
                 }
             },
             {
@@ -38,10 +38,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.02,
+                    "throughput_qps": 0.01,
-                    "latency_mean": 50,
+                    "latency_mean": 65.6,
-                    "peak_memory_mb_max": 65000,
+                    "peak_memory_mb_max": 80777,
-                    "peak_memory_mb_mean": 65000
+                    "peak_memory_mb_mean": 80777
                 }
             }
         ]
@@ -72,10 +72,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.21,
+                    "throughput_qps": 0.10,
-                    "latency_mean": 10,
+                    "latency_mean": 7.2,
-                    "peak_memory_mb_max": 56000,
+                    "peak_memory_mb_max": 68100,
-                    "peak_memory_mb_mean": 56000
+                    "peak_memory_mb_mean": 68100
                 }
             },
             {
@@ -89,10 +89,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.06,
+                    "throughput_qps": 0.03,
-                    "latency_mean": 20,
+                    "latency_mean": 24.0,
-                    "peak_memory_mb_max": 56000,
+                    "peak_memory_mb_max": 68100,
-                    "peak_memory_mb_mean": 56000
+                    "peak_memory_mb_mean": 68100
                 }
             }
         ]
@@ -133,10 +133,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.26,
+                    "throughput_qps": 0.1,
-                    "latency_mean": 10,
+                    "latency_mean": 6.5,
-                    "peak_memory_mb_max": 58000,
+                    "peak_memory_mb_max": 72600,
-                    "peak_memory_mb_mean": 58000
+                    "peak_memory_mb_mean": 72600
                 }
             },
             {
@@ -150,10 +150,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.1,
+                    "throughput_qps": 0.05,
-                    "latency_mean": 10,
+                    "latency_mean": 16.0,
-                    "peak_memory_mb_max": 67000,
+                    "peak_memory_mb_max": 81000,
-                    "peak_memory_mb_mean": 67000
+                    "peak_memory_mb_mean": 81000
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json

@@ -21,10 +21,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.05,
+                    "throughput_qps": 0.02,
-                    "latency_mean": 20,
+                    "latency_mean": 40.0,
-                    "peak_memory_mb_max": 63000,
+                    "peak_memory_mb_max": 70000,
-                    "peak_memory_mb_mean": 63000
+                    "peak_memory_mb_mean": 70000
                 }
             },
             {
@@ -38,10 +38,10 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.03,
+                    "throughput_qps": 0.005,
-                    "latency_mean": 30,
+                    "latency_mean": 80.0,
-                    "peak_memory_mb_max": 63000,
+                    "peak_memory_mb_max": 70000,
-                    "peak_memory_mb_mean": 63000
+                    "peak_memory_mb_mean": 70000
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen_image_vllm_omni.json

@@ -21,9 +21,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.35,
+                    "throughput_qps": 0.30,
-                    "latency_mean": 3,
+                    "latency_mean": 3.50,
-                    "peak_memory_mb_mean": 56000
+                    "peak_memory_mb_mean": 67000
                 }
             },
             {
@@ -37,9 +37,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.03,
+                    "throughput_qps": 0.037,
-                    "latency_mean": 26,
+                    "latency_mean": 27.0,
-                    "peak_memory_mb_mean": 65000
+                    "peak_memory_mb_mean": 74000
                 }
             }
         ]
@@ -67,9 +67,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.35,
+                    "throughput_qps": 0.30,
-                    "latency_mean": 3,
+                    "latency_mean": 3.50,
-                    "peak_memory_mb_mean": 56000
+                    "peak_memory_mb_mean": 67000
                 }
             },
             {
@@ -83,9 +83,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.03,
+                    "throughput_qps": 0.037,
-                    "latency_mean": 26,
+                    "latency_mean": 27.0,
-                    "peak_memory_mb_mean": 65000
+                    "peak_memory_mb_mean": 74000
                 }
             }
         ]
@@ -116,9 +116,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.11,
+                    "throughput_qps": 0.1,
-                    "latency_mean": 9,
+                    "latency_mean": 9.1,
-                    "peak_memory_mb_mean": 55000
+                    "peak_memory_mb_mean": 61000
                 }
             }
         ]
@@ -159,9 +159,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.46,
+                    "throughput_qps": 0.1,
-                    "latency_mean": 3,
+                    "latency_mean": 3.50,
-                    "peak_memory_mb_mean": 56000
+                    "peak_memory_mb_mean": 67000
                 }
             },
             {
@@ -175,9 +175,9 @@
                 "max-concurrency": 1,
                 "enable-negative-prompt": true,
                 "baseline": {
-                    "throughput_qps": 0.18,
+                    "throughput_qps": 0.1,
-                    "latency_mean": 6,
+                    "latency_mean": 6.15,
-                    "peak_memory_mb_mean": 67000
+                    "peak_memory_mb_mean": 74000
                 }
             }
         ]

tests/dfx/perf/tests/test_tts.json

@@ -4,9 +4,7 @@
     "server_params": {
       "model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base",
       "stage_overrides": {
-        "1": {
+        "1": {"devices": "1"}
-          "devices": "1"
-        }
       }
     },
     "benchmark_params": [
@@ -17,21 +15,13 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "linyueqian/seed-tts-eval-subset",
-        "num_prompts": [
+        "num_prompts": [20],
-          20
+        "max_concurrency": [1],
-        ],
-        "max_concurrency": [
-          1
-        ],
         "seed_tts_locale": "en",
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [
+          "median_audio_ttfp_ms": [350],
-            3000
+          "median_audio_rtf": [0.25]
-          ],
-          "median_audio_rtf": [
-            0.6
-          ]
         }
       },
       {
@@ -41,34 +31,14 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "linyueqian/seed-tts-eval-subset",
-        "num_prompts": [
+        "num_prompts": [80, 128, 128],
-          80,
+        "max_concurrency": [8, 16, 64],
-          128,
-          128
-        ],
-        "max_concurrency": [
-          8,
-          16,
-          64
-        ],
         "seed_tts_locale": "en",
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [
+          "median_audio_ttfp_ms": [3500, 30000, 7000],
-            1000,
+          "median_audio_rtf": [0.75, 10.0, 1.5],
-            3000,
+          "audio_throughput": [10.0, 1.0, 40.0]
-            14000
-          ],
-          "median_audio_rtf": [
-            0.70,
-            1.15,
-            3.90
-          ],
-          "audio_throughput": [
-            12,
-            14,
-            14
-          ]
         }
       },
       {
@@ -79,19 +49,13 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "linyueqian/seed-tts-eval-subset",
-        "num_prompts": [
+        "num_prompts": [200],
-          200
+        "max_concurrency": [4],
-        ],
-        "max_concurrency": [
-          4
-        ],
         "seed_tts_locale": "en",
         "seed_tts_wer_eval": true,
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "mean_audio_rtf": [
+          "mean_audio_rtf": [0.45]
-            0.45
-          ]
         }
       }
     ]
@@ -101,9 +65,7 @@
     "server_params": {
       "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
       "stage_overrides": {
-        "1": {
+        "1": {"devices": "1"}
-          "devices": "1"
-        }
       }
     },
     "benchmark_params": [
@@ -114,26 +76,14 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [
+        "num_prompts": [20],
-          20
+        "max_concurrency": [1],
-        ],
-        "max_concurrency": [
-          1
-        ],
         "seed_tts_locale": "en",
-        "extra_body": {
+        "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
-          "voice": "Vivian",
-          "language": "English",
-          "task_type": "CustomVoice"
-        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [
+          "median_audio_ttfp_ms": [150],
-            2000
+          "median_audio_rtf": [0.15]
-          ],
-          "median_audio_rtf": [
-            0.40
-          ]
         }
       },
       {
@@ -143,39 +93,15 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [
+        "num_prompts": [80, 128, 128],
-          80,
+        "max_concurrency": [8, 16, 64],
-          128,
-          128
-        ],
-        "max_concurrency": [
-          8,
-          16,
-          64
-        ],
         "seed_tts_locale": "en",
-        "extra_body": {
+        "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
-          "voice": "Vivian",
-          "language": "English",
-          "task_type": "CustomVoice"
-        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [
+          "median_audio_ttfp_ms": [1500, 30000, 3000],
-            300,
+          "median_audio_rtf": [0.30, 10.0, 0.60],
-            1000,
+          "audio_throughput": [30.0, 1.0, 100.0]
-            5000
-          ],
-          "median_audio_rtf": [
-            0.25,
-            0.35,
-            1.10
-          ],
-          "audio_throughput": [
-            30,
-            40,
-            40
-          ]
         }
       },
       {
@@ -186,24 +112,14 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [
+        "num_prompts": [200],
-          200
+        "max_concurrency": [4],
-        ],
-        "max_concurrency": [
-          4
-        ],
         "seed_tts_locale": "en",
-        "extra_body": {
+        "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
-          "voice": "Vivian",
-          "language": "English",
-          "task_type": "CustomVoice"
-        },
         "seed_tts_wer_eval": true,
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "mean_audio_rtf": [
+          "mean_audio_rtf": [0.35]
-            0.35
-          ]
         }
       },
       {
@@ -213,29 +129,15 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
-        "num_prompts": [
+        "num_prompts": [100],
-          100
+        "request_rate": [2.0],
-        ],
-        "request_rate": [
-          2.0
-        ],
         "seed_tts_locale": "en",
-        "extra_body": {
+        "extra_body": {"voice": "Vivian", "language": "English", "task_type": "CustomVoice"},
-          "voice": "Vivian",
-          "language": "English",
-          "task_type": "CustomVoice"
-        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [
+          "median_audio_ttfp_ms": [3000],
-            100
+          "median_audio_rtf": [0.25],
-          ],
+          "audio_throughput": [4.0]
-          "median_audio_rtf": [
-            0.20
-          ],
-          "audio_throughput": [
-            10
-          ]
         }
       },
       {
@@ -245,25 +147,14 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_design",
-        "num_prompts": [
+        "num_prompts": [20],
-          20
+        "max_concurrency": [1],
-        ],
-        "max_concurrency": [
-          1
-        ],
         "seed_tts_locale": "en",
-        "extra_body": {
+        "extra_body": {"task_type": "VoiceDesign", "language": "English"},
-          "task_type": "VoiceDesign",
-          "language": "English"
-        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [
+          "median_audio_ttfp_ms": [150],
-            100
+          "median_audio_rtf": [0.15]
-          ],
-          "median_audio_rtf": [
-            0.15
-          ]
         }
       },
       {
@@ -273,38 +164,15 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_design",
-        "num_prompts": [
+        "num_prompts": [80, 128, 128],
-          80,
+        "max_concurrency": [8, 16, 64],
-          128,
-          128
-        ],
-        "max_concurrency": [
-          8,
-          16,
-          64
-        ],
         "seed_tts_locale": "en",
-        "extra_body": {
+        "extra_body": {"task_type": "VoiceDesign", "language": "English"},
-          "task_type": "VoiceDesign",
-          "language": "English"
-        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [
+          "median_audio_ttfp_ms": [1500, 30000, 3500],
-            100,
+          "median_audio_rtf": [0.35, 10.0, 0.70],
-            1000,
+          "audio_throughput": [25.0, 1.0, 80.0]
-            5000
-          ],
-          "median_audio_rtf": [
-            0.25,
-            0.40,
-            1.20
-          ],
-          "audio_throughput": [
-            30,
-            40,
-            40
-          ]
         }
       },
       {
@@ -314,28 +182,15 @@
         "backend": "openai-audio-speech",
         "endpoint": "/v1/audio/speech",
         "dataset_path": "benchmarks/build_dataset/seed_tts_design",
-        "num_prompts": [
+        "num_prompts": [100],
-          100
+        "request_rate": [2.0],
-        ],
-        "request_rate": [
-          2.0
-        ],
         "seed_tts_locale": "en",
-        "extra_body": {
+        "extra_body": {"task_type": "VoiceDesign", "language": "English"},
-          "task_type": "VoiceDesign",
-          "language": "English"
-        },
         "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
         "baseline": {
-          "median_audio_ttfp_ms": [
+          "median_audio_ttfp_ms": [3500],
-            100
+          "median_audio_rtf": [0.30],
-          ],
+          "audio_throughput": [4.0]
-          "median_audio_rtf": [
-            0.20
-          ],
-          "audio_throughput": [
-            9
-          ]
         }
       }
     ]

tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json

@@ -30,9 +30,9 @@
                     }
                 ],
                 "baseline": {
-                    "throughput_qps": 0.04,
+                    "throughput_qps": 0.034,
-                    "latency_mean": 30,
+                    "latency_mean": 26.0,
-                    "peak_memory_mb_mean": 75000
+                    "peak_memory_mb_mean": 80000
                 }
             }
         ]
@@ -72,9 +72,9 @@
                     }
                 ],
                 "baseline": {
-                    "throughput_qps": 0.06,
+                    "throughput_qps": 0.042,
-                    "latency_mean": 20,
+                    "latency_mean": 21.6,
-                    "peak_memory_mb_mean": 48000
+                    "peak_memory_mb_mean": 55300
                 }
             },
             {
@@ -97,9 +97,9 @@
                     }
                 ],
                 "baseline": {
-                    "throughput_qps": 0.01,
+                    "throughput_qps": 0.0085,
-                    "latency_mean": 90,
+                    "latency_mean": 101.6,
-                    "peak_memory_mb_mean": 56000
+                    "peak_memory_mb_mean": 65200
                 }
             }
         ]

tests/diffusion/distributed/test_autoencoder_kl_hunyuan.py

@@ -1,79 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import (
-    DistributedAutoencoderKLHunyuan,
-)
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class _DummyDistributedAutoencoderKLHunyuan(DistributedAutoencoderKLHunyuan):
-    def __init__(self):
-        torch.nn.Module.__init__(self)
-        self.tile_latent_min_size = 8
-        self.tile_sample_min_size = 8
-        self.tile_overlap_factor = 0.25
-        self.use_spatial_tiling = False
-
-    @property
-    def dtype(self) -> torch.dtype:
-        return torch.float32
-
-    def decoder(self, x: torch.Tensor) -> torch.Tensor:
-        return x + 10
-
-    def encoder(self, x: torch.Tensor) -> torch.Tensor:
-        return x + 20
-
-    def blend_v(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
-        return b
-
-    def blend_h(self, _a: torch.Tensor, b: torch.Tensor, _blend_extent: int) -> torch.Tensor:
-        return b
-
-
-def test_hunyuan_vae_use_tiling_aliases_spatial_tiling():
-    # Verify use_tiling property maps to use_spatial_tiling.
-    vae = _DummyDistributedAutoencoderKLHunyuan()
-
-    assert not vae.use_tiling
-
-    vae.use_tiling = True
-
-    assert vae.use_spatial_tiling
-
-
-def test_hunyuan_vae_decode_tiles_round_trip():
-    # Validate decode tile split/exec/merge returns expected reconstructed tensor.
-    vae = _DummyDistributedAutoencoderKLHunyuan()
-    z = torch.arange(144, dtype=torch.float32).reshape(1, 1, 1, 12, 12)
-
-    tile_tasks, grid_spec = vae.tile_split(z)
-    decoded_tiles = {task.grid_coord: vae.tile_exec(task) for task in tile_tasks}
-    output = vae.tile_merge(decoded_tiles, grid_spec)
-
-    assert grid_spec.split_dims == (3, 4)
-    assert grid_spec.grid_shape == (2, 2)
-    assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6}
-    assert len(tile_tasks) == 4
-    assert torch.equal(output, z + 10)
-
-
-def test_hunyuan_vae_encode_tiles_round_trip():
-    # Validate encode tile split/exec/merge returns expected latent tensor.
-    vae = _DummyDistributedAutoencoderKLHunyuan()
-    x = torch.arange(144, dtype=torch.float32).reshape(1, 1, 1, 12, 12)
-
-    tile_tasks, grid_spec = vae.encode_tile_split(x)
-    encoded_tiles = {task.grid_coord: vae.encode_tile_exec(task) for task in tile_tasks}
-    output = vae.encode_tile_merge(encoded_tiles, grid_spec)
-
-    assert grid_spec.split_dims == (3, 4)
-    assert grid_spec.grid_shape == (2, 2)
-    assert grid_spec.tile_spec == {"blend_extent": 2, "row_limit": 6}
-    assert len(tile_tasks) == 4
-    assert torch.equal(output, x + 20)

tests/e2e/accuracy/test_helios_distilled.py

@@ -0,0 +1,368 @@
+from __future__ import annotations
+
+import json
+import math
+import os
+import shutil
+import signal
+import socket
+import subprocess
+import sys
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import reset_artifact_dir
+from tests.helpers.mark import hardware_test
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+MODEL_NAME = "BestWishYsh/Helios-Distilled"
+PROMPT = "A cat wearing sunglasses dances on a beach at sunset, cinematic lighting."
+HEIGHT = 384
+WIDTH = 640
+NUM_FRAMES = 99
+FPS = 16
+SEED = 42
+GUIDANCE_SCALE = 1.0
+PYRAMID_NUM_INFERENCE_STEPS = (1, 1, 1)
+MIN_PSNR_DB = 35.0
+MAX_MAE = 3.0
+MAX_TIME_RATIO = 1.03
+WARMUP_REQUESTS = 1
+WARMUP_NUM_INFERENCE_STEPS = 1
+
+
+def _run(command: list[str], *, cwd: Path, env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]:
+    print("+ " + " ".join(command), flush=True)
+    return subprocess.run(
+        command,
+        cwd=str(cwd),
+        env=env,
+        check=True,
+        text=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+    )
+
+
+def _read_json(path: Path) -> dict[str, Any]:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _open_port(host: str = "127.0.0.1") -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind((host, 0))
+        return int(sock.getsockname()[1])
+
+
+def _wait_for_service(base_url: str, process: subprocess.Popen, *, timeout_s: float = 1200.0) -> None:
+    import requests
+
+    deadline = time.monotonic() + timeout_s
+    while time.monotonic() < deadline:
+        if process.poll() is not None:
+            raise RuntimeError(f"Server exited with code {process.returncode} before becoming ready.")
+        try:
+            response = requests.get(f"{base_url}/health", timeout=1)
+            if response.status_code == 200:
+                return
+        except requests.RequestException:
+            pass
+        time.sleep(2)
+    raise TimeoutError(f"Server at {base_url} did not become ready within {timeout_s:.0f}s.")
+
+
+def _stop_process(process: subprocess.Popen) -> None:
+    if process.poll() is not None:
+        return
+    try:
+        os.killpg(process.pid, signal.SIGTERM)
+        process.wait(timeout=30)
+    except Exception:
+        os.killpg(process.pid, signal.SIGKILL)
+        process.wait(timeout=30)
+
+
+def _extract_frames(video_path: Path, output_dir: Path) -> list[Image.Image]:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    frame_pattern = output_dir / "frame_%05d.png"
+    _run(
+        [
+            "ffmpeg",
+            "-hide_banner",
+            "-loglevel",
+            "error",
+            "-i",
+            str(video_path),
+            "-vsync",
+            "0",
+            str(frame_pattern),
+        ],
+        cwd=output_dir,
+    )
+    frame_paths = sorted(output_dir.glob("frame_*.png"))
+    if not frame_paths:
+        raise RuntimeError(f"No frames extracted from {video_path}")
+    return [Image.open(path).convert("RGB").copy() for path in frame_paths]
+
+
+def _image_metrics(reference: np.ndarray, candidate: np.ndarray) -> dict[str, float]:
+    diff = reference.astype(np.float64) - candidate.astype(np.float64)
+    mse = float(np.mean(diff * diff))
+    mae = float(np.mean(np.abs(diff)))
+    max_abs = float(np.max(np.abs(diff)))
+    psnr_db = float("inf") if mse == 0.0 else 20.0 * math.log10(255.0) - 10.0 * math.log10(mse)
+    ref_flat = reference.astype(np.float64).reshape(-1)
+    cand_flat = candidate.astype(np.float64).reshape(-1)
+    ref_norm = float(np.linalg.norm(ref_flat))
+    cand_norm = float(np.linalg.norm(cand_flat))
+    if ref_norm == 0.0 and cand_norm == 0.0:
+        cosine = 1.0
+    else:
+        cosine = float(np.dot(ref_flat, cand_flat) / (ref_norm * cand_norm))
+    return {
+        "mae": mae,
+        "mse": mse,
+        "max_abs": max_abs,
+        "psnr_db": psnr_db,
+        "cosine_similarity": cosine,
+    }
+
+
+def _video_metrics(reference_frames: list[Image.Image], candidate_frames: list[Image.Image]) -> dict[str, Any]:
+    if len(reference_frames) != len(candidate_frames):
+        raise ValueError(f"Frame count mismatch: reference={len(reference_frames)} candidate={len(candidate_frames)}")
+
+    ref_stack = np.stack([np.asarray(frame, dtype=np.uint8) for frame in reference_frames], axis=0)
+    cand_stack = np.stack([np.asarray(frame, dtype=np.uint8) for frame in candidate_frames], axis=0)
+    mid = len(reference_frames) // 2
+    return {
+        "num_frames": len(reference_frames),
+        "frame0_metrics": _image_metrics(ref_stack[0], cand_stack[0]),
+        "mid_frame_index": mid,
+        "mid_frame_metrics": _image_metrics(ref_stack[mid], cand_stack[mid]),
+        "all_frames_metrics": _image_metrics(ref_stack, cand_stack),
+    }
+
+
+def _reference_repo(output_root: Path) -> Path:
+    configured = os.environ.get("VLLM_HELIOS_REFERENCE_REPO")
+    if configured:
+        reference = Path(configured).expanduser().resolve()
+        if not reference.exists():
+            raise FileNotFoundError(f"VLLM_HELIOS_REFERENCE_REPO does not exist: {reference}")
+        return reference
+
+    reference = output_root / "origin_main"
+    if reference.exists():
+        subprocess.run(["git", "worktree", "remove", "--force", str(reference)], cwd=str(REPO_ROOT), check=False)
+        if reference.exists():
+            shutil.rmtree(reference)
+    _run(["git", "fetch", "origin", "main"], cwd=REPO_ROOT)
+    _run(["git", "worktree", "add", "--detach", str(reference), "origin/main"], cwd=REPO_ROOT)
+    return reference
+
+
+def _server_env(repo: Path) -> dict[str, str]:
+    env = dict(os.environ)
+    env["PYTHONPATH"] = str(repo) + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
+    env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+    return env
+
+
+def _benchmark_config() -> str:
+    return json.dumps(
+        [
+            {
+                "width": WIDTH,
+                "height": HEIGHT,
+                "num_inference_steps": 50,
+                "num_frames": NUM_FRAMES,
+                "fps": FPS,
+                "guidance_scale": GUIDANCE_SCALE,
+                "is_enable_stage2": True,
+                "pyramid_num_stages": 3,
+                "pyramid_num_inference_steps_list": list(PYRAMID_NUM_INFERENCE_STEPS),
+                "is_amplify_first_chunk": False,
+                "weight": 1,
+            }
+        ]
+    )
+
+
+def _start_server(*, repo: Path, variant_dir: Path) -> tuple[subprocess.Popen, str]:
+    port = _open_port()
+    base_url = f"http://127.0.0.1:{port}"
+    log_file = (variant_dir / "server.log").open("w", encoding="utf-8")
+    command = [
+        sys.executable,
+        "-m",
+        "vllm_omni.entrypoints.cli.main",
+        "serve",
+        MODEL_NAME,
+        "--host",
+        "127.0.0.1",
+        "--port",
+        str(port),
+        "--omni",
+    ]
+    print("+ " + " ".join(command), flush=True)
+    process = subprocess.Popen(
+        command,
+        cwd=str(repo),
+        env=_server_env(repo),
+        stdout=log_file,
+        stderr=subprocess.STDOUT,
+        preexec_fn=os.setsid,
+        text=True,
+    )
+    try:
+        _wait_for_service(base_url, process)
+    except Exception:
+        log_file.close()
+        _stop_process(process)
+        raise
+    log_file.close()
+    return process, base_url
+
+
+def _run_serving_benchmark(*, repo: Path, variant_dir: Path, base_url: str) -> dict[str, Any]:
+    metrics_path = variant_dir / "metrics.json"
+    response_dir = variant_dir / "responses"
+    command = [
+        sys.executable,
+        str(REPO_ROOT / "benchmarks/diffusion/diffusion_benchmark_serving.py"),
+        "--base-url",
+        base_url,
+        "--endpoint",
+        "/v1/videos",
+        "--model",
+        MODEL_NAME,
+        "--dataset",
+        "random",
+        "--task",
+        "t2v",
+        "--num-prompts",
+        "1",
+        "--max-concurrency",
+        "1",
+        "--warmup-requests",
+        str(WARMUP_REQUESTS),
+        "--warmup-num-inference-steps",
+        str(WARMUP_NUM_INFERENCE_STEPS),
+        "--warmup-concurrency",
+        "1",
+        "--random-prompt",
+        PROMPT,
+        "--seed",
+        str(SEED),
+        "--disable-tqdm",
+        "--random-request-config",
+        _benchmark_config(),
+        "--output-file",
+        str(metrics_path),
+        "--save-response-dir",
+        str(response_dir),
+    ]
+    completed = _run(command, cwd=repo, env=_server_env(repo))
+    (variant_dir / "benchmark.log").write_text(completed.stdout, encoding="utf-8")
+    metrics = _read_json(metrics_path)
+    saved_responses = metrics.get("saved_responses") or []
+    if len(saved_responses) != 1:
+        raise RuntimeError(f"Expected one saved video response, got {saved_responses}")
+    return {
+        "repo": str(repo),
+        "metrics": metrics,
+        "video": saved_responses[0],
+        "log": str(variant_dir / "benchmark.log"),
+        "server_log": str(variant_dir / "server.log"),
+    }
+
+
+def _run_variant(*, label: str, repo: Path, output_root: Path) -> dict[str, Any]:
+    variant_dir = output_root / label
+    variant_dir.mkdir(parents=True, exist_ok=True)
+    process, base_url = _start_server(repo=repo, variant_dir=variant_dir)
+    try:
+        return _run_serving_benchmark(repo=repo, variant_dir=variant_dir, base_url=base_url)
+    finally:
+        _stop_process(process)
+
+
+@pytest.mark.benchmark
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_helios_distilled_pr_matches_main_acc_perf(accuracy_artifact_root: Path) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("Helios-Distilled accuracy/perf comparison requires CUDA.")
+
+    output_root = reset_artifact_dir(accuracy_artifact_root / "helios_distilled_acc_perf")
+    reference = _reference_repo(output_root)
+    output_json = output_root / "result.json"
+
+    reference_result = _run_variant(label="reference", repo=reference, output_root=output_root)
+    candidate_result = _run_variant(label="candidate", repo=REPO_ROOT, output_root=output_root)
+
+    with TemporaryDirectory() as tmp:
+        tmp_dir = Path(tmp)
+        reference_frames = _extract_frames(Path(reference_result["video"]), tmp_dir / "reference_frames")
+        candidate_frames = _extract_frames(Path(candidate_result["video"]), tmp_dir / "candidate_frames")
+        output_metrics = _video_metrics(reference_frames, candidate_frames)
+
+    reference_time = float(reference_result["metrics"]["latency_mean"])
+    candidate_time = float(candidate_result["metrics"]["latency_mean"])
+    time_ratio = candidate_time / reference_time
+    result = {
+        "reference": reference_result,
+        "candidate": candidate_result,
+        "sampling": {
+            "model": MODEL_NAME,
+            "prompt": PROMPT,
+            "height": HEIGHT,
+            "width": WIDTH,
+            "num_frames": NUM_FRAMES,
+            "fps": FPS,
+            "seed": SEED,
+            "guidance_scale": GUIDANCE_SCALE,
+            "pyramid_num_inference_steps_list": list(PYRAMID_NUM_INFERENCE_STEPS),
+            "is_amplify_first_chunk": False,
+        },
+        "performance": {
+            "reference_measured_time_s": reference_time,
+            "candidate_measured_time_s": candidate_time,
+            "candidate_over_reference_time_ratio": time_ratio,
+        },
+        "output_metrics": output_metrics,
+        "thresholds": {
+            "min_psnr_db": MIN_PSNR_DB,
+            "max_mae": MAX_MAE,
+            "max_time_ratio": MAX_TIME_RATIO,
+        },
+    }
+    output_json.write_text(json.dumps(result, indent=2, sort_keys=True, allow_nan=True), encoding="utf-8")
+    all_metrics = output_metrics["all_frames_metrics"]
+    print(
+        json.dumps(
+            {
+                "output_json": str(output_json),
+                "candidate_over_reference_time_ratio": time_ratio,
+                "all_frames_psnr_db": all_metrics["psnr_db"],
+                "all_frames_mae": all_metrics["mae"],
+                "all_frames_cosine_similarity": all_metrics["cosine_similarity"],
+            },
+            indent=2,
+            allow_nan=True,
+        )
+    )
+
+    assert all_metrics["psnr_db"] >= MIN_PSNR_DB
+    assert all_metrics["mae"] <= MAX_MAE
+    assert time_ratio <= MAX_TIME_RATIO
+    assert output_json.exists(), f"Expected Helios comparison result at {output_json}"

vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_hunyuan.py

@@ -1,145 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-from typing import Any
-
-import torch
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
-    DistributedOperator,
-    DistributedVaeMixin,
-    GridSpec,
-    TileTask,
-)
-from vllm_omni.diffusion.models.hunyuan_image3.autoencoder import AutoencoderKLConv3D
-
-logger = init_logger(__name__)
-
-
-class DistributedAutoencoderKLHunyuan(AutoencoderKLConv3D, DistributedVaeMixin):
-    @classmethod
-    def from_config(cls, config: Any, **kwargs: Any):
-        model = super().from_config(config, **kwargs)
-        model.init_distributed()
-        return model
-
-    @classmethod
-    def from_pretrained(cls, *args: Any, **kwargs: Any):
-        model = super().from_pretrained(*args, **kwargs)
-        model.init_distributed()
-        return model
-
-    @property
-    def use_tiling(self) -> bool:
-        return self.use_spatial_tiling
-
-    @use_tiling.setter
-    def use_tiling(self, use_tiling: bool) -> None:
-        self.use_spatial_tiling = use_tiling
-
-    def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
-        _, _, _, height, width = z.shape
-        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
-        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
-        row_limit = int(self.tile_sample_min_size - blend_extent)
-
-        tiletask_list = []
-        for i in range(0, height, overlap_size):
-            for j in range(0, width, overlap_size):
-                tile = z[:, :, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
-                tiletask_list.append(
-                    TileTask(
-                        len(tiletask_list),
-                        (i // overlap_size, j // overlap_size),
-                        tile,
-                        workload=tile.shape[3] * tile.shape[4],
-                    )
-                )
-
-        grid_spec = GridSpec(
-            split_dims=(3, 4),
-            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
-            tile_spec={"blend_extent": blend_extent, "row_limit": row_limit},
-            output_dtype=self.dtype,
-        )
-        return tiletask_list, grid_spec
-
-    def tile_exec(self, task: TileTask) -> torch.Tensor:
-        return self.decoder(task.tensor)
-
-    def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
-        grid_h, grid_w = grid_spec.grid_shape
-        result_rows = []
-        for i in range(grid_h):
-            result_row = []
-            for j in range(grid_w):
-                tile = coord_tensor_map[(i, j)]
-                if i > 0:
-                    tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_extent"])
-                if j > 0:
-                    tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_extent"])
-                result_row.append(tile[:, :, :, : grid_spec.tile_spec["row_limit"], : grid_spec.tile_spec["row_limit"]])
-            result_rows.append(torch.cat(result_row, dim=-1))
-        return torch.cat(result_rows, dim=-2)
-
-    def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
-        _, _, _, height, width = x.shape
-        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
-        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
-        row_limit = int(self.tile_latent_min_size - blend_extent)
-
-        tiletask_list = []
-        for i in range(0, height, overlap_size):
-            for j in range(0, width, overlap_size):
-                tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
-                tiletask_list.append(
-                    TileTask(
-                        len(tiletask_list),
-                        (i // overlap_size, j // overlap_size),
-                        tile,
-                        workload=tile.shape[3] * tile.shape[4],
-                    )
-                )
-
-        grid_spec = GridSpec(
-            split_dims=(3, 4),
-            grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
-            tile_spec={"blend_extent": blend_extent, "row_limit": row_limit},
-            output_dtype=self.dtype,
-        )
-        return tiletask_list, grid_spec
-
-    def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
-        return self.encoder(task.tensor)
-
-    def encode_tile_merge(
-        self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
-    ) -> torch.Tensor:
-        return self.tile_merge(coord_tensor_map, grid_spec)
-
-    def spatial_tiled_encode(self, x: torch.Tensor):
-        if not self.is_distributed_enabled():
-            return super().spatial_tiled_encode(x)
-
-        logger.debug("Encode running with distributed executor")
-        return self.distributed_executor.execute(
-            x,
-            DistributedOperator(
-                split=self.encode_tile_split,
-                exec=self.encode_tile_exec,
-                merge=self.encode_tile_merge,
-            ),
-            broadcast_result=True,
-        )
-
-    def spatial_tiled_decode(self, z: torch.Tensor):
-        if not self.is_distributed_enabled():
-            return super().spatial_tiled_decode(z)
-
-        logger.debug("Decode running with distributed executor")
-        return self.distributed_executor.execute(
-            z,
-            DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
-            broadcast_result=True,
-        )

vllm_omni/diffusion/models/helios/pipeline_helios.py

@@ -29,7 +29,6 @@ from vllm_omni.diffusion.models.helios.scheduling_helios import HeliosScheduler
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import (
@@ -504,7 +503,12 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
 
         # Prepare frame indices
         if keep_first_frame:
-            indices = torch.arange(0, sum([1, *history_sizes, num_latent_frames_per_chunk]))
+            indices = torch.arange(
+                0,
+                sum([1, *history_sizes, num_latent_frames_per_chunk]),
+                device=device,
+                dtype=torch.float32,
+            )
             (
                 indices_prefix,
                 indices_latents_history_long,
@@ -514,7 +518,12 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
             ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0)
             indices_latents_history_short = torch.cat([indices_prefix, indices_latents_history_1x], dim=0)
         else:
-            indices = torch.arange(0, sum([*history_sizes, num_latent_frames_per_chunk]))
+            indices = torch.arange(
+                0,
+                sum([*history_sizes, num_latent_frames_per_chunk]),
+                device=device,
+                dtype=torch.float32,
+            )
             (
                 indices_latents_history_long,
                 indices_latents_history_mid,
@@ -531,6 +540,7 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
             attention_kwargs = {}
 
         vae_dtype = self.vae.dtype
+        latent_chunks_to_decode = []
 
         # Chunked denoising loop
         for k in range(num_latent_chunk):
@@ -645,16 +655,15 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
                 slice(-num_latent_frames_per_chunk, None),
             )
 
-            current_latents = real_history_latents[index_slice].to(vae_dtype) / latents_std + latents_mean
+            if output_type != "latent":
-            current_video = self.vae.decode(current_latents, return_dict=False)[0]
+                current_latents = real_history_latents[index_slice].to(vae_dtype) / latents_std + latents_mean
+                latent_chunks_to_decode.append(current_latents)
 
-            if history_video is None:
+        if output_type != "latent" and latent_chunks_to_decode:
-                history_video = current_video
+            batched_latents = torch.cat(latent_chunks_to_decode, dim=0)
-            else:
+            batched_video = self.vae.decode(batched_latents, return_dict=False)[0]
-                history_video = torch.cat([history_video, current_video], dim=2)
+            history_video = torch.cat(list(batched_video.split(batch_size, dim=0)), dim=2)
 
-        if current_omni_platform.is_available():
-            current_omni_platform.empty_cache()
         self._current_timestep = None
 
         if output_type == "latent":
@@ -688,6 +697,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
     ) -> torch.Tensor:
         """Single-stage denoising loop for one chunk."""
         batch_size = latents.shape[0]
+        latents_history_short = latents_history_short.to(transformer_dtype)
+        latents_history_mid = latents_history_mid.to(transformer_dtype)
+        latents_history_long = latents_history_long.to(transformer_dtype)
         do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None
 
         with self.progress_bar(total=len(timesteps)) as pbar:
@@ -702,9 +714,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
                     "indices_latents_history_short": indices_latents_history_short,
                     "indices_latents_history_mid": indices_latents_history_mid,
                     "indices_latents_history_long": indices_latents_history_long,
-                    "latents_history_short": latents_history_short.to(transformer_dtype),
+                    "latents_history_short": latents_history_short,
-                    "latents_history_mid": latents_history_mid.to(transformer_dtype),
+                    "latents_history_mid": latents_history_mid,
-                    "latents_history_long": latents_history_long.to(transformer_dtype),
+                    "latents_history_long": latents_history_long,
                     "attention_kwargs": attention_kwargs,
                     "return_dict": False,
                 }
@@ -783,6 +795,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
     ) -> torch.Tensor:
         """Pyramid multi-stage denoising for one chunk."""
         batch_size, num_channel, num_frames_lat, height, width = latents.shape
+        latents_history_short = latents_history_short.to(transformer_dtype)
+        latents_history_mid = latents_history_mid.to(transformer_dtype)
+        latents_history_long = latents_history_long.to(transformer_dtype)
 
         # Downsample latents to the smallest pyramid level
         latents_flat = latents.permute(0, 2, 1, 3, 4).reshape(batch_size * num_frames_lat, num_channel, height, width)
@@ -853,9 +868,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip
                         "indices_latents_history_short": indices_latents_history_short,
                         "indices_latents_history_mid": indices_latents_history_mid,
                         "indices_latents_history_long": indices_latents_history_long,
-                        "latents_history_short": latents_history_short.to(transformer_dtype),
+                        "latents_history_short": latents_history_short,
-                        "latents_history_mid": latents_history_mid.to(transformer_dtype),
+                        "latents_history_mid": latents_history_mid,
-                        "latents_history_long": latents_history_long.to(transformer_dtype),
+                        "latents_history_long": latents_history_long,
                         "attention_kwargs": attention_kwargs,
                         "return_dict": False,
                     }

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -30,6 +30,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer
 
+from .autoencoder import AutoencoderKLConv3D
 from .hunyuan_image3_tokenizer import TokenizerWrapper
 from .hunyuan_image3_transformer import (
     CausalMMOutputWithPast,
@@ -342,13 +343,7 @@ class HunyuanImage3Pipeline(
         quant_config = od_config.quantization_config
         self.model = HunyuanImage3Model(self.hf_config, quant_config=quant_config)
         self.transformer = self.model
-        # Lazy import to break circular dependency:
+        self.vae = AutoencoderKLConv3D.from_config(self.hf_config.vae)
-        # autoencoder_kl_hunyuan -> hunyuan_image3/__init__ -> pipeline_hunyuan_image3 -> autoencoder_kl_hunyuan
-        from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import (  # noqa: PLC0415
-            DistributedAutoencoderKLHunyuan,
-        )
-
-        self.vae = DistributedAutoencoderKLHunyuan.from_config(self.hf_config.vae)
         self.vae.use_spatial_tiling = self.od_config.vae_use_tiling
         self._pipeline = None
         self._tkwrapper = TokenizerWrapper(od_config.model)