[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.

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/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
-            current_video = self.vae.decode(current_latents, return_dict=False)[0]
+            if output_type != "latent":
+                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:
-                history_video = current_video
-            else:
-                history_video = torch.cat([history_video, current_video], dim=2)
+        if output_type != "latent" and latent_chunks_to_decode:
+            batched_latents = torch.cat(latent_chunks_to_decode, dim=0)
+            batched_video = self.vae.decode(batched_latents, return_dict=False)[0]
+            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_mid": latents_history_mid.to(transformer_dtype),
-                    "latents_history_long": latents_history_long.to(transformer_dtype),
+                    "latents_history_short": latents_history_short,
+                    "latents_history_mid": latents_history_mid,
+                    "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_mid": latents_history_mid.to(transformer_dtype),
-                        "latents_history_long": latents_history_long.to(transformer_dtype),
+                        "latents_history_short": latents_history_short,
+                        "latents_history_mid": latents_history_mid,
+                        "latents_history_long": latents_history_long,
                         "attention_kwargs": attention_kwargs,
                         "return_dict": False,
                     }