From 80c89e00cb47c5ad71ef66a1ccd1be482036a7f4 Mon Sep 17 00:00:00 2001 From: zhanqiuhu <49648934+ZhanqiuHu@users.noreply.github.com> Date: Mon, 2 Mar 2026 22:02:13 -0500 Subject: [PATCH 1/7] [Bugfix] Fix Qwen3-TTS code predictor crash due to missing vLLM config context (#1619) Signed-off-by: Zhanqiu --- .../qwen3_tts_code_predictor_vllm.py | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index a3133c179..ac05d23c8 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -370,16 +370,17 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): def _maybe_init_kv_cache(self, device: torch.device) -> None: if self._kv_cache is not None: return - max_seq_len = int(getattr(self.config, "num_code_groups", 16) or 16) - # Upper bound on batch size: vLLM scheduler max_num_seqs (fallback 8). - max_batch = int(getattr(self._vllm_config.scheduler_config, "max_num_seqs", 8) or 8) - max_batch = max(1, max_batch) - self._kv_cache = _LocalPredictorKVCache( - vllm_config=self._vllm_config, - max_seq_len=max_seq_len, - max_batch_size=max_batch, - device=device, - ) + with set_current_vllm_config(self._vllm_config): + max_seq_len = int(getattr(self.config, "num_code_groups", 16) or 16) + # Upper bound on batch size: vLLM scheduler max_num_seqs (fallback 8). + max_batch = int(getattr(self._vllm_config.scheduler_config, "max_num_seqs", 8) or 8) + max_batch = max(1, max_batch) + self._kv_cache = _LocalPredictorKVCache( + vllm_config=self._vllm_config, + max_seq_len=max_seq_len, + max_batch_size=max_batch, + device=device, + ) @torch.inference_mode() def reset_cache(self) -> None: From 90a48f44449c5097056da8c8fa91e3d3d6cc7f85 Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Tue, 3 Mar 2026 11:56:42 +0800 Subject: [PATCH 2/7] [Debug] Enable curl retry aligned with openai (#1539) Signed-off-by: tzhouam --- .../qwen2_5_omni/run_curl_multimodal_generation.sh | 12 +++++++----- .../qwen3_omni/run_curl_multimodal_generation.sh | 12 +++++++----- 2 files changed, 14 insertions(+), 10 deletions(-) diff --git a/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh b/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh index c5a265b62..ed09bac14 100644 --- a/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh +++ b/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh @@ -160,10 +160,7 @@ esac echo "Running query type: $QUERY_TYPE" echo "" - -output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d @- < Date: Tue, 3 Mar 2026 12:52:51 +0800 Subject: [PATCH 3/7] [Doc] Fix links in the configuration doc (#1615) Signed-off-by: Yuanheng Zhao --- docs/configuration/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/configuration/README.md b/docs/configuration/README.md index 02440b95d..e2073a96a 100644 --- a/docs/configuration/README.md +++ b/docs/configuration/README.md @@ -2,7 +2,7 @@ This section lists the most common options for running vLLM-Omni. -For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html) +For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html) Currently, the main options are maintained by stage configs for each model. @@ -19,3 +19,4 @@ For introduction, please check [Introduction for stage config](./stage_configs.m - **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss - **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models - **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models +- **[CPU Offloading](../user_guide/diffusion/cpu_offload_diffusion.md)** - Enable CPU offloading (model-level and layerwise) for for DiT models From e7f2c88701799245f6ea202b6ab38fa8cd0aa01a Mon Sep 17 00:00:00 2001 From: Alicia <115451386+congw729@users.noreply.github.com> Date: Tue, 3 Mar 2026 15:00:33 +0800 Subject: [PATCH 4/7] [CI] Add scripts for bechmark collection and email distribution. (#1307) Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com> --- .buildkite/test-nightly.yml | 27 +- pyproject.toml | 3 +- tests/perf/scripts/run_benchmark.py | 17 +- tools/nightly/generate_nightly_perf_excel.py | 437 +++++++++++++++++++ tools/nightly/send_nightly_perf_email.py | 233 ++++++++++ 5 files changed, 704 insertions(+), 13 deletions(-) create mode 100644 tools/nightly/generate_nightly_perf_excel.py create mode 100644 tools/nightly/send_nightly_perf_email.py diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 7b535033e..8b21acfa7 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -1,5 +1,5 @@ steps: - - label: "Omni Model Test with H100" + - label: ":full_moon: Omni Model Test with H100" timeout_in_minutes: 90 depends_on: upload-nightly-pipeline if: build.env("NIGHTLY") == "1" @@ -41,7 +41,7 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: "Omni Model Test" + - label: ":full_moon: Omni Model Test" timeout_in_minutes: 60 depends_on: upload-nightly-pipeline if: build.env("NIGHTLY") == "1" @@ -56,18 +56,22 @@ steps: always-pull: true shm-size: "8gb" propagate-environment: true + shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" volumes: - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Omni Model Perf Test" - timeout_in_minutes: 120 + - label: ":full_moon: Omni Model Perf Test with H100" + key: nightly-performance + timeout_in_minutes: 180 depends_on: upload-nightly-pipeline if: build.env("NIGHTLY") == "1" commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export BENCHMARK_DIR=tests - pytest -s -v tests/perf/scripts/run_benchmark.py + - buildkite-agent artifact upload "tests/*.json" agents: queue: "mithril-h100-pool" plugins: @@ -96,3 +100,18 @@ steps: hostPath: path: /mnt/hf-cache type: DirectoryOrCreate + + - label: ":email: Nightly Perf Collection & Email" + key: nightly-perf-distribution + depends_on: nightly-performance + if: build.env("NIGHTLY") == "1" + commands: + - pip install openpyxl + - export DEFAULT_INPUT_DIR=tests + - export DEFAULT_OUTPUT_DIR=tests + - buildkite-agent artifact download "tests/*.json" . --step nightly-performance + - python tools/nightly/generate_nightly_perf_excel.py + - python tools/nightly/send_nightly_perf_email.py + - buildkite-agent artifact upload "tests/*.xlsx" + agents: + queue: "cpu_queue_premerge" diff --git a/pyproject.toml b/pyproject.toml index 73ad69c25..95e43d871 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,7 +49,8 @@ dev = [ "imageio[ffmpeg]>=0.6.0", "opencv-python>=4.12.0.88", "mooncake-transfer-engine==0.3.8.post1", - "av" # for ComfyUI tests + "av", # for ComfyUI tests + "openpyxl>=3.0.0", # for nightly CI ] docs = [ diff --git a/tests/perf/scripts/run_benchmark.py b/tests/perf/scripts/run_benchmark.py index 688666162..113c7e826 100644 --- a/tests/perf/scripts/run_benchmark.py +++ b/tests/perf/scripts/run_benchmark.py @@ -1,9 +1,5 @@ -import os - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" - import json +import os import subprocess import threading from datetime import datetime @@ -14,6 +10,9 @@ import pytest from tests.conftest import OmniServer, modify_stage_config +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + def load_configs(config_path: str) -> list[dict[str, Any]]: try: @@ -123,6 +122,8 @@ def run_benchmark(args: list, test_name: str, flow, dataset_name: str, num_promp "--endpoint", "/v1/chat/completions", "--save-result", + "--result-dir", + os.environ.get("BENCHMARK_DIR", "tests"), "--result-filename", result_filename, ] @@ -137,9 +138,9 @@ def run_benchmark(args: list, test_name: str, flow, dataset_name: str, num_promp for line in iter(process.stderr.readline, ""): print(line, end=" ") - if "--result-dir" in args: - index = args.index("--result-dir") - result_dir = args[index + 1] + if "--result-dir" in command: + index = command.index("--result-dir") + result_dir = command[index + 1] else: result_dir = "./" diff --git a/tools/nightly/generate_nightly_perf_excel.py b/tools/nightly/generate_nightly_perf_excel.py new file mode 100644 index 000000000..a9376475d --- /dev/null +++ b/tools/nightly/generate_nightly_perf_excel.py @@ -0,0 +1,437 @@ +#!/usr/bin/env python3 +""" +Generate a nightly Excel performance report from JSON results. + +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +from collections.abc import Iterable, Sequence +from datetime import datetime, timezone +from typing import Any + +from openpyxl import Workbook +from openpyxl.styles import PatternFill +from openpyxl.utils import get_column_letter + +LOGGER = logging.getLogger(__name__) + +GREY_BLOCK_FILL = PatternFill(start_color="D3D3D3", fill_type="solid") + +# Benchmark metric columns: grey the latest row's cell when value changed vs previous date. +BENCHMARK_COLUMNS: tuple[str, ...] = ( + "num_prompts", + "request_rate", + "burstiness", + "max_concurrency", + "duration", + "completed", + "failed", + "request_throughput", + "output_throughput", + "total_token_throughput", + "mean_ttft_ms", + "p99_ttft_ms", + "mean_tpot_ms", + "p99_tpot_ms", + "mean_itl_ms", + "p99_itl_ms", + "mean_e2el_ms", + "p99_e2el_ms", + "mean_audio_rtf", + "p99_audio_rtf", + "mean_audio_duration_s", + "p99_audio_duration_s", +) +# Columns that get float coercion and number format in Excel. Excludes request_rate ("inf" str) +# and max_concurrency (null); leave those as-is. If they become float in the future, they are +# still written correctly from JSON without coercion here. +NUMERIC_FORMAT_COLUMNS: tuple[str, ...] = tuple( + c for c in BENCHMARK_COLUMNS if c not in ("request_rate", "max_concurrency") +) + +_COLUMNS_FILENAME = "nightly_perf_summary_columns.txt" +DEFAULT_INPUT_DIR = os.getenv("DEFAULT_INPUT_DIR") if os.getenv("DEFAULT_INPUT_DIR") else "tests" +DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR") if os.getenv("DEFAULT_OUTPUT_DIR") else "tests" + + +def _load_summary_columns(script_dir: str) -> list[str]: + """Load summary column names from a file next to this script; fallback to default if missing.""" + path = os.path.join(script_dir, _COLUMNS_FILENAME) + default = [ + "date", + "endpoint_type", + "backend", + "model_id", + "tokenizer_id", + "num_prompts", + "request_rate", + "burstiness", + "max_concurrency", + "duration", + "completed", + "failed", + "request_throughput", + "output_throughput", + "total_token_throughput", + "mean_ttft_ms", + "p99_ttft_ms", + "mean_tpot_ms", + "p99_tpot_ms", + "mean_itl_ms", + "p99_itl_ms", + "mean_e2el_ms", + "p99_e2el_ms", + "mean_audio_rtf", + "p99_audio_rtf", + "mean_audio_duration_s", + "p99_audio_duration_s", + "commit_sha", + "build_id", + "build_url", + "source_file", + ] + if not os.path.isfile(path): + return default + columns: list[str] = [] + with open(path, encoding="utf-8") as f: + for line in f: + s = line.strip() + if s and not s.startswith("#"): + columns.append(s) + return columns if columns else default + + +def _vllm_omni_root() -> str: + """Resolve vllm-omni repo root: directory that contains a 'tests' subdir (and usually 'tools').""" + path = os.path.dirname(os.path.abspath(__file__)) + while path and path != os.path.dirname(path): + if os.path.isdir(os.path.join(path, "tests")): + return path + path = os.path.dirname(path) + return os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..")) + + +def _default_input_dir() -> str: + """Default: vllm-omni root / DEFAULT_INPUT_DIR (where performance JSON files live).""" + root = _vllm_omni_root() + return os.path.join(root, DEFAULT_INPUT_DIR) + + +def _default_output_file() -> str: + """Default: vllm-omni root / DEFAULT_OUTPUT_DIR / nightly_perf_.xlsx.""" + ts = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S") + return os.path.join(_vllm_omni_root(), DEFAULT_OUTPUT_DIR, f"nightly_perf_{ts}.xlsx") + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser( + description="Read performance JSON files from vllm-omni/tests/ and generate an Excel report." + ) + parser.add_argument( + "--input-dir", + type=str, + default=_default_input_dir(), + help="Directory containing performance JSON files; default is /DEFAULT_INPUT_DIR.", + ) + parser.add_argument( + "--output-file", + type=str, + default=_default_output_file(), + help="Output path of the Excel report; \ + default is /DEFAULT_OUTPUT_DIR/nightly_perf_.xlsx.", + ) + parser.add_argument( + "--commit-sha", + type=str, + default=None, + help="Optional commit SHA; defaults to environment variable BUILDKITE_COMMIT if unset.", + ) + parser.add_argument( + "--build-id", + type=str, + default=None, + help="Optional build ID; defaults to environment variable BUILDKITE_BUILD_ID if unset.", + ) + parser.add_argument( + "--build-url", + type=str, + default=None, + help="Optional build URL; defaults to environment variable BUILDKITE_BUILD_URL if unset.", + ) + return parser.parse_args() + + +def _load_json_file(path: str) -> dict[str, Any] | None: + """Safely load a single JSON file; return None and log a warning on failure.""" + try: + with open(path, encoding="utf-8") as f: + data = json.load(f) + except (OSError, json.JSONDecodeError) as exc: + LOGGER.warning("failed to load json '%s': %s", path, exc) + return None + + if not isinstance(data, dict): + LOGGER.warning("json root in '%s' is not an object, skip", path) + return None + + return data + + +def _iter_json_records(input_dir: str) -> Iterable[dict[str, Any]]: + """Iterate over JSON files in the input directory and yield normalized records. + commit_sha/build_id/build_url are not set here; they are applied later only to + rows with the latest date (see _apply_build_metadata_to_latest_only). + """ + if not os.path.isdir(input_dir): + LOGGER.warning("input dir '%s' does not exist or is not a directory", input_dir) + return + + for entry in sorted(os.listdir(input_dir)): + if not entry.endswith(".json"): + continue + full_path = os.path.join(input_dir, entry) + if not os.path.isfile(full_path): + continue + + data = _load_json_file(full_path) + if data is None: + continue + + record: dict[str, Any] = dict(data) + record.setdefault("date", datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")) + record["source_file"] = os.path.basename(full_path) + yield record + + +def _collect_records(input_dir: str) -> list[dict[str, Any]]: + """Collect all JSON records into a list.""" + records: list[dict[str, Any]] = [] + for record in _iter_json_records(input_dir): + records.append(record) + return records + + +def _apply_build_metadata_to_latest_only( + records: Sequence[dict[str, Any]], + commit_sha: str | None, + build_id: str | None, + build_url: str | None, +) -> None: + """Set commit_sha, build_id, build_url only on rows with the latest date. + Other rows get None so that build info is not duplicated for older benchmark data. + """ + if not records: + return + max_date = max((r.get("date") or "") for r in records) + for r in records: + if (r.get("date") or "") == max_date: + r["commit_sha"] = commit_sha + r["build_id"] = build_id + r["build_url"] = build_url + else: + r["commit_sha"] = None + r["build_id"] = None + r["build_url"] = None + + +def _sort_records_for_summary(records: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Sort so that same model_id is grouped, newest date first within each group.""" + by_date_desc = sorted(records, key=lambda r: (r.get("date") or ""), reverse=True) + return sorted(by_date_desc, key=lambda r: (r.get("model_id") or "")) + + +def _values_differ(a: Any, b: Any) -> bool: + """Return True if two values are considered different; avoid direct float equality.""" + if a is None and b is None: + return False + if a is None or b is None: + return True + if isinstance(a, float) and isinstance(b, float): + if a != a and b != b: + return False + if a != a or b != b: + return True + return abs(a - b) > 1e-9 + return a != b + + +def _apply_benchmark_change_highlight( + ws, + summary_columns: Sequence[str], + records: Sequence[dict[str, Any]], +) -> None: + """Grey cells in the latest row of each model when a benchmark metric changed vs previous date.""" + if not records: + return + col_to_index = {c: i + 1 for i, c in enumerate(summary_columns)} + # Walk by model_id blocks (records already sorted by model_id, date desc). + i = 0 + while i < len(records): + model_id = records[i].get("model_id") + block_start = i + while i < len(records) and records[i].get("model_id") == model_id: + i += 1 + block_end = i + newest_idx = block_start + prev_idx = block_start + 1 if block_start + 1 < block_end else None + if prev_idx is None: + continue + excel_row = newest_idx + 2 + for col in BENCHMARK_COLUMNS: + if col not in col_to_index: + continue + cur_val = records[newest_idx].get(col) + prev_val = records[prev_idx].get(col) + if _values_differ(cur_val, prev_val): + ws.cell(row=excel_row, column=col_to_index[col]).fill = GREY_BLOCK_FILL + + +def _build_raw_columns( + records: Sequence[dict[str, Any]], + summary_columns: Sequence[str], +) -> list[str]: + """Infer the column set for the raw sheet based on all records.""" + keys: set[str] = set() + for record in records: + keys.update(record.keys()) + # Ensure summary columns appear first; remaining columns sorted alphabetically. + ordered_keys: list[str] = [] + for key in summary_columns: + if key in keys: + ordered_keys.append(key) + keys.discard(key) + ordered_keys.extend(sorted(keys)) + return ordered_keys + + +def _to_float_if_numeric(value: Any) -> Any: + """Coerce to float when possible so Excel treats as number; avoid #### from narrow columns.""" + if value is None: + return value + if isinstance(value, (int, float)): + return float(value) if isinstance(value, int) else value + if isinstance(value, str): + try: + return float(value) + except (ValueError, TypeError): + return value + return value + + +def _write_sheet( + ws, + columns: Sequence[str], + rows: Iterable[dict[str, Any]], + numeric_columns: Sequence[str] = (), +) -> None: + """Write column names and row data into the given worksheet.""" + numeric_set = set(numeric_columns) + ws.append(list(columns)) + for record in rows: + row_values = [] + for col in columns: + v = record.get(col) + if col in numeric_set: + v = _to_float_if_numeric(v) + row_values.append(v) + ws.append(row_values) + + +def _format_benchmark_columns( + ws, + columns: Sequence[str], + num_data_rows: int, +) -> None: + """Set number format and column width for numeric benchmark columns so values display (no ####).""" + numeric_set = set(NUMERIC_FORMAT_COLUMNS) + for c, col_name in enumerate(columns): + if col_name not in numeric_set: + continue + col_letter = get_column_letter(c + 1) + ws.column_dimensions[col_letter].width = 14 + for r in range(2, 2 + num_data_rows): + cell = ws.cell(row=r, column=c + 1) + if cell.value is not None and isinstance(cell.value, (int, float)): + cell.number_format = "0.0000" + elif isinstance(cell.value, str): + try: + cell.value = float(cell.value) + cell.number_format = "0.0000" + except (ValueError, TypeError): + pass + + +def _ensure_parent_dir(path: str) -> None: + """Ensure that the parent directory of the output file exists.""" + parent = os.path.dirname(os.path.abspath(path)) + if not parent: + return + os.makedirs(parent, exist_ok=True) + + +def generate_excel_report( + input_dir: str, + output_file: str, + commit_sha: str | None, + build_id: str | None, + build_url: str | None, +) -> None: + """Main logic: load JSON records and generate an Excel report.""" + script_dir = os.path.dirname(os.path.abspath(__file__)) + summary_columns = _load_summary_columns(script_dir) + + records = _collect_records(input_dir) + if not records: + LOGGER.warning("no valid json records found under '%s'", input_dir) + + sorted_records = _sort_records_for_summary(records) + _apply_build_metadata_to_latest_only(sorted_records, commit_sha, build_id, build_url) + + wb = Workbook() + ws_summary = wb.active + ws_summary.title = "summary" + + _write_sheet(ws_summary, summary_columns, sorted_records, numeric_columns=NUMERIC_FORMAT_COLUMNS) + _format_benchmark_columns(ws_summary, summary_columns, len(sorted_records)) + _apply_benchmark_change_highlight(ws_summary, summary_columns, sorted_records) + + if sorted_records: + raw_columns = _build_raw_columns(sorted_records, summary_columns) + ws_raw = wb.create_sheet(title="raw") + _write_sheet(ws_raw, raw_columns, sorted_records) + + _ensure_parent_dir(output_file) + wb.save(output_file) + LOGGER.info("excel report saved to '%s'", output_file) + + +def main() -> None: + """Command-line entrypoint.""" + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + ) + + args = parse_args() + + commit_sha = args.commit_sha or os.getenv("BUILDKITE_COMMIT") + build_id = args.build_id or os.getenv("BUILDKITE_BUILD_ID") + build_url = args.build_url or os.getenv("BUILDKITE_BUILD_URL") + + generate_excel_report( + input_dir=args.input_dir, + output_file=args.output_file, + commit_sha=commit_sha, + build_id=build_id, + build_url=build_url, + ) + + +if __name__ == "__main__": + main() diff --git a/tools/nightly/send_nightly_perf_email.py b/tools/nightly/send_nightly_perf_email.py new file mode 100644 index 000000000..af11cc73a --- /dev/null +++ b/tools/nightly/send_nightly_perf_email.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +""" +Send the nightly performance Excel report by email to DAILY_EMAIL_LIST. + +Reads SMTP and recipient config from environment variables. Use --dry-run to +validate config and print subject/body without sending. +""" + +from __future__ import annotations + +import argparse +import logging +import os +import smtplib +import sys +from email.mime.application import MIMEApplication +from email.mime.multipart import MIMEMultipart +from email.mime.text import MIMEText + +LOGGER = logging.getLogger(__name__) + +# Env keys for SMTP and recipients. +ENV_SMTP_HOST = "SMTP_HOST" +ENV_SMTP_PORT = "SMTP_PORT" +ENV_SMTP_USERNAME = "SMTP_USERNAME" +ENV_SMTP_PASSWORD = "SMTP_PASSWORD" +ENV_DAILY_EMAIL_LIST = "DAILY_EMAIL_LIST" +ENV_EMAIL_SENDER = "EMAIL_SENDER" +ENV_EMAIL_SUBJECT_PREFIX = "EMAIL_SUBJECT_PREFIX" +ENV_BUILD_URL = "BUILDKITE_BUILD_URL" +ENV_COMMIT = "BUILDKITE_COMMIT" + +DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR") + +# Do not attach Excel if size >= this (bytes); body will suggest downloading from build URL. +MAX_ATTACHMENT_BYTES = 20 * 1024 * 1024 + +SMTP_RETRIES = 3 +SMTP_RETRY_DELAY_SEC = 5 + + +def _get_required_env() -> dict[str, str]: + """Read required env vars; raise SystemExit with clear message if any missing.""" + required = { + ENV_SMTP_HOST: os.environ.get(ENV_SMTP_HOST), + ENV_SMTP_PORT: os.environ.get(ENV_SMTP_PORT), + ENV_SMTP_USERNAME: os.environ.get(ENV_SMTP_USERNAME), + ENV_SMTP_PASSWORD: os.environ.get(ENV_SMTP_PASSWORD), + ENV_DAILY_EMAIL_LIST: os.environ.get(ENV_DAILY_EMAIL_LIST), + } + missing = [k for k, v in required.items() if not (v and str(v).strip())] + if missing: + raise SystemExit(f"Missing required env vars: {', '.join(missing)}. Set them (e.g. in Buildkite secrets).") + return {k: str(v).strip() for k, v in required.items()} + + +def _get_latest_file(folder_path: str) -> str: + """Get the latest modified file from the folder path.""" + files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(".xlsx")] + if not files: + raise SystemExit(f"No Excel files found in {folder_path}") + return max(files, key=os.path.getmtime) + + +def _recipients_list(comma_separated: str) -> list[str]: + """Parse DAILY_EMAIL_LIST into a list of addresses.""" + return [a.strip() for a in comma_separated.split(",") if a.strip()] + + +def _build_body( + date_str: str, + commit_sha: str | None, + build_url: str | None, + attachment_skipped: bool, +) -> str: + """Plain-text body with build metadata; note when attachment was skipped.""" + lines = [ + f"Nightly performance report date: {date_str}", + "", + f"Commit: {commit_sha or 'N/A'}", + f"Build: {build_url or 'N/A'}", + "", + ] + if attachment_skipped: + lines.append( + "Report file was too large to attach. Please download the Excel from the build artifacts (Build URL above)." + ) + lines.append("") + return "\n".join(lines) + + +def _build_subject(prefix: str | None, date_str: str) -> str: + """Subject line: optional prefix + date.""" + base = f"Nightly Perf {date_str}" + if prefix and prefix.strip(): + return f"{prefix.strip()} {base}" + return base + + +def _send_mail( + report_file: str, + date_str: str, + dry_run: bool, +) -> None: + """Load config, build message, and send (or dry-run).""" + cfg = _get_required_env() + recipients = _recipients_list(cfg[ENV_DAILY_EMAIL_LIST]) + if not recipients: + raise SystemExit("DAILY_EMAIL_LIST is empty after parsing.") + + commit_sha = os.environ.get(ENV_COMMIT) + build_url = os.environ.get(ENV_BUILD_URL) + sender = os.environ.get(ENV_EMAIL_SENDER) or cfg[ENV_SMTP_USERNAME] + prefix = os.environ.get(ENV_EMAIL_SUBJECT_PREFIX) + + size = os.path.getsize(report_file) if os.path.isfile(report_file) else 0 + attach_excel = size < MAX_ATTACHMENT_BYTES and size > 0 + attachment_skipped = os.path.isfile(report_file) and not attach_excel + + body = _build_body( + commit_sha=commit_sha, build_url=build_url, date_str=date_str, attachment_skipped=attachment_skipped + ) + subject = _build_subject(prefix=prefix, date_str=date_str) + + msg = MIMEMultipart() + msg["Subject"] = subject + msg["From"] = sender + msg["To"] = ", ".join(recipients) + msg.attach(MIMEText(body, "plain", "utf-8")) + + if attach_excel: + with open(report_file, "rb") as f: + part = MIMEApplication(f.read(), _subtype="vnd.openxmlformats-officedocument.spreadsheetml.sheet") + part.add_header("Content-Disposition", "attachment", filename=os.path.basename(report_file)) + msg.attach(part) + + if dry_run: + LOGGER.info("dry-run: not sending mail") + print("To:", recipients, file=sys.stderr) + print("Subject:", subject, file=sys.stderr) + print("Attachment:", "yes" if attach_excel else "no (size limit)", file=sys.stderr) + print("Body preview:", body[:300] + ("..." if len(body) > 300 else ""), file=sys.stderr) + return + + port = int(cfg[ENV_SMTP_PORT], 10) + last_err: Exception | None = None + for attempt in range(SMTP_RETRIES): + try: + with smtplib.SMTP(cfg[ENV_SMTP_HOST], port=port, timeout=30) as smtp: + smtp.starttls() + smtp.login(cfg[ENV_SMTP_USERNAME], cfg[ENV_SMTP_PASSWORD]) + smtp.sendmail(sender, recipients, msg.as_string()) + LOGGER.info("sent nightly perf email to %d recipient(s)", len(recipients)) + return + except Exception as e: + last_err = e + LOGGER.warning("SMTP attempt %d/%d failed: %s", attempt + 1, SMTP_RETRIES, e) + if attempt < SMTP_RETRIES - 1: + import time + + time.sleep(SMTP_RETRY_DELAY_SEC) + raise SystemExit(f"Failed to send email after {SMTP_RETRIES} attempts.") from last_err + + +def _date_from_filename(path: str) -> str: + """Try to derive a date string from report filename (e.g. nightly_perf_20260211-020704.xlsx -> 2026-02-11).""" + base = os.path.splitext(os.path.basename(path))[0] + if base.startswith("nightly_perf_") and len(base) > 13: + raw = base.replace("nightly_perf_", "")[:8] + if len(raw) == 8 and raw.isdigit(): + return f"{raw[:4]}-{raw[4:6]}-{raw[6:8]}" + return base or "unknown" + + +def _vllm_omni_root() -> str: + """Resolve vllm-omni repo root: directory that contains a 'tests' subdir (and usually 'tools').""" + path = os.path.dirname(os.path.abspath(__file__)) + while path and path != os.path.dirname(path): + if os.path.isdir(os.path.join(path, "tests")): + return path + path = os.path.dirname(path) + return os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..")) + + +def _default_output_dir() -> str: + """Default: vllm-omni root / DEFAULT_OUTPUT_DIR (where performance .xlsx files live).""" + root = _vllm_omni_root() + return os.path.join(root, DEFAULT_OUTPUT_DIR) + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser( + description="Send nightly performance Excel report by email (config from env).", + ) + parser.add_argument( + "--report-file", + type=str, + default=_default_output_dir(), + help="Folder/file path to the nightly_perf_*.xlsx file; default is DEFAULT_OUTPUT_DIR.", + ) + parser.add_argument( + "--date", + type=str, + default=None, + help="Date string for subject/body (default: inferred from report filename).", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print recipient, subject, and body; do not send.", + ) + return parser.parse_args() + + +def main() -> None: + """Entrypoint.""" + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + ) + args = parse_args() + + report_file = _get_latest_file(args.report_file) if os.path.isdir(args.report_file) else args.report_file + if not os.path.isfile(report_file): + raise SystemExit(f"Report file not found: {report_file}") + + date_str = args.date or _date_from_filename(report_file) + _send_mail(report_file=report_file, date_str=date_str, dry_run=args.dry_run) + + +if __name__ == "__main__": + main() From d275c0c993e09bef5e6dfd560c67f6db69ff5f95 Mon Sep 17 00:00:00 2001 From: dengyunyang <584797741@qq.com> Date: Tue, 3 Mar 2026 15:48:51 +0800 Subject: [PATCH 5/7] [FEATURE] Tile/Patch parallelism refactor for easily support other models (#1366) Signed-off-by: dengyunyang <584797741@qq.com> --- docs/design/feature/vae_parallel.md | 289 ++++++++++++++++++ .../diffusion/parallelism_acceleration.md | 10 +- docs/user_guide/diffusion_acceleration.md | 12 +- .../image_to_video/image_to_video.py | 19 +- .../text_to_video/text_to_video.py | 11 +- .../test_distributed_vae_executor.py | 165 ++++++++++ .../test_vae_decode_parallelism.py | 154 ++++++++++ .../distributed/autoencoders/__init__.py | 0 .../autoencoders/autoencoder_kl.py | 210 +++++++++++++ .../autoencoders/autoencoder_kl_qwenimage.py | 120 ++++++++ .../autoencoders/autoencoder_kl_wan.py | 142 +++++++++ .../autoencoders/distributed_vae_executor.py | 193 ++++++++++++ .../models/qwen_image/pipeline_qwen_image.py | 10 +- .../diffusion/models/sd3/pipeline_sd3.py | 8 +- .../models/wan2_2/pipeline_wan2_2.py | 4 +- .../models/wan2_2/pipeline_wan2_2_i2v.py | 4 +- .../models/z_image/pipeline_z_image.py | 8 +- vllm_omni/diffusion/registry.py | 37 +-- 18 files changed, 1336 insertions(+), 60 deletions(-) create mode 100644 docs/design/feature/vae_parallel.md create mode 100644 tests/diffusion/distributed/test_distributed_vae_executor.py create mode 100644 tests/e2e/offline_inference/test_vae_decode_parallelism.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/__init__.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py create mode 100644 vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py diff --git a/docs/design/feature/vae_parallel.md b/docs/design/feature/vae_parallel.md new file mode 100644 index 000000000..9009ece72 --- /dev/null +++ b/docs/design/feature/vae_parallel.md @@ -0,0 +1,289 @@ +# VAE Patch Parallelism + +This document describes how to add **VAE Patch Parallelism** support to a diffusion model. +We use **Qwen-Image** as the reference implementation. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Step-by-Step Implementation](#step-by-step-implementation) +- [Testing](#testing) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + +### What is Vae Patch parallel? + +**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is: + ++ Split into multiple spatial tiles + ++ Distributed across multiple ranks + ++ Decoded in parallel + ++ Merged to reconstruct the final output + +This approach: + ++ Distributes computation across multiple devices + ++ Reduces peak memory usage per device + ++ Accelerates decoding latency + +### Architecture +We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding. + +The executor is model-agnostic and accepts three function parameters: + ++ split – Partition the latent into tiles + ++ exec – Decode a single tile + ++ merge – Combine decoded tiles into the final output + +#### Execution Flow + ++ Call split(z) to generate a list of TileTask and a GridSpec + ++ Dispatch tasks across ranks using workload-based balancing + ++ Each rank executes exec(task) on its assigned tiles + ++ Gather decoded tile results to rank 0 + ++ Rank 0 performs merge(...) + ++ (Optional) Broadcast final result to all ranks + +This design separates: + ++ Distributed execution logic + ++ Model-specific tiling and merging logic + +#### Why split / exec / merge is necessary? + +The latent tensor cannot be arbitrarily partitioned. + +During decoding: + ++ Each output pixel may depend on neighboring pixels + ++ The receptive field is model-dependent + +Therefore: + ++ Tiles must include overlap + ++ Merge must perform blending to avoid seams + +## Step-by-Step Implementation + +### Step 1: Implement DistributedAutoencoderKLQwenImage +`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version: + + +``` +class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model +``` +**Key points**: ++ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin ++ Call init_distributed() after loading weights + +### Step 2: Implement split/exec/merge +Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`: +``` +class GridSpec: + split_dims: tuple[int, ...] # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W)) + grid_shape: tuple[int, ...] # Tile grid layout (num_rows, num_cols) + tile_spec: dict = field(default_factory=dict) # Metadata required for merging + output_dtype: torch.dtype | None = None # Final output dtype +``` +``` +class TileTask: + tile_id: int # task id + grid_coord: tuple[int, ...] # Tile position in grid + tensor: torch.Tensor | list[torch.Tensor] # The tile tensor + workload: int | float = 1 # Used for load balancing (e.g., tile area) +``` +And tiled base split/exec/merge as follow: +``` +def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + # mostly copy from AutoencoderKL + _, _, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + tiletask_list = [] + for i in range(0, height, tile_latent_stride_height): + for j in range(0, width, tile_latent_stride_width): + time_list = [] + for k in range(num_frames): + self._conv_idx = [0] + tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + time_list.append(tile) + tiletask_list.append( + TileTask( + len(tiletask_list), + (i // tile_latent_stride_height, j // tile_latent_stride_width), + time_list, + workload=time_list[0].shape[3] * time_list[0].shape[4], + ) + ) + tile_spec = { + "sample_height": sample_height, + "sample_width": sample_width, + "blend_height": blend_height, + "blend_width": blend_width, + } + 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=tile_spec, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + +def tile_exec(self, task: TileTask) -> torch.Tensor: + """Decode a single latent tile into RGB space.""" + self.clear_cache() + time = [] + for k in range(len(task.tensor)): + self._conv_idx = [0] + tile = self.post_quant_conv(task.tensor[k]) + decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx) + time.append(decoded) + result = torch.cat(time, dim=2) + return result + +def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + """Merge decoded tiles into a full image.""" + grid_h, grid_w = grid_spec.grid_shape + result_rows = [] + self.clear_cache() + + 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_height"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) + result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=3)[ + :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"] + ] + return dec +``` + +### Step 3: Override tiled_decode +We need to override tiled_decode, the main logic is: ++ check distributed is enabled ++ select split/exec/merge ++ Invoke self.distributed_decoder.execute to decode +``` +def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): + if not self.is_distributed_enabled(): + return super().tiled_decode(z, return_dict=return_dict) + + logger.info("Decode run with distributed executor") + result = self.distributed_decoder.execute( + z, + DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge), + broadcast_result=True, + ) + if not return_dict: + return (result,) + + return DecoderOutput(sample=result) +``` +`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0. + +### Step 4: Modify Pipeline +Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage +``` +class YourModelPipeline(nn.Module): + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + ... +- self.vae = AutoencoderKL.from_pretrained( +- model, subfolder="vae", local_files_only=local_files_only).to(self.device) ++ self.vae = DistributedAutoencoderKL.from_pretrained( ++ model, subfolder="vae", local_files_only=local_files_only ++ ).to(self.device) +``` + +## Testing +Verify numerical consistency between: ++ vae_patch_parallel_size = 1 + ++ vae_patch_parallel_size = N + +Example: +torch.allclose(output_1, output_n, atol=1e-5) + +Testing requirements: ++ Fix random seed ++ Use identical tiling strategy + +```python +m = Omni( + model=model_name, + vae_use_tiling=True, + parallel_config=DiffusionParallelConfig( + tensor_parallel_size=2, + vae_patch_parallel_size=1, # or 2 + ), + ) +``` +When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead. + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Notes | +|-------|------|-------| +| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL | +| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan | +| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage | + +--- + +## Summary + +Adding Vae Patch Parallel support to diffusion model: + +1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge +2. **Change vae model in pipeline to Distributed Vae** +3. **Test** - Verify with `tensor_parallel_size=N` quality diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index b0225e2ff..4f4a209d0 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -27,7 +27,7 @@ The following table shows which models are currently supported by parallelism me | **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | | **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | | **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | -| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | @@ -51,9 +51,9 @@ The following table shows which models are currently supported by parallelism me ### VideoGen -| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | -|-------|------------------|:----------:|:--------------:|:---------------:|:----:| -| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | +| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel | +|-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:| +| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ | ### Tensor Parallelism @@ -85,7 +85,7 @@ outputs = omni.generate( VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP). !!! note "Enablement and feature gate" - - VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`, `stabilityai/stable-diffusion-3.5`). + - VAE patch parallelism is currently **enabled only for validated pipelines** (check [ImageGen](#imagegen) and [VideoGen](#videogen) for more information). - If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.) #### Offline Inference diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 9fd14a2b0..2c36c966c 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -58,8 +58,8 @@ The following table shows which models are currently supported by each accelerat | **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | | **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | | **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | -| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | -| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | @@ -67,13 +67,13 @@ The following table shows which models are currently supported by each accelerat | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | -| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | +| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ### VideoGen -| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | -|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:| -| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | +| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel | +|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:| +| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ### Quantization diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index a491ddfc3..4e42b8b29 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -20,6 +20,7 @@ Usage: import argparse import os +import time from pathlib import Path import numpy as np @@ -105,6 +106,12 @@ def parse_args() -> argparse.Namespace: choices=[1, 2], help="Number of GPUs used for classifier free guidance parallel size.", ) + parser.add_argument( + "--vae-patch-parallel-size", + type=int, + default=1, + help="Number of GPUs used for VAE patch/tile parallelism (decode).", + ) parser.add_argument( "--enforce-eager", action="store_true", @@ -173,6 +180,7 @@ def main(): ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size, tensor_parallel_size=args.tensor_parallel_size, + vae_patch_parallel_size=args.vae_patch_parallel_size, use_hsdp=args.use_hsdp, hsdp_shard_size=args.hsdp_shard_size, hsdp_replicate_size=args.hsdp_replicate_size, @@ -199,10 +207,14 @@ def main(): print(f" Model: {args.model}") print(f" Inference steps: {args.num_inference_steps}") print(f" Frames: {args.num_frames}") - print(f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}") + print( + f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}," + f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}" + ) print(f" Video size: {args.width}x{args.height}") print(f"{'=' * 60}\n") + generation_start = time.perf_counter() # omni.generate() returns Generator[OmniRequestOutput, None, None] frames = omni.generate( { @@ -220,6 +232,11 @@ def main(): num_frames=args.num_frames, ), ) + generation_end = time.perf_counter() + generation_time = generation_end - generation_start + + # Print profiling results + print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)") # Extract video frames from OmniRequestOutput if isinstance(frames, list) and len(frames) > 0: diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 40fafa100..5db1b73dc 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -109,6 +109,12 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for tensor parallelism (TP) inside the DiT.", ) + parser.add_argument( + "--vae-patch-parallel-size", + type=int, + default=1, + help="Number of GPUs used for VAE patch/tile parallelism (decode).", + ) return parser.parse_args() @@ -141,6 +147,7 @@ def main(): ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size, tensor_parallel_size=args.tensor_parallel_size, + vae_patch_parallel_size=args.vae_patch_parallel_size, ) # Check if profiling is requested via environment variable @@ -172,7 +179,9 @@ def main(): print(f" Inference steps: {args.num_inference_steps}") print(f" Frames: {args.num_frames}") print( - f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}" + f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}," + f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}," + f" vae_patch_parallel_size={args.vae_patch_parallel_size}" ) print(f" Video size: {args.width}x{args.height}") print(f"{'=' * 60}\n") diff --git a/tests/diffusion/distributed/test_distributed_vae_executor.py b/tests/diffusion/distributed/test_distributed_vae_executor.py new file mode 100644 index 000000000..b53014f39 --- /dev/null +++ b/tests/diffusion/distributed/test_distributed_vae_executor.py @@ -0,0 +1,165 @@ +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedVaeExecutor, + DistributedVaeMixin, + GridSpec, + TileTask, +) + + +class E2EOperator: + """tiles with (2, 3) -- (H,W)""" + + def split(self, z): + rows_num = 2 + cols_num = 3 + + h_size = z.shape[0] + w_size = z.shape[1] + + tasks = [] + for i in range(rows_num): + for j in range(cols_num): + tasks.append( + TileTask( + tile_id=len(tasks), + grid_coord=(i, j), + tensor=z[ + ((i * h_size) // rows_num) : (((i + 1) * h_size) // rows_num), + ((j * w_size) // cols_num) : (((j + 1) * w_size) // cols_num), + ], + ) + ) + + grid_spec = GridSpec( + split_dims=(0, 1), + grid_shape=(2, 3), + ) + return tasks, grid_spec + + def exec(self, task: TileTask): + return torch.full_like(task.tensor, fill_value=task.tile_id) + + def merge(self, coord_tensor_map, grid_spec): + tiles = [] + for r in range(grid_spec.grid_shape[0]): + row_tiles = [] + for c in range(grid_spec.grid_shape[1]): + coord = (r, c) + row_tiles.append(coord_tensor_map[coord]) + tiles.append(torch.cat(row_tiles, dim=1)) + return torch.cat(tiles, dim=0) + + +class DummyMixin(DistributedVaeMixin): + def __init__(self): + self.use_tiling = True + self.distributed_decoder = MagicMock() + self.distributed_decoder.parallel_size = 2 + self.distributed_decoder.group = None + + +@pytest.fixture(autouse=True) +def mock_dist(): + with ( + patch.object(dist, "get_world_size", return_value=2), + patch.object(dist, "get_rank", return_value=0), + patch.object(dist, "is_initialized", return_value=True), + patch.object(dist, "gather", return_value=None), + patch.object(dist, "broadcast", return_value=None), + ): + yield + + +@pytest.fixture(autouse=True) +def mock_dit_group(): + with patch( + "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group", + new=MagicMock(return_value=None), + ): + yield + + +@pytest.fixture(autouse=True) +def mock_dist_vae_executor(): + with ( + patch.object(DistributedVaeExecutor, "gather_tensors", side_effect=lambda x: [x]), + patch.object(DistributedVaeExecutor, "broadcast_tensor", side_effect=lambda x: x), + ): + yield + + +# ============================ +# Unitest +# ============================ + + +def test_balance_tasks(): + executor = DistributedVaeExecutor() + workloads = [2, 5, 13, 8, 2, 4] # each is 17 + tasks = [TileTask(0, (i,), torch.tensor([i]), workload=load) for i, load in enumerate(workloads)] + assigned = executor._balance_tasks(tasks, num_rank=2) + assert len(assigned) == 2 + total_work = [sum(t.workload for t in group) for group in assigned] + assert total_work[0] == total_work[1] + + +def test_compute_global_padding_shape(): + executor = DistributedVaeExecutor() + executor.rank = 0 + + local_results = [(0, torch.zeros((2, 3))), (1, torch.zeros((4, 2)))] + shape = executor._compute_global_padding_shape(local_results, 2, "cpu") + + assert shape == [2, 4, 3] + + +def test_pack_and_unpack(): + executor = DistributedVaeExecutor() + executor.world_size = 1 + + grid_spec = GridSpec(split_dims=(0, 1), grid_shape=(2, 2)) + + # ====================== + # pack + # ====================== + local_results = [(0, torch.tensor([[1, 2], [3, 4]]))] + + global_shape = [3, 3, 3] # (tiles, H, W) + + tile_tensor, meta_tensor = executor._pack_local_tiles( + local_results, global_shape, grid_spec, device="cpu", dtype=torch.int64 + ) + + # check pack + assert tile_tensor.shape == torch.Size(global_shape) + assert meta_tensor.shape == (global_shape[0], len(grid_spec.split_dims) + 1) + assert meta_tensor[0, 0] == 0 + assert meta_tensor[0, 1] == 2 + assert meta_tensor[0, 2] == 2 + + # ====================== + # unpack + # ====================== + meta_gather = [meta_tensor] + tile_gather = [tile_tensor] + + tid_coord_map = {0: (0, 0)} + + coord_tensor_map = executor._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map) + + # check unpack + assert torch.equal(coord_tensor_map[(0, 0)], torch.tensor([[1, 2], [3, 4]])) + + +def test_is_distributed_enabled(): + mixin = DummyMixin() + assert mixin.is_distributed_enabled() is True + + mixin.use_tiling = False + assert mixin.is_distributed_enabled() is False diff --git a/tests/e2e/offline_inference/test_vae_decode_parallelism.py b/tests/e2e/offline_inference/test_vae_decode_parallelism.py new file mode 100644 index 000000000..0bd3e9218 --- /dev/null +++ b/tests/e2e/offline_inference/test_vae_decode_parallelism.py @@ -0,0 +1,154 @@ +import os +import sys +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image +from vllm.distributed.parallel_state import cleanup_dist_env_and_memory + +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +# ruff: noqa: E402 +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +import time + +from vllm_omni import Omni +from vllm_omni.platforms import current_omni_platform + +# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +all_model_configs = [ + { + "model_name": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", + "out_width": 1280, + "out_height": 720, + "out_frames": 16, + }, + { + "model_name": "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + "out_width": 1280, + "out_height": 720, + "out_frames": 16, + }, + { + "model_name": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", + "out_width": 1280, + "out_height": 720, + "out_frames": 16, + }, + { + "model_name": "Qwen/Qwen-Image", + "out_width": 3072, + "out_height": 3072, + "out_frames": 1, + }, + # { + # "model_name": "stepfun-ai/NextStep-1.1", + # "out_width": 3072, + # "out_height": 3072, + # "out_frames": 1, + # }, +] + + +def is_nextstep_model(model_name: str) -> bool: + """Check if the model is a NextStep model by reading its config.""" + from vllm.transformers_utils.config import get_hf_file_to_dict + + try: + cfg = get_hf_file_to_dict("config.json", model_name) + if cfg and cfg.get("model_type") == "nextstep": + return True + except Exception: + pass + return False + + +def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1): + try: + parallel_config = DiffusionParallelConfig( + tensor_parallel_size=tp, + vae_patch_parallel_size=vae_patch_parallel_size, + ) + + omni_kwargs = { + "model": model_configs["model_name"], + "vae_use_tiling": using_tile, + "parallel_config": parallel_config, + } + use_nextstep = is_nextstep_model(model_configs["model_name"]) + if use_nextstep: + # NextStep-1.1 requires explicit pipeline class + omni_kwargs["model_class_name"] = "NextStep11Pipeline" + m = Omni(**omni_kwargs) + image = Image.new("RGB", (out_width, out_height), (0, 0, 0)) + start = time.perf_counter() + outputs = m.generate( + { + "prompt": "A cat sitting on a table", + "multi_modal_data": {"image": image}, + }, + sampling_params_list=OmniDiffusionSamplingParams( + height=model_configs["out_height"], + width=model_configs["out_width"], + num_frames=model_configs["out_frames"], + num_inference_steps=2, + generator=torch.Generator(current_omni_platform.device_type).manual_seed(42), + ), + ) + end = time.perf_counter() + first_output = outputs[0] + req_out = first_output.request_output[0] + frames = req_out.images[0] + if isinstance(frames, torch.Tensor): + frames = frames.detach().cpu().numpy() + elif isinstance(frames, Image.Image): + frames = np.array(frames) + # frames shape: (batch, num_frames, height, width, channels) + cost = (end - start) * 1000 + return frames, cost + finally: + m.close() + cleanup_dist_env_and_memory() + + +@pytest.mark.parametrize("model_configs", all_model_configs) +def test_vae_parallel_model(model_configs: str): + parallel_result, parallel_time = model_run( + model_configs=model_configs, + tp=2, + out_width=1280, + out_height=704, + out_frames=5, + using_tile=True, + vae_patch_parallel_size=2, + ) + non_parallel_result, non_parallel_time = model_run( + model_configs=model_configs, + tp=2, + out_width=1280, + out_height=704, + out_frames=5, + using_tile=True, + vae_patch_parallel_size=1, + ) + result_diff = np.abs(non_parallel_result - parallel_result) + + mean_threshold = 3e-2 + max_threshold = 3e-2 # they should be totally same + print( + f"{model_configs['model_name']} TP = 2 (tile + parallel vs tile): " + f"mean_abs_diff={result_diff.mean():.6e}, max_abs_diff={result_diff.max():.6e}; " + f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; " + f"parallel generate take time: {parallel_time:.2f} ms, non-parallel take time: {non_parallel_time:.2f} ms" + ) + # assert non_parallel_time > parallel_time + assert result_diff.mean() < mean_threshold + assert result_diff.max() < max_threshold diff --git a/vllm_omni/diffusion/distributed/autoencoders/__init__.py b/vllm_omni/diffusion/distributed/autoencoders/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py new file mode 100644 index 000000000..7df2d6a8a --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from typing import Any + +import torch +from diffusers.models.autoencoders import AutoencoderKL as Diffusers_AutoencoderKL +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.nextstep_1_1.modeling_flux_vae import AutoencoderKL as Next_Step_AutoencoderKL + +logger = init_logger(__name__) + + +# We use base class because some model re-implement AutoencoderKL, but split is share. +class DistributedAutoencoderKL_base(DistributedVaeMixin): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + # mostly copy from AutoencoderKL + 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 = self.tile_sample_min_size - blend_extent + + # Split z into overlapping 64x64 tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + tiletask_list = [] + for i in range(0, z.shape[2], overlap_size): + for j in range(0, z.shape[3], 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[2] * tile.shape[3], + ) + ) + + tile_spec = { + "blend_extent": blend_extent, + "row_limit": row_limit, + } + grid_spec = GridSpec( + split_dims=(2, 3), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec=tile_spec, + ) + return tiletask_list, grid_spec + + def tile_exec(self, task: TileTask) -> torch.Tensor: + """Decode a single latent tile into RGB space.""" + tile = task.tensor + if self.config.use_post_quant_conv: + tile = self.post_quant_conv(tile) + decoded = self.decoder(tile) + return decoded + + def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + """Merge decoded tiles into a full image.""" + + 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=3)) + + dec = torch.cat(result_rows, dim=2) + return dec + + def patch_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + overlap_latent = int(self.tile_latent_min_size * float(self.tile_overlap_factor)) + halo_base = max(0, overlap_latent // 2) + + _, _, latent_h, latent_w = z.shape + scale = int(2 ** (len(self.config.block_out_channels) - 1)) + max_parallel_size = self.distributed_decoder.parallel_size + + root = int(math.sqrt(max_parallel_size)) + for rows in range(root, 0, -1): + if max_parallel_size % rows == 0: + grid_rows, grid_cols = rows, max_parallel_size // rows + break + tiletask_list = [] + halo_size = dict() + for i in range(grid_rows): + for j in range(grid_cols): + h0 = (i * latent_h) // grid_rows + h1 = ((i + 1) * latent_h) // grid_rows + w0 = (j * latent_w) // grid_cols + w1 = ((j + 1) * latent_w) // grid_cols + + core_h = max(0, h1 - h0) + core_w = max(0, w1 - w0) + halo = max(halo_base, min(core_h, core_w) // 2) + ph0 = max(0, h0 - halo) + ph1 = min(latent_h, h1 + halo) + pw0 = max(0, w0 - halo) + pw1 = min(latent_w, w1 + halo) + tile = z[:, :, ph0:ph1, pw0:pw1] + tiletask_list.append(TileTask(len(tiletask_list), (i, j), tile, tile.shape[2] * tile.shape[3])) + halo_size[(i, j)] = { + "up": h0 - ph0, + "down": ph1 - h1, + "left": w0 - pw0, + "right": pw1 - w1, + } + + tile_spec = { + "halo_size": halo_size, + "scale": scale, + } + grid_spec = GridSpec( + split_dims=(2, 3), + grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1), + tile_spec=tile_spec, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def patch_exec(self, task: TileTask) -> torch.Tensor: + return self.tile_exec(task) + + def patch_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): + halo = grid_spec.tile_spec["halo_size"][(i, j)] + scale = grid_spec.tile_spec["scale"] + + tile = coord_tensor_map[(i, j)] + halo_up = halo["up"] * scale + halo_down = halo["down"] * scale + halo_left = halo["left"] * scale + halo_right = halo["right"] * scale + + core_tile = tile[ + :, + :, + halo_up : (None if halo_down == 0 else -halo_down), + halo_left : (None if halo_right == 0 else -halo_right), + ] + result_row.append(core_tile) + result_rows.append(torch.cat(result_row, dim=3)) + dec = torch.cat(result_rows, dim=2) + return dec + + def _strategy_select(self, z: torch.Tensor): + tile_latent_min_size = getattr(self, "tile_latent_min_size", None) + tile_overlap_factor = getattr(self, "tile_overlap_factor", None) + if tile_latent_min_size is None or tile_overlap_factor is None: + return None, None, None + if z.shape[-1] > tile_latent_min_size or z.shape[-2] > tile_latent_min_size: + return self.tile_split, self.tile_exec, self.tile_merge + + return self.patch_split, self.patch_exec, self.patch_merge + + # Normally, we should override tiled_decode. However, since we also need to + # support the patch split strategy, we override decode instead. + def decode(self, z: torch.Tensor, return_dict: bool = True, *args: Any, **kwargs: Any): + if not self.is_distributed_enabled(): + return super().decode(z, return_dict=return_dict, *args, **kwargs) + + split, exec, merge = self._strategy_select(z) + + if split is not None: + strategy = "tile" if split == self.tile_split else "patch" + logger.info(f"Decode run with distributed executor, split strategy is {strategy}") + result = self.distributed_decoder.execute( + z, DistributedOperator(split=split, exec=exec, merge=merge), broadcast_result=False + ) + if not return_dict: + return (result,) + + from diffusers.models.autoencoders.vae import DecoderOutput + + return DecoderOutput(sample=result) + else: + return super().decode(z, return_dict=return_dict, *args, **kwargs) + + +class DistributedAutoencoderKL(DistributedAutoencoderKL_base, Diffusers_AutoencoderKL): + pass + + +# Next_Step_AutoencoderKL not support tiling now, so we currently disable it. + +# class DistributedAutoencoderKL_NextStep(DistributedAutoencoderKL_base, Next_Step_AutoencoderKL): +# pass diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py new file mode 100644 index 000000000..7549bbd3d --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import torch +from diffusers.models.autoencoders import AutoencoderKLQwenImage +from diffusers.models.autoencoders.vae import DecoderOutput +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeMixin, + GridSpec, + TileTask, +) + +logger = init_logger(__name__) + + +class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + # mostly copy from AutoencoderKL + _, _, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + + blend_height = self.tile_sample_min_height - self.tile_sample_stride_height + blend_width = self.tile_sample_min_width - self.tile_sample_stride_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + tiletask_list = [] + for i in range(0, height, tile_latent_stride_height): + for j in range(0, width, tile_latent_stride_width): + time_list = [] + for k in range(num_frames): + self._conv_idx = [0] + tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + time_list.append(tile) + tiletask_list.append( + TileTask( + len(tiletask_list), + (i // tile_latent_stride_height, j // tile_latent_stride_width), + time_list, + workload=time_list[0].shape[3] * time_list[0].shape[4], + ) + ) + tile_spec = { + "sample_height": sample_height, + "sample_width": sample_width, + "blend_height": blend_height, + "blend_width": blend_width, + } + 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=tile_spec, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def tile_exec(self, task: TileTask) -> torch.Tensor: + """Decode a single latent tile into RGB space.""" + self.clear_cache() + time = [] + for k in range(len(task.tensor)): + self._conv_idx = [0] + tile = self.post_quant_conv(task.tensor[k]) + decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx) + time.append(decoded) + result = torch.cat(time, dim=2) + return result + + def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + """Merge decoded tiles into a full image.""" + grid_h, grid_w = grid_spec.grid_shape + self.clear_cache() + + 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_height"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) + result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=3)[ + :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"] + ] + return dec + + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): + if not self.is_distributed_enabled(): + return super().tiled_decode(z, return_dict=return_dict) + + logger.info("Decode run with distributed executor") + result = self.distributed_decoder.execute( + z, + DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge), + broadcast_result=True, + ) + if not return_dict: + return (result,) + + return DecoderOutput(sample=result) diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py new file mode 100644 index 000000000..7defbae79 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import torch +from diffusers.models.autoencoders import AutoencoderKLWan +from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify +from diffusers.models.autoencoders.vae import DecoderOutput +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + DistributedVaeMixin, + GridSpec, + TileTask, +) + +logger = init_logger(__name__) + + +class DistributedAutoencoderKLWan(AutoencoderKLWan, DistributedVaeMixin): + @classmethod + def from_pretrained(cls, *args: Any, **kwargs: Any): + model = super().from_pretrained(*args, **kwargs) + model.init_distributed() + return model + + def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]: + _, _, num_frames, height, width = z.shape + sample_height = height * self.spatial_compression_ratio + sample_width = width * self.spatial_compression_ratio + + tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio + tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio + tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio + tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio + tile_sample_stride_height = self.tile_sample_stride_height + tile_sample_stride_width = self.tile_sample_stride_width + if self.config.patch_size is not None: + sample_height = sample_height // self.config.patch_size + sample_width = sample_width // self.config.patch_size + tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size + tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size + blend_height = self.tile_sample_min_height // self.config.patch_size - tile_sample_stride_height + blend_width = self.tile_sample_min_width // self.config.patch_size - tile_sample_stride_width + else: + blend_height = self.tile_sample_min_height - tile_sample_stride_height + blend_width = self.tile_sample_min_width - tile_sample_stride_width + + tiletask_list = [] + for i in range(0, height, tile_latent_stride_height): + for j in range(0, width, tile_latent_stride_width): + time_list = [] + for k in range(num_frames): + self._conv_idx = [0] + tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width] + time_list.append(tile) + tiletask_list.append( + TileTask( + len(tiletask_list), + (i // tile_latent_stride_height, j // tile_latent_stride_width), + time_list, + workload=time_list[0].shape[3] * time_list[0].shape[4], + ) + ) + tile_spec = { + "sample_height": sample_height, + "sample_width": sample_width, + "blend_height": blend_height, + "blend_width": blend_width, + "tile_sample_stride_height": tile_sample_stride_height, + "tile_sample_stride_width": tile_sample_stride_width, + } + 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=tile_spec, + output_dtype=self.dtype, + ) + return tiletask_list, grid_spec + + def tile_exec(self, task: TileTask) -> torch.Tensor: + """Decode a single latent tile into RGB space.""" + self.clear_cache() + time = [] + for k in range(len(task.tensor)): + self._conv_idx = [0] + tile = self.post_quant_conv(task.tensor[k]) + decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=(k == 0)) + time.append(decoded) + result = torch.cat(time, dim=2) + return result + + def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor: + """Merge decoded tiles into a full image.""" + grid_h, grid_w = grid_spec.grid_shape + result_rows = [] + self.clear_cache() + 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_height"]) + if j > 0: + tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"]) + result_row.append( + tile[ + :, + :, + :, + : grid_spec.tile_spec["tile_sample_stride_height"], + : grid_spec.tile_spec["tile_sample_stride_width"], + ] + ) + result_rows.append(torch.cat(result_row, dim=-1)) + + dec = torch.cat(result_rows, dim=3)[ + :, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"] + ] + + if self.config.patch_size is not None: + dec = unpatchify(dec, patch_size=self.config.patch_size) + + dec = torch.clamp(dec, min=-1.0, max=1.0) + return dec + + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): + if not self.is_distributed_enabled(): + return super().tiled_decode(z, return_dict=return_dict) + + logger.info("Decode run with distributed executor") + result = self.distributed_decoder.execute( + z, + DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge), + broadcast_result=False, + ) + if not return_dict: + return (result,) + + return DecoderOutput(sample=result) diff --git a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py new file mode 100644 index 000000000..bdf664741 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch +import torch.distributed as dist +from vllm.logger import init_logger + +from vllm_omni.diffusion.distributed.parallel_state import get_dit_group + +logger = init_logger(__name__) + + +@dataclass +class GridSpec: + """The Grid shape split""" + + split_dims: tuple[int, ...] # For example: (3,4) for (B, C, T, H, W), (2,3) for (B, C, H, W) + grid_shape: tuple[int, ...] # For example: (nh, nw) for (B, C, T, H, W), (nh, nw) for (B, C, H, W) + tile_spec: dict = field(default_factory=dict) + output_dtype: torch.dtype | None = None + + +@dataclass +class TileTask: + tile_id: int + grid_coord: tuple[int, ...] # The coordinate of the tile in GridSpec.grid_shape + tensor: torch.Tensor | list[torch.Tensor] # The tile tensor + workload: int | float = 1 + + +@dataclass +class DistributedOperator: + split: callable + exec: callable + merge: callable + + +class DistributedVaeExecutor: + """ + Abstract util class for distributed patch/tile parallel VAE decoding. + """ + + def __init__(self): + self.group = get_dit_group() + self.world_size = dist.get_world_size(self.group) + self.rank = dist.get_rank(self.group) + self.parallel_size = 1 + + def set_parallel_size(self, parallel_size: int): + self.parallel_size = parallel_size + + def gather_tensors(self, tensor: torch.Tensor): + gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] if self.rank == 0 else None + dist.gather(tensor, gather_list=gather_list, dst=0, group=self.group) + return gather_list + + def broadcast_tensor(self, tensor: torch.Tensor): + dist.broadcast(tensor, src=0, group=self.group) + return tensor + + def _compute_global_padding_shape(self, local_results, output_ndim: int, device): + local_tile_max_dims = [0] * output_ndim + for _, tile_tensor in local_results: + for dim_idx, dim_size in enumerate(tile_tensor.shape): + local_tile_max_dims[dim_idx] = max(local_tile_max_dims[dim_idx], dim_size) + local_shape_stat = torch.tensor([len(local_results), *local_tile_max_dims], device=device) + dist.all_reduce( + local_shape_stat, + op=dist.ReduceOp.MAX, + group=self.group, + ) + return local_shape_stat.tolist() + + def _pack_local_tiles(self, local_results, global_padding_shape, grid_spec, device, dtype): + tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype) + meta_tensor = torch.full( + (global_padding_shape[0], len(grid_spec.split_dims) + 1), -1, device=device, dtype=torch.int64 + ) + for idx, (tid, t_tensor) in enumerate(local_results): + meta_tensor[idx, 0] = tid + for i, dim in enumerate(grid_spec.split_dims): + meta_tensor[idx, i + 1] = t_tensor.shape[dim] + slices = tuple(slice(0, s) for s in t_tensor.shape) + tile_tensor[idx][slices] = t_tensor + return tile_tensor, meta_tensor + + def _unpack_tiles(self, meta_gather, tile_gather, grid_spec, tid_coord_map): + coord_tensor_map = {} + for r in range(self.world_size): + meta_src = meta_gather[r] + tiles_src = tile_gather[r] + for idx in range(meta_src.shape[0]): + tid = int(meta_src[idx, 0]) + if tid < 0: + continue + slices = [slice(None)] * tiles_src[idx].ndim + for i, dim in enumerate(grid_spec.split_dims): + slices[dim] = slice(0, int(meta_src[idx, i + 1])) + slices = tuple(slices) + coord_tensor_map[tid_coord_map[tid]] = tiles_src[idx][slices] + + return coord_tensor_map + + def _balance_tasks(self, task_list, num_rank): + workloads = [0] * num_rank + assigned = [[] for _ in range(num_rank)] + + for task in sorted(task_list, key=lambda t: t.workload, reverse=True): + r = workloads.index(min(workloads)) + assigned[r].append(task) + workloads[r] += task.workload + + return assigned + + def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True): + pp_size = min(self.parallel_size, self.world_size) + + # 1. Split into tiles + tiletask_list, grid_spec = operator.split(z) + tid_coord_map = {task.tile_id: task.grid_coord for task in tiletask_list} + + # 2. local decode + assigned = self._balance_tasks(tiletask_list, pp_size) + local_tasks = assigned[self.rank] if pp_size <= self.world_size else [] + local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks] + + # 3. compute shape per rank + global_padding_shape = self._compute_global_padding_shape(local_results, z.ndim, z.device) + + # 5. prepare tile tensors + output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype + local_tile_tensor, local_meta_tensor = self._pack_local_tiles( + local_results, global_padding_shape, grid_spec, z.device, output_dtype + ) + + # 6. gather tiles & meta + meta_gather = self.gather_tensors(local_meta_tensor) + tile_gather = self.gather_tensors(local_tile_tensor) + + if self.rank != 0: + result = torch.empty(0, device=z.device) # Dummy return for non-zero ranks + else: + # 7. reconstruct full tensor (rank 0) + coord_tensor_map = self._unpack_tiles(meta_gather, tile_gather, grid_spec, tid_coord_map) + result = operator.merge(coord_tensor_map, grid_spec) + + if broadcast_result: + result = self._sync_final_result(result, z.ndim, z.device, output_dtype) + return result + + def _sync_final_result(self, rank0_result, output_ndim, output_device, output_dtype): + shape_tensor = torch.empty((output_ndim,), device=output_device, dtype=torch.int64) + if self.rank == 0: + shape_tensor.copy_(torch.tensor(tuple(rank0_result.shape), device=output_device, dtype=torch.int64)) + dist.broadcast(shape_tensor, src=0, group=self.group) + + if self.rank != 0: + sync_result = torch.empty(tuple(shape_tensor.tolist()), device=output_device, dtype=output_dtype) + else: + sync_result = rank0_result + dist.broadcast(sync_result, src=0, group=self.group) + return sync_result + + +class DistributedVaeMixin: + def init_distributed(self): + self.distributed_decoder = DistributedVaeExecutor() + + def set_parallel_size(self, parallel_size: int) -> bool: + return self.distributed_decoder.set_parallel_size(parallel_size) + + def is_distributed_enabled(self) -> bool: + if ( + self.distributed_decoder.parallel_size <= 1 + or not dist.is_initialized() + or not getattr(self, "use_tiling", False) + ): + return False + world_size = dist.get_world_size(group=self.distributed_decoder.group) + pp_size = min(int(self.distributed_decoder.parallel_size), int(world_size)) + if pp_size <= 1: + return False + if self.distributed_decoder.parallel_size > pp_size: + logger.warning( + f"vae_patch_parallel_size={self.distributed_decoder.parallel_size} " + f"is greater than dit_group={world_size};" + f" using dit_group size={world_size}" + ) + return True diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index a61044304..02758ee19 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -13,9 +13,6 @@ import numpy as np import torch import torch.distributed from diffusers.image_processor import VaeImageProcessor -from diffusers.models.autoencoders.autoencoder_kl_qwenimage import ( - AutoencoderKLQwenImage, -) from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) @@ -25,6 +22,7 @@ from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( @@ -269,9 +267,9 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin): self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only ) - self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( - self.device - ) + self.vae = DistributedAutoencoderKLQwenImage.from_pretrained( + model, subfolder="vae", local_files_only=local_files_only + ).to(self.device) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config) self.transformer = QwenImageTransformer2DModel( diff --git a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py index 3668c132f..9c9722c36 100644 --- a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py +++ b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py @@ -6,7 +6,6 @@ from collections.abc import Iterable import torch from diffusers.image_processor import VaeImageProcessor -from diffusers.models.autoencoders import AutoencoderKL from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) @@ -16,6 +15,7 @@ from transformers import CLIPTextModelWithProjection, CLIPTokenizer, T5EncoderMo from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -174,9 +174,9 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin): ) self.transformer = SD3Transformer2DModel(od_config=od_config) - self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( - self.device - ) + self.vae = DistributedAutoencoderKL.from_pretrained( + model, subfolder="vae", local_files_only=local_files_only + ).to(self.device) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 23eebc7fd..63de7395f 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -11,13 +11,13 @@ from typing import Any, cast import PIL.Image import torch -from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, UMT5EncoderModel from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -267,7 +267,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin): self.text_encoder = UMT5EncoderModel.from_pretrained( model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) - self.vae = AutoencoderKLWan.from_pretrained( + self.vae = DistributedAutoencoderKLWan.from_pretrained( model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only ).to(self.device) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 4b32d5a5f..2d750c762 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -11,13 +11,13 @@ from typing import Any, cast import numpy as np import PIL.Image import torch -from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -210,7 +210,7 @@ class Wan22I2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin): self.image_encoder = None # VAE - self.vae = AutoencoderKLWan.from_pretrained( + self.vae = DistributedAutoencoderKLWan.from_pretrained( model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only ).to(self.device) diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index 2c6ef4c86..b16d24bf7 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -24,7 +24,6 @@ from typing import Any import torch import torch.nn as nn from diffusers.image_processor import VaeImageProcessor -from diffusers.models.autoencoders import AutoencoderKL from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from diffusers.utils import logging from diffusers.utils.torch_utils import randn_tensor @@ -32,6 +31,7 @@ from transformers import AutoModel, AutoTokenizer from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.z_image.z_image_transformer import ( @@ -171,9 +171,9 @@ class ZImagePipeline(nn.Module): self.text_encoder = AutoModel.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only ) - self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( - self._execution_device - ) + self.vae = DistributedAutoencoderKL.from_pretrained( + model, subfolder="vae", local_files_only=local_files_only + ).to(self._execution_device) # Get vLLM quantization config for linear layers quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config) self.transformer = ZImageTransformer2DModel(quant_config=quant_config) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index ed99b1473..c806dbae9 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -8,6 +8,7 @@ from vllm.logger import init_logger from vllm.model_executor.models.registry import _LazyRegisteredModel, _ModelRegistry from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelConfig, get_sp_plan_from_model from vllm_omni.diffusion.hooks.sequence_parallel import apply_sequence_parallel @@ -123,13 +124,6 @@ DiffusionModelRegistry = _ModelRegistry( } ) -_VAE_PATCH_PARALLEL_ALLOWLIST = { - # Only enable for models we have validated end-to-end. - "StableDiffusion3Pipeline", - "ZImagePipeline", - "NextStep11Pipeline", -} - _NO_CACHE_ACCELERATION = { # Pipelines that do not support cache acceleration (cache_dit / tea_cache). "NextStep11Pipeline", @@ -161,17 +155,14 @@ def initialize_model( model = model_class(od_config=od_config) vae_pp_size = od_config.parallel_config.vae_patch_parallel_size - if vae_pp_size > 1 and od_config.model_class_name not in _VAE_PATCH_PARALLEL_ALLOWLIST: + is_distributed_vae = hasattr(model, "vae") and isinstance(model.vae, DistributedVaeMixin) + if vae_pp_size > 1 and not is_distributed_vae: logger.warning( - "vae_patch_parallel_size=%d is set but VAE patch parallelism is only enabled for %s; ignoring.", + "vae_patch_parallel_size=%d is set but VAE patch parallelism is NOT enabled for %s; ignoring.", vae_pp_size, - sorted(_VAE_PATCH_PARALLEL_ALLOWLIST), + od_config.model_class_name, ) - if ( - vae_pp_size > 1 - and od_config.model_class_name in _VAE_PATCH_PARALLEL_ALLOWLIST - and not od_config.vae_use_tiling - ): + if vae_pp_size > 1 and is_distributed_vae and not od_config.vae_use_tiling: logger.info( "vae_patch_parallel_size=%d requires vae_use_tiling; automatically enabling it.", vae_pp_size, @@ -184,20 +175,8 @@ def initialize_model( if hasattr(model.vae, "use_tiling"): model.vae.use_tiling = od_config.vae_use_tiling - if ( - vae_pp_size > 1 - and hasattr(model, "vae") - and od_config.model_class_name in _VAE_PATCH_PARALLEL_ALLOWLIST - and od_config.vae_use_tiling - ): - from vllm_omni.diffusion.distributed.parallel_state import get_dit_group - from vllm_omni.diffusion.distributed.vae_patch_parallel import maybe_wrap_vae_decode_with_patch_parallelism - - maybe_wrap_vae_decode_with_patch_parallelism( - model, - vae_patch_parallel_size=vae_pp_size, - group_getter=get_dit_group, - ) + if is_distributed_vae: + model.vae.set_parallel_size(vae_pp_size) # Apply sequence parallelism if enabled # This follows diffusers' pattern where enable_parallelism() is called From b13761ed264e8160b622c1f1d0a0e62376efc2b2 Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Tue, 3 Mar 2026 18:40:53 +0800 Subject: [PATCH 6/7] [Bugfix] Fix filepath resolution for model with subdir and GLM-Image generation (#1609) Signed-off-by: yuanheng Signed-off-by: Yuanheng Zhao --- .../offline_inference/glm_image/README.md | 8 ++--- .../offline_inference/glm_image/end2end.py | 5 +-- vllm_omni/entrypoints/stage_utils.py | 33 ++++++++++++++----- 3 files changed, 29 insertions(+), 17 deletions(-) diff --git a/examples/offline_inference/glm_image/README.md b/examples/offline_inference/glm_image/README.md index 73ae0e046..1e7e75144 100644 --- a/examples/offline_inference/glm_image/README.md +++ b/examples/offline_inference/glm_image/README.md @@ -40,8 +40,7 @@ GLM-Image uses a 2-stage pipeline: ```bash python end2end.py \ - --model-path /path/to/glm-image \ - --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \ + --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \ --prompt "A beautiful sunset over the ocean with sailing boats" \ --height 1024 \ --width 1024 \ @@ -52,8 +51,7 @@ python end2end.py \ ```bash python end2end.py \ - --model-path /path/to/glm-image \ - --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \ + --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \ --prompt "Transform this scene into a winter wonderland" \ --image input.png \ --output output_i2i.png @@ -64,7 +62,7 @@ python end2end.py \ ```bash python end2end.py \ --model-path /path/to/glm-image \ - --config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \ + --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \ --prompt "A photorealistic cat sitting on a window sill" \ --height 1024 \ --width 1024 \ diff --git a/examples/offline_inference/glm_image/end2end.py b/examples/offline_inference/glm_image/end2end.py index 884ebddee..77de043a4 100644 --- a/examples/offline_inference/glm_image/end2end.py +++ b/examples/offline_inference/glm_image/end2end.py @@ -178,9 +178,6 @@ def main(args: argparse.Namespace) -> None: print("=" * 60) # Validate arguments - if not args.model_path: - raise ValueError("--model-path is required") - if not args.prompt: raise ValueError("--prompt is required") @@ -365,7 +362,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--model-path", type=str, - required=True, + default="zai-org/GLM-Image", help="Path to GLM-Image model directory or HuggingFace model ID", ) parser.add_argument( diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index fade4c844..110bf48cd 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -322,6 +322,21 @@ def _to_dict(x: Any) -> dict[str, Any]: return {} +def _resolve_model_to_local_path(model: str) -> str: + """Resolve an HF Hub model ID to its local cache snapshot path.""" + if os.path.isdir(model): + return model + + try: + from huggingface_hub import snapshot_download + + # no network access is attempted, check local model path only + return snapshot_download(model, local_files_only=True) + except Exception: + logger.warning(f"Could not resolve {model} to a local snapshot path; using as-is", exc_info=True) + return model + + def _resolve_model_tokenizer_paths( model: str, engine_args: dict, @@ -331,30 +346,32 @@ def _resolve_model_tokenizer_paths( Some models (e.g., GLM-Image) have tokenizer in root and model in subdirectory. This function handles model_subdir and tokenizer_subdir engine_args. + When the base model path is an HF Hub ID rather than an absolute local path, + the ID is first resolved to the local snapshot directory so that subdirectory + joins produce valid filesystem paths. + Args: - model: Base model path + model: Base model path or HF Hub model ID engine_args: Engine arguments (modified in-place to remove subdir args and set tokenizer if needed) Returns: Resolved model path (may be subdirectory of original) """ - import os - model_subdir = engine_args.pop("model_subdir", None) tokenizer_subdir = engine_args.pop("tokenizer_subdir", None) - base_model_path = model + resolved_base = _resolve_model_to_local_path(model) if model_subdir: - model = os.path.join(model, model_subdir) + model = os.path.join(resolved_base, model_subdir) logger.info(f"Using model subdirectory: {model}") if tokenizer_subdir is not None: - tokenizer_path = os.path.join(base_model_path, tokenizer_subdir) if tokenizer_subdir else base_model_path + tokenizer_path = os.path.join(resolved_base, tokenizer_subdir) if tokenizer_subdir else resolved_base engine_args["tokenizer"] = tokenizer_path logger.info(f"Using tokenizer from: {tokenizer_path}") elif model_subdir and "tokenizer" not in engine_args: - engine_args["tokenizer"] = base_model_path - logger.info(f"Using tokenizer from base model path: {base_model_path}") + engine_args["tokenizer"] = resolved_base + logger.info(f"Using tokenizer from base model path: {resolved_base}") return model From 742d845b8ea1c018cd5370a4b47e31bad407e6c3 Mon Sep 17 00:00:00 2001 From: Junhong Liu Date: Tue, 3 Mar 2026 22:06:52 +0800 Subject: [PATCH 7/7] Make chunk_size and left_context_size configurable via YAML for async chunking (#1423) Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com> Signed-off-by: Junhong Liu --- .../core/sched/omni_generation_scheduler.py | 3 ++ vllm_omni/core/sched/output.py | 1 + .../chunk_transfer_adapter.py | 5 ++++ .../models/qwen3_omni/qwen3_omni.py | 20 ++++++++++--- .../models/qwen3_omni/qwen3_omni_code2wav.py | 25 +++++++++-------- .../models/qwen3_tts/qwen3_tts_code2wav.py | 24 ++++++++++++---- .../qwen3_omni_moe_async_chunk.yaml | 16 +++++++++++ .../stage_input_processors/qwen3_omni.py | 28 +++++++++++++------ .../stage_input_processors/qwen3_tts.py | 25 +++++++++-------- vllm_omni/worker/gpu_model_runner.py | 10 +++++-- 10 files changed, 113 insertions(+), 44 deletions(-) diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 75ffa8083..96f9f2402 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -49,6 +49,7 @@ class OmniGenerationScheduler(VLLMScheduler): scheduled_spec_decode_tokens: dict[str, list[int]] = {} scheduled_encoder_inputs: dict[str, list[int]] = {} cached_prompt_token_ids: dict[str, list[int]] = {} + cached_additional_information: dict[str, dict | None] = {} # Temporary queue: preserve waiting order, do not disturb non-diffusion requests skipped_waiting_requests = create_request_queue(self.policy) @@ -105,6 +106,7 @@ class OmniGenerationScheduler(VLLMScheduler): req_to_new_blocks[request.request_id] = new_blocks num_scheduled_tokens[request.request_id] = num_new_tokens cached_prompt_token_ids[request.request_id] = request.prompt_token_ids + cached_additional_information[request.request_id] = getattr(request, "additional_information", None) token_budget -= num_new_tokens scheduled_running_reqs.append(request) req_index += 1 @@ -225,6 +227,7 @@ class OmniGenerationScheduler(VLLMScheduler): num_computed_tokens=cached_reqs_data.num_computed_tokens, num_output_tokens=cached_reqs_data.num_output_tokens, prompt_token_ids=cached_prompt_token_ids, + additional_information=cached_additional_information, ) total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index b86a97534..f09e75c2b 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -71,6 +71,7 @@ class OmniCachedRequestData(CachedRequestData): """ prompt_token_ids: dict[str, list[int]] + additional_information: dict[str, dict | None] @dataclass diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index d6fc81c6d..d80323bbd 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -155,6 +155,11 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): new_ids = payload_data.get("code_predictor_codes", []) request.prompt_token_ids = new_ids + # Pass additional fields (like left_context_size) to the request + # Only pass chunk context metadata in additional_information + request.additional_information = {} + if "left_context_size" in payload_data: + request.additional_information["left_context_size"] = payload_data["left_context_size"] request.num_computed_tokens = 0 # Empty chunk with more data expected: keep polling. diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index f798a21d2..7506c117d 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -5,6 +5,7 @@ from collections.abc import Iterable from functools import cached_property +from typing import Any import torch import torch.nn as nn @@ -160,6 +161,7 @@ class Qwen3OmniMoeForConditionalGeneration( self.requires_raw_input_tokens = True elif self.model_stage == "code2wav": + self.enable_update_additional_information = True self.thinker = None self.talker = None # Initialize code2wav (codec codes → audio waveform) @@ -254,7 +256,7 @@ class Qwen3OmniMoeForConditionalGeneration( codec: torch.Tensor | None = None, sampling_metadata: SamplingMetadata | None = None, logits_index: int | None = None, - additional_information: dict[str, object] | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, **kwargs: object, ) -> torch.Tensor | IntermediateTensors | OmniOutput: """ @@ -359,7 +361,15 @@ class Qwen3OmniMoeForConditionalGeneration( codes = input_ids_flatten.reshape(1, 16, -1) # Generate audio from codec codes - audio_tensors = self.generate_audio(codes, voice_type, seq_token_counts) + # Get every request's left_context_size from runtime_additional_information (passed via kwargs) + left_context_size = [] + if runtime_additional_information is not None: + for info in runtime_additional_information: + if "left_context_size" in info: + left_context_size.append(info["left_context_size"]) + else: + logger.debug("No additional_information provided to code2wav stage.") + audio_tensors = self.generate_audio(codes, voice_type, left_context_size, seq_token_counts) return audio_tensors @@ -440,6 +450,7 @@ class Qwen3OmniMoeForConditionalGeneration( self, code: torch.Tensor, voice_type: str, + left_context_size: list[int] | None = None, seq_token_counts: list[int] | None = None, ) -> list[torch.Tensor]: """ @@ -448,6 +459,7 @@ class Qwen3OmniMoeForConditionalGeneration( Args: code: [batch, num_quantizers, T] - RVQ codec codes voice_type: Voice type (not used in Qwen3, kept for compatibility) + left_context_size: Left context size for streaming decode seq_token_counts: Token count for each request in batch Returns: @@ -471,10 +483,10 @@ class Qwen3OmniMoeForConditionalGeneration( talker_codes = talker_codes.expand(1, 16, -1) if self.vllm_config.model_config.async_chunk: + # Only use left_context_size from additional information audio_tensors = self.code2wav.chunked_decode_streaming( talker_codes, - chunk_size=25, - left_context_size=25, + left_context_size=left_context_size, seq_token_counts=seq_token_counts, ) else: diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index 54bc1c1c5..41eeec5e7 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -213,8 +213,7 @@ class Qwen3OmniMoeCode2Wav(nn.Module): def chunked_decode_streaming( self, codes: torch.Tensor, - chunk_size: int = 25, - left_context_size: int = 25, + left_context_size: list[int] | None = None, seq_token_counts: list[int] | None = None, ) -> list[torch.Tensor]: """ @@ -222,9 +221,10 @@ class Qwen3OmniMoeCode2Wav(nn.Module): Uses overlapping chunks with left context to avoid boundary artifacts. + No longer need chunk size here, which is different from chunked_decode + Args: codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes - chunk_size: Number of codec frames per chunk left_context_size: Number of overlapping frames for context seq_token_counts: Token count for each request in batch @@ -233,6 +233,13 @@ class Qwen3OmniMoeCode2Wav(nn.Module): codes. For ``batch_size == 1``, this is a list containing a single tensor with shape ``[1, waveform_len]``. """ + if not (left_context_size and seq_token_counts and len(left_context_size) == len(seq_token_counts)): + logger.warning( + "chunked_decode_streaming: missing/invalid left_context_size or seq_token_counts; " + "defaulting to left_context_size=zeros(len=codes.shape[0])." + ) + left_context_size = [0] * codes.shape[0] + # Decode chunk wavs = [] batch_wav = self(codes) if seq_token_counts is not None: @@ -241,14 +248,10 @@ class Qwen3OmniMoeCode2Wav(nn.Module): # Fallback: assume all batch elements share the same sequence length. code_seq_lens = [codes.shape[-1]] * codes.shape[0] for idx, code_seq_len in enumerate(code_seq_lens): - # TODO: need to optimize algorithms, current only support - # chunk_size = left_context_size = 25 - if code_seq_len <= chunk_size: - context_size = 0 - else: - context_size = left_context_size - # Remove context from output (context_size * total_upsample samples) - wav_chunk = batch_wav[idx, :, context_size * self.total_upsample : code_seq_len * self.total_upsample] + # Remove context from output (left_context_size * total_upsample samples) + wav_chunk = batch_wav[ + idx, :, left_context_size[idx] * self.total_upsample : code_seq_len * self.total_upsample + ] wavs.append(wav_chunk) return wavs diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index a87027de1..851455bda 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -33,6 +33,7 @@ class Qwen3TTSCode2Wav(nn.Module): self.have_multimodal_outputs = True self.has_preprocess = False self.has_postprocess = False + self.enable_update_additional_information = True # Generation-only stage (no logits / sampling). self.requires_raw_input_tokens = True @@ -146,6 +147,7 @@ class Qwen3TTSCode2Wav(nn.Module): positions: torch.Tensor | None = None, intermediate_tensors: Any = None, inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, **kwargs: Any, ) -> OmniOutput: """Decode codec codes into audio waveform. @@ -177,18 +179,28 @@ class Qwen3TTSCode2Wav(nn.Module): ids = input_ids.reshape(-1).to(dtype=torch.long) request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts")) - # Parse each request: extract ctx_frames, validate, reshape codes. - # input_ids layout per request: [codec_context_frames, *flat_codes] + # Parse each request: extract left_context_size, validate, reshape codes. + # input_ids layout per request: [*flat_codes] # where flat_codes is codebook-major [q*F]. - parsed = [] # (ctx_frames, actual_frames) + parsed = [] # (left_context_size, actual_frames) valid_codes = [] valid_indices = [] + # Get left_context_size from runtime_additional_information (passed via kwargs) + left_context_size = [0] * len(request_ids_list) + if runtime_additional_information is not None: + for i, info in enumerate(runtime_additional_information): + if i >= len(left_context_size): + break + if "left_context_size" in info: + left_context_size[i] = info["left_context_size"] + else: + logger.debug("No additional_information provided to code2wav stage.") for i, req_ids in enumerate(request_ids_list): - if req_ids.numel() < 2: + if req_ids.numel() < 1: parsed.append((0, 0)) continue - ctx_frames = int(req_ids[0].item()) - flat = req_ids[1:] + ctx_frames = left_context_size[i] + flat = req_ids n = flat.numel() # Warmup / dummy_run: not divisible by num_quantizers. if n == 0 or n % q != 0: diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml index da996494e..8d8d00d6b 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml @@ -29,6 +29,9 @@ stage_args: final_output: true final_output_type: text is_comprehension: true + # Use named connector to apply runtime.connectors.extra. + output_connectors: + to_stage_1: connector_of_shared_memory default_sampling_params: temperature: 0.4 top_p: 0.9 @@ -60,6 +63,9 @@ stage_args: engine_input_source: [0] # final_output: true # final_output_type: text + # Distributed connector configuration + input_connectors: + from_stage_0: connector_of_shared_memory default_sampling_params: temperature: 0.9 top_k: 50 @@ -99,3 +105,13 @@ stage_args: seed: 42 detokenize: True repetition_penalty: 1.1 + +runtime: + + connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + # Align with Omni: small chunks with sufficient context overlap. + codec_chunk_frames: 25 # code2wav decode chunk size + codec_left_context_frames: 25 # code2wav left context size diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index 3a42159a8..31a5ddd47 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -220,6 +220,12 @@ def talker2code2wav_async_chunk( if "code_predictor_codes" not in pooling_output: return None + connector = getattr(transfer_manager, "connector", None) + raw_cfg = getattr(connector, "config", {}) or {} + cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} + chunk_size_config = int(cfg.get("codec_chunk_frames", 25)) + left_context_size_config = int(cfg.get("codec_left_context_frames", 25)) + code_predictor_codes = pooling_output["code_predictor_codes"] if code_predictor_codes is None: @@ -244,23 +250,27 @@ def talker2code2wav_async_chunk( return None request_id = request.external_req_id - chunk_size = left_context_size = 25 transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) length = len(transfer_manager.code_prompt_token_ids[request_id]) - chunk_length = length % chunk_size + chunk_length = length % chunk_size_config if chunk_length != 0 and not is_finished: return None - context_length = chunk_length if chunk_length != 0 else chunk_size + context_length = chunk_length if chunk_length != 0 else chunk_size_config + # ensure left context does not exceed available length + left_context_size = max(0, min(length - context_length, left_context_size_config)) end_index = min(length, left_context_size + context_length) + codes = ( + torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]) + .transpose(0, 1) + .reshape(-1) + .tolist() + ) + info = { - "code_predictor_codes": ( - torch.tensor(transfer_manager.code_prompt_token_ids[request_id][-end_index:]) - .transpose(0, 1) - .reshape(-1) - .tolist() - ), + "code_predictor_codes": codes, + "left_context_size": left_context_size, "finished": torch.tensor(is_finished, dtype=torch.bool), } return info diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index 72e17bf4f..8c21052e9 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -40,12 +40,12 @@ def talker2code2wav_async_chunk( connector = getattr(transfer_manager, "connector", None) raw_cfg = getattr(connector, "config", {}) or {} cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} - chunk_size = int(cfg.get("codec_chunk_frames", 25)) - left_context_size = int(cfg.get("codec_left_context_frames", 25)) - if chunk_size <= 0 or left_context_size < 0: + chunk_size_config = int(cfg.get("codec_chunk_frames", 25)) + left_context_size_config = int(cfg.get("codec_left_context_frames", 25)) + if chunk_size_config <= 0 or left_context_size_config < 0: raise ValueError( - f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, " - f"codec_left_context_frames={left_context_size}" + f"Invalid codec chunk config: codec_chunk_frames={chunk_size_config}, " + f"codec_left_context_frames={left_context_size_config}" ) length = len(transfer_manager.code_prompt_token_ids[request_id]) @@ -59,22 +59,23 @@ def talker2code2wav_async_chunk( } return None - chunk_length = length % chunk_size + chunk_length = length % chunk_size_config if chunk_length != 0 and not finished: return None - context_length = chunk_length if chunk_length != 0 else chunk_size - end_index = min(length, left_context_size + context_length) - ctx_frames = max(0, int(end_index - context_length)) + context_length = chunk_length if chunk_length != 0 else chunk_size_config + end_index = min(length, left_context_size_config + context_length) + left_context_size = max(0, int(end_index - context_length)) window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] # Pack context + chunk into codebook-major flat codes for adapter. code_predictor_codes = torch.tensor(window_frames).transpose(0, 1).reshape(-1).tolist() - # Build final prompt_token_ids with ctx_frames header for Qwen3-TTS Code2Wav. - # The model expects input_ids layout: [ctx_frames, *flat_codes]. + # Build final prompt_token_ids and left_context_size header for Qwen3-TTS Code2Wav. + # The model expects input_ids layout: [*flat_codes]. return { - "code_predictor_codes": [int(ctx_frames)] + code_predictor_codes, + "code_predictor_codes": code_predictor_codes, + "left_context_size": left_context_size, "finished": torch.tensor(finished, dtype=torch.bool), } diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 32ab75429..af2344ce7 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1155,12 +1155,18 @@ class OmniGPUModelRunner(GPUModelRunner): # Prefill: overlay prompt_embeds and collect additional_information self._collect_additional_information_for_prefill(num_scheduled_tokens_np) + # Keep per-request additional_information in sync for both new and + # cached requests. This is required for stages without preprocess + # (e.g., code2wav) so runtime_additional_information can be refreshed + # from scheduler cached infos on every step. + if hasattr(self.model, "has_preprocess") or hasattr(self.model, "enable_update_additional_information"): + if self.vllm_config.model_config.async_chunk: + self._update_additional_information(scheduler_output) + if hasattr(self.model, "has_preprocess") and self.model.has_preprocess: # Overlay custom prompt_embeds per request for the prompt portion; # collect additional_information (tensor/list) for prefill portion only decode_req_ids = [] - if self.vllm_config.model_config.async_chunk: - self._update_additional_information(scheduler_output) for req_index, req_id in enumerate(self.input_batch.req_ids): req_infos = self.model_intermediate_buffer.get(req_id, {})